"""Hue instability diagnostics for near-neutral chroma.

Goal: quantify how sensitive hue (atan2(b,a)) becomes as chroma shrinks, and how
that propagates to the model output.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

from helmlab.spaces.analytical import AnalyticalSpace


def _wrap_pi(x: np.ndarray) -> np.ndarray:
    return (x + np.pi) % (2.0 * np.pi) - np.pi


def _circular_std_rad(angles: np.ndarray) -> float:
    """Circular standard deviation (radians) based on mean resultant length."""
    angles = np.asarray(angles, dtype=np.float64).ravel()
    if angles.size == 0:
        return float("nan")
    c = np.mean(np.cos(angles))
    s = np.mean(np.sin(angles))
    R = float(np.hypot(c, s))
    R = max(1e-12, min(1.0, R))
    return float(np.sqrt(-2.0 * np.log(R)))


def _xyY_to_XYZ(x: np.ndarray, y: np.ndarray, Y: float) -> np.ndarray:
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    X = x * Y / np.maximum(y, 1e-12)
    Z = (1.0 - x - y) * Y / np.maximum(y, 1e-12)
    return np.stack([X, np.full_like(X, Y, dtype=np.float64), Z], axis=-1)


def _raw_lab(space: AnalyticalSpace, XYZ: np.ndarray) -> np.ndarray:
    """Stage 1-3 only: XYZ -> Lab_raw (before hue correction / enrichment)."""
    XYZ = np.asarray(XYZ, dtype=np.float64)
    p = space.params
    LMS = XYZ @ p.M1.T
    LMS_c = np.sign(LMS) * np.abs(LMS) ** p.gamma
    return LMS_c @ p.M2.T


def _raw_Lab_to_XYZ(space: AnalyticalSpace, Lab_raw: np.ndarray) -> np.ndarray:
    """Inverse of stage 1-3 only: Lab_raw -> XYZ."""
    Lab_raw = np.asarray(Lab_raw, dtype=np.float64)
    LMS_c = Lab_raw @ space._M2_inv.T
    inv_gamma = 1.0 / space.params.gamma
    LMS = np.sign(LMS_c) * np.abs(LMS_c) ** inv_gamma
    return LMS @ space._M1_inv.T


@dataclass(frozen=True)
class HueJitterRow:
    C_target: float
    C_raw_mean: float
    hue_std_deg: float
    out_L_std: float
    out_a_std: float
    out_b_std: float


def measure_hue_jitter(
    space: AnalyticalSpace,
    *,
    base_hue_deg: float = 240.0,
    L_raw: float = 0.6,
    C_targets: np.ndarray | None = None,
    n_samples: int = 500,
    noise_xyz_rel: float = 1e-6,
) -> list[HueJitterRow]:
    """Generate near-neutral chromaticity perturbations around D65 and measure jitter.

    We generate a base point by specifying *raw* Lab (stage 1-3 output) with
    very small chroma, then invert back to XYZ. To probe near-neutral hue
    instability, we add tiny isotropic noise in XYZ and measure:
      - raw hue circular std (atan2(b,a))
      - output coordinate variance after the full pipeline
    """
    if C_targets is None:
        C_targets = np.logspace(-6, -2, 9)
    C_targets = np.asarray(C_targets, dtype=np.float64)

    rng = np.random.default_rng(0)
    theta0 = math.radians(base_hue_deg)
    rows: list[HueJitterRow] = []

    for C_t in C_targets:
        Lab_base = np.array([L_raw, float(C_t) * math.cos(theta0), float(C_t) * math.sin(theta0)], dtype=np.float64)
        XYZ_base = _raw_Lab_to_XYZ(space, Lab_base)
        sigma = float(noise_xyz_rel) * float(np.max(np.abs(XYZ_base)) + 1e-12)
        XYZ = XYZ_base + rng.normal(0.0, sigma, size=(n_samples, 3))

        Lab_raw = _raw_lab(space, XYZ)
        a_raw = Lab_raw[:, 1]
        b_raw = Lab_raw[:, 2]
        C_raw = np.hypot(a_raw, b_raw)
        h_raw = np.arctan2(b_raw, a_raw)
        hue_std = _circular_std_rad(_wrap_pi(h_raw - theta0))

        out = space.from_XYZ(XYZ)
        rows.append(HueJitterRow(
            C_target=float(C_t),
            C_raw_mean=float(np.mean(C_raw)),
            hue_std_deg=float(np.degrees(hue_std)),
            out_L_std=float(np.std(out[:, 0])),
            out_a_std=float(np.std(out[:, 1])),
            out_b_std=float(np.std(out[:, 2])),
        ))

    return rows


def print_table(rows: list[HueJitterRow], title: str) -> None:
    print(f"\n== {title} ==")
    print("C_target    C_raw_mean  hue_std(deg)   std(L)     std(a)     std(b)")
    for r in rows:
        print(f"{r.C_target:9.1e}  {r.C_raw_mean:9.2e}    {r.hue_std_deg:9.2f}  "
              f"{r.out_L_std:9.2e}  {r.out_a_std:9.2e}  {r.out_b_std:9.2e}")


def main() -> None:
    base = AnalyticalSpace(neutral_correction=True, ab_rotate_deg=-28.2)
    gated = AnalyticalSpace(neutral_correction=True, ab_rotate_deg=-28.2, hue_lightness_c0=0.1)

    rows_base = measure_hue_jitter(base)
    rows_gated = measure_hue_jitter(gated)

    print_table(rows_base, "baseline (paper config)")
    print_table(rows_gated, "hue_lightness_c0=0.1")


if __name__ == "__main__":
    main()
