"""Generate golden test vectors for the ``tm30`` crate.

Every case runs the production Python reference ``compute_tm30`` (HPCS
``tm30.py`` on colour-science 0.4.7) and writes one JSON file into
``crates/tm30/tests/golden`` holding the exact input (wavelengths, values)
and every output field. NaN is encoded as ``null``. Floats are written with
``repr`` so they round-trip exactly.

Usage (from the repository root):

    python -W ignore crates/tm30/tools/golden.py [--tm30-dir DIR] [--lamps-dir DIR]

``--tm30-dir`` is the folder containing ``tm30.py`` (defaults to the HPCS
scratch copy), ``--lamps-dir`` the buildfor.life ``comparison-data/lamps``
folder with ``L*/spd.csv`` files (read only).
"""

from __future__ import annotations

import argparse
import csv
import glob
import json
import math
import os
import sys
import warnings

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.normpath(os.path.join(HERE, "..", "tests", "golden"))

DEFAULT_TM30_DIR = os.path.join(
    os.path.expanduser("~"), "AppData", "Local", "hermes", "cache", "scratch", "hpcs-py"
)
DEFAULT_LAMPS_DIR = "C:/dev/buildfor_life_web/apps/web/comparison-data/lamps"

# HPCS 6500 grid: 350 points, 380..1050 nm.
HPCS_GRID = np.linspace(380.0, 1050.0, 350)


def clean(x):
    x = float(x)
    return None if math.isnan(x) else x


