"""
ANSI/IES TM-30-18 computation from an HPCS 6500 spectrum.

Wraps colour-science to derive the fidelity index Rf, gamut index Rg, and the
per-hue-bin data needed to draw a TM-30 color vector graphic (CVG): the test
and reference chromaticity averages in CAM02-UCS a'b' (raw and normalized so
the reference gamut is the unit circle), plus per-bin chroma shift, hue shift,
and local fidelity.

The instrument spectrum is relative; TM-30 is scale-invariant, so no absolute
calibration is required. Only the 380-780 nm range is used.
"""

import warnings

import numpy as np

try:
    # colour warns at import time about optional scipy/matplotlib features;
    # the TM-30 path used here needs neither.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        import colour
        from colour.quality import colour_fidelity_index_ANSIIESTM3018
except ImportError as e:  # pragma: no cover
    raise ImportError(
        "TM-30 computation requires colour-science; run `uv sync`"
    ) from e

VISIBLE_MIN_NM = 380
VISIBLE_MAX_NM = 780
N_BINS = 16


def compute_tm30(wavelengths_nm, values):
    """Compute TM-30-18 quantities from a measured spectrum.

    Args:
        wavelengths_nm: sequence of wavelengths in nm (any step; 380-1050
            from the HPCS 6500 is fine, the IR tail is discarded).
        values: spectral power at each wavelength (relative units).

    Returns dict:
        Rf, Rg: floats
        CCT_K, Duv: floats (as derived by the TM-30 reference selection)
        bins: list of 16 dicts, hue bin 1..16, each with
            ref_a, ref_b, test_a, test_b        raw CAM02-UCS bin averages
            ref_a_norm ... test_b_norm          normalized (reference = unit circle)
            Rcs_pct                             chroma shift, percent
            Rhs                                 hue shift (rad, CAM02-UCS)
            Rf_h                                local fidelity for the bin
    """
    wl = np.asarray(wavelengths_nm, dtype=float)
    vals = np.asarray(values, dtype=float)
    mask = (wl >= VISIBLE_MIN_NM) & (wl <= VISIBLE_MAX_NM)
    if mask.sum() < 10:
        raise ValueError("spectrum does not cover the visible range")

    # Resample onto a uniform 1 nm grid: the instrument's ~1.92 nm grid is
    # non-uniform in colour's eyes and triggers interpolator paths that need
    # scipy; a uniform grid keeps the pipeline on plain numpy.
    grid = np.arange(VISIBLE_MIN_NM, VISIBLE_MAX_NM + 1, 1, dtype=float)
    resampled = np.interp(grid, wl[mask], vals[mask])

    sd = colour.SpectralDistribution(dict(zip(grid, resampled)))
    # colour emits informational "aligning shape" runtime warnings while it
    # adapts observers/illuminants to our 1 nm grid; not actionable.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        spec = colour_fidelity_index_ANSIIESTM3018(sd, additional_data=True)

    averages_test = np.asarray(spec.averages_test)
    averages_reference = np.asarray(spec.averages_reference)
    average_norms = np.asarray(spec.average_norms)
    R_cs = np.asarray(spec.R_cs)
    R_hs = np.asarray(spec.R_hs)

    # Local fidelity per hue bin: mean of the per-CES R_f over the samples
    # assigned to the bin (spec.bins holds the bin index of each CES).
    sample_bins = np.asarray(spec.bins)
    R_s = np.asarray(spec.R_s)
    Rf_h = np.full(N_BINS, np.nan)
    for j in range(N_BINS):
        members = R_s[sample_bins == j]
        if members.size:
            Rf_h[j] = members.mean()

    bins = []
    for j in range(N_BINS):
        norm = average_norms[j] if average_norms[j] else 1.0
        bins.append(
            {
                "bin": j + 1,
                "ref_a": float(averages_reference[j, 0]),
                "ref_b": float(averages_reference[j, 1]),
                "test_a": float(averages_test[j, 0]),
                "test_b": float(averages_test[j, 1]),
                "ref_a_norm": float(averages_reference[j, 0] / norm),
                "ref_b_norm": float(averages_reference[j, 1] / norm),
                "test_a_norm": float(averages_test[j, 0] / norm),
                "test_b_norm": float(averages_test[j, 1] / norm),
                "Rcs_pct": float(R_cs[j]),
                "Rhs": float(R_hs[j]),
                "Rf_h": float(Rf_h[j]),
            }
        )

    return {
        "Rf": float(spec.R_f),
        "Rg": float(spec.R_g),
        "CCT_K": float(spec.CCT),
        "Duv": float(spec.D_uv),
        "bins": bins,
    }


if __name__ == "__main__":
    # Sanity check against a known illuminant.
    sd = colour.SDS_ILLUMINANTS["FL2"]
    wl = sd.wavelengths
    result = compute_tm30(wl, sd.values)
    print(f"FL2: Rf={result['Rf']:.1f} Rg={result['Rg']:.1f} "
          f"CCT={result['CCT_K']:.0f} K Duv={result['Duv']:.4f}")
    for b in result["bins"][:4]:
        print(b)