def read_lamp_csv(path):
    wl, vals = [], []
    with open(path, newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            wl.append(float(row["wavelength_nm"]))
            vals.append(float(row["value"]))
    return np.array(wl), np.array(vals)


def build_cases(lamps_dir):
    import colour
    from colour.colorimetry import SDS_ILLUMINANTS, SpectralShape, sd_blackbody
    from colour.colorimetry import sd_CIE_illuminant_D_series
    from colour.temperature import CCT_to_xy_CIE_D

    cases = []

    # Real lamps measured with the HPCS 6500 (380-1050 nm, ~1.92 nm steps).
    for path in sorted(glob.glob(os.path.join(lamps_dir, "L*", "spd.csv"))):
        lamp = os.path.basename(os.path.dirname(path))
        wl, vals = read_lamp_csv(path)
        cases.append((f"lamp_{lamp}", f"{lamp}/spd.csv", wl, vals))

    # colour-science illuminant library.
    names = (
        [f"FL{i}" for i in range(1, 13)]
        + ["A", "D50", "D65"]
        + [k for k in SDS_ILLUMINANTS.keys() if k.startswith("LED-")]
        + [f"HP{i}" for i in range(1, 6)]
    )
    for name in names:
        sd = SDS_ILLUMINANTS[name]
        cases.append(
            (
                "illuminant_" + name.replace("-", "_").replace(".", "_"),
                f"colour.SDS_ILLUMINANTS[{name!r}] {sd.shape}",
                sd.wavelengths,
                sd.values,
            )
        )

    # Planckian radiators, covering the < 4000 K (pure Planckian reference)
    # and 4000-5000 K (Planckian / daylight mixture) regimes.
    shape = SpectralShape(360, 830, 1)
    for T in (2700, 4300, 4700):
        sd = sd_blackbody(T, shape)
        cases.append((f"blackbody_{T}K", f"sd_blackbody({T}, {shape})", sd.wavelengths, sd.values))

    # CIE daylight (> 5000 K regime).
    for T in (6500, 10000):
        xy = CCT_to_xy_CIE_D(T * 1.4388 / 1.4380)
        sd = sd_CIE_illuminant_D_series(xy)
        cases.append((f"daylight_{T}K", f"sd_CIE_illuminant_D_series(CCT_to_xy_CIE_D({T}*1.4388/1.4380))", sd.wavelengths, sd.values))

    # Grid variations.
    # Coarse 10 nm grid.
    d50 = SDS_ILLUMINANTS["D50"]
    wl10 = np.arange(380.0, 781.0, 10.0)
    cases.append(("grid_10nm_D50", "D50 sampled every 10 nm", wl10, d50[wl10]))
    # HPCS 6500 grid (380-1050 nm, 670/349 nm steps), blackbody and LED.
    bb = sd_blackbody(3000, SpectralShape(360, 1100, 1))
    cases.append(("hpcs_grid_blackbody_3000K", "sd_blackbody(3000) on 380-1050 nm / 350 points", HPCS_GRID, bb[HPCS_GRID]))
    led = SDS_ILLUMINANTS["LED-B3"]
    cases.append(
        (
            "hpcs_grid_LED_B3",
            "LED-B3 linearly resampled on 380-1050 nm / 350 points (0 beyond 780)",
            HPCS_GRID,
            np.interp(HPCS_GRID, led.wavelengths, led.values, right=0.0),
        )
    )
    # Irregular, non-uniform grid.
    rng = np.random.default_rng(20180)
    wl_irr = np.sort(np.concatenate([[371.3], rng.uniform(372.0, 800.0, 240), [801.7]]))
    fl11 = SDS_ILLUMINANTS["FL11"]
    cases.append(
        (
            "irregular_grid_FL11",
            "FL11 linearly resampled on 242 random wavelengths 371-802 nm",
            wl_irr,
            np.interp(wl_irr, fl11.wavelengths, fl11.values),
        )
    )
    # Partial coverage (400-700 nm only): np.interp holds the end values.
    fl4 = SDS_ILLUMINANTS["FL4"]
    wl_p = np.arange(400.0, 701.0, 2.5)
    cases.append(
        (
            "partial_400_700_FL4",
            "FL4 on 400-700 nm / 2.5 nm",
            wl_p,
            np.interp(wl_p, fl4.wavelengths, fl4.values),
        )
    )
    # Negated spectrum: colour normalises by Y, so a negative scale must give
    # exactly the positive-scale result.
    fl2 = SDS_ILLUMINANTS["FL2"]
    cases.append(("negated_FL2", "-1 * FL2 (scale invariance incl. sign)", fl2.wavelengths, -fl2.values))
    # Noisy real lamp with negative values (dark-noise like).
    if cases and cases[0][0].startswith("lamp_"):
        _, _, wl, vals = cases[0]
        noise = rng.normal(0.0, 0.01 * float(np.max(vals)), vals.shape)
        cases.append(("noisy_" + cases[0][0], "first lamp + 1% gaussian noise (negative values)", wl, vals + noise))

    return cases


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--tm30-dir", default=DEFAULT_TM30_DIR)
    ap.add_argument("--lamps-dir", default=DEFAULT_LAMPS_DIR)
    args = ap.parse_args()

    sys.path.insert(0, args.tm30_dir)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        import colour  # noqa: F401
        from tm30 import compute_tm30

        cases = build_cases(args.lamps_dir)

    os.makedirs(OUT, exist_ok=True)
    for old in glob.glob(os.path.join(OUT, "*.json")):
        os.remove(old)

    for name, source, wl, vals in cases:
        wl = [float(x) for x in np.asarray(wl, dtype=float)]
        vals = [float(x) for x in np.asarray(vals, dtype=float)]
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            res = compute_tm30(wl, vals)
        expected = {
            "Rf": clean(res["Rf"]),
            "Rg": clean(res["Rg"]),
            "CCT_K": clean(res["CCT_K"]),
            "Duv": clean(res["Duv"]),
            "bins": [
                {k: (v if k == "bin" else clean(v)) for k, v in b.items()}
                for b in res["bins"]
            ],
        }
        doc = {
            "name": name,
            "source": source,
            "generator": "crates/tm30/tools/golden.py (tm30.py + colour-science "
            + __import__("colour").__version__
            + ", numpy "
            + np.__version__
            + ")",
            "wavelengths": wl,
            "values": vals,
            "expected": expected,
        }
        path = os.path.join(OUT, name + ".json")
        with open(path, "w", newline="\n") as f:
            json.dump(doc, f, indent=1, allow_nan=False)
        print(
            f"{name:32s} CCT={expected['CCT_K']:9.2f} Duv={expected['Duv']: .5f} "
            f"Rf={expected['Rf']:6.2f} Rg={expected['Rg']:6.2f}"
        )
    print(f"{len(cases)} cases -> {OUT}")
    return 0


if __name__ == "__main__":
    sys.exit(main())