"""Gradient benchmark for UI-focused interpolation.

Generates random gradients and compares interpolation methods using CIEDE2000
as an evaluation metric (computed on the XYZ of the *displayed* sRGB steps).

Reported metrics:
  - step uniformity: coefficient of variation of per-step ΔE00
  - monotonicity: count of violations of distance-from-start monotonic increase
  - smoothness: mean |Δ(step ΔE00)| normalized by mean step ΔE00
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from helmlab.metrics.delta_e import delta_e_2000
from helmlab.spaces.analytical import AnalyticalParams, AnalyticalSpace
from helmlab.spaces.base import ColorSpace
from helmlab.spaces.oklch import OKLCH
from helmlab.utils.conversions import Lab_to_XYZ, XYZ_to_Lab
from helmlab.utils.gamut import gamut_map
from helmlab.utils.srgb_convert import XYZ_to_sRGB, clamp_srgb, sRGB_to_XYZ


class CIELabSpace(ColorSpace):
    name = "CIELab"

    def from_XYZ(self, XYZ: np.ndarray) -> np.ndarray:
        return XYZ_to_Lab(XYZ)

    def to_XYZ(self, coords: np.ndarray) -> np.ndarray:
        return Lab_to_XYZ(coords)


def _hsv_to_rgb(h: float, s: float, v: float) -> np.ndarray:
    """HSV in [0,1] -> sRGB in [0,1] (gamma-encoded)."""
    h = float(h) % 1.0
    s = float(np.clip(s, 0.0, 1.0))
    v = float(np.clip(v, 0.0, 1.0))
    i = int(h * 6.0)
    f = h * 6.0 - i
    p = v * (1.0 - s)
    q = v * (1.0 - f * s)
    t = v * (1.0 - (1.0 - f) * s)
    i = i % 6
    if i == 0:
        r, g, b = v, t, p
    elif i == 1:
        r, g, b = q, v, p
    elif i == 2:
        r, g, b = p, v, t
    elif i == 3:
        r, g, b = p, q, v
    elif i == 4:
        r, g, b = t, p, v
    else:
        r, g, b = v, p, q
    return np.array([r, g, b], dtype=np.float64)


def _interp_gradient_XYZ(
    space: ColorSpace,
    XYZ0: np.ndarray,
    XYZ1: np.ndarray,
    *,
    steps: int = 32,
    gamut: str = "srgb",
) -> np.ndarray:
    """Interpolate linearly in `space`, gamut-map, return XYZ steps."""
    c0 = space.from_XYZ(XYZ0)
    c1 = space.from_XYZ(XYZ1)
    ts = np.linspace(0.0, 1.0, steps, dtype=np.float64)
    coords = (1.0 - ts)[:, None] * c0[None, :] + ts[:, None] * c1[None, :]
    coords = gamut_map(coords, space, gamut=gamut)
    return space.to_XYZ(coords)


@dataclass(frozen=True)
class GradientMetrics:
    uniformity_cv: float
    smoothness: float
    monotonic_violations: int
    max_backtrack: float


def _metrics_for_gradient(XYZ_steps: np.ndarray, tol: float = 1e-6) -> GradientMetrics:
    XYZ_steps = np.asarray(XYZ_steps, dtype=np.float64)
    d_step = delta_e_2000(XYZ_steps[:-1], XYZ_steps[1:])
    mean_d = float(np.mean(d_step)) if d_step.size else float("nan")
    uniformity_cv = float(np.std(d_step) / mean_d) if mean_d > 0.0 else float("nan")
    smoothness = (
        float(np.mean(np.abs(np.diff(d_step))) / mean_d)
        if mean_d > 0.0 and d_step.size > 1
        else float("nan")
    )

    XYZ0 = XYZ_steps[0]
    XYZ0_rep = np.repeat(XYZ0[None, :], XYZ_steps.shape[0], axis=0)
    d_from_start = delta_e_2000(XYZ0_rep, XYZ_steps)
    diffs = np.diff(d_from_start)
    monotonic_violations = int(np.sum(diffs < -tol))
    max_backtrack = float(np.min(diffs)) if diffs.size else 0.0

    return GradientMetrics(
        uniformity_cv=uniformity_cv,
        smoothness=smoothness,
        monotonic_violations=monotonic_violations,
        max_backtrack=max_backtrack,
    )


def _make_endpoints(rng: np.random.Generator, kind: str) -> tuple[np.ndarray, np.ndarray]:
    """Return (srgb0, srgb1), gamma-encoded in [0,1]."""
    if kind == "achromatic_chromatic":
        g = float(rng.uniform(0.05, 0.95))
        srgb0 = np.array([g, g, g], dtype=np.float64)
        h = float(rng.uniform(0.0, 1.0))
        s = float(rng.uniform(0.45, 1.0))
        v = float(rng.uniform(0.15, 0.95))
        srgb1 = _hsv_to_rgb(h, s, v)
        return srgb0, srgb1

    if kind == "hue_ramp":
        h0 = float(rng.uniform(0.0, 1.0))
        dh = float(rng.uniform(60.0 / 360.0, 180.0 / 360.0))
        h1 = (h0 + dh) % 1.0
        s = float(rng.uniform(0.35, 1.0))
        v = float(rng.uniform(0.20, 0.95))
        return _hsv_to_rgb(h0, s, v), _hsv_to_rgb(h1, s, v)

    raise ValueError("kind must be 'achromatic_chromatic' or 'hue_ramp'")


def run_benchmark(*, n_gradients: int = 200, steps: int = 32, seed: int = 0) -> None:
    rng = np.random.default_rng(seed)

    base_params = AnalyticalParams.load("src/helmlab/data/analytical_params.json")
    helmlab_base = AnalyticalSpace(base_params, neutral_correction=True, ab_rotate_deg=-28.2)

    # Near-neutral improved candidate (from local search):
    # cp pre-cs gating + Lh gate + retuned metric.
    improved_params = AnalyticalParams.from_dict(base_params.to_dict())
    improved_params.dist_power = 1.1321138755705196
    improved_params.dist_wC = 0.8887074032076554
    improved_params.dist_compress = 10.0
    improved_params.dist_post_power = 0.8930521602495525
    helmlab_improved = AnalyticalSpace(
        improved_params,
        neutral_correction=True,
        ab_rotate_deg=-28.2,
        chroma_power_c0=0.05,
        chroma_power_gate_source="pre_cs",
        Lh_c0=0.50,
    )

    methods: list[tuple[str, ColorSpace]] = [
        ("Helmlab (baseline)", helmlab_base),
        ("Helmlab (NN-improved)", helmlab_improved),
        ("OKLab (linear)", OKLCH()),
        ("CIELab (linear)", CIELabSpace()),
    ]

    kinds = ["achromatic_chromatic", "hue_ramp"]
    per_method: dict[str, list[GradientMetrics]] = {name: [] for name, _ in methods}

    for i in range(n_gradients):
        kind = kinds[i % len(kinds)]
        srgb0, srgb1 = _make_endpoints(rng, kind)
        srgb0 = clamp_srgb(srgb0)
        srgb1 = clamp_srgb(srgb1)
        XYZ0 = sRGB_to_XYZ(srgb0)
        XYZ1 = sRGB_to_XYZ(srgb1)

        for name, space in methods:
            XYZ_steps = _interp_gradient_XYZ(space, XYZ0, XYZ1, steps=steps, gamut="srgb")
            # Quantize to realizable sRGB and re-evaluate in XYZ.
            srgb_steps = clamp_srgb(XYZ_to_sRGB(XYZ_steps))
            XYZ_steps = sRGB_to_XYZ(srgb_steps)
            per_method[name].append(_metrics_for_gradient(XYZ_steps))

    def summarize(vals: list[float]) -> tuple[float, float]:
        arr = np.array(vals, dtype=np.float64)
        return float(np.nanmean(arr)), float(np.nanpercentile(arr, 90.0))

    header = (
        f"{'Method':24s}  "
        f"{'uniform_CV(mean/p90)':>20s}  "
        f"{'smooth(mean/p90)':>18s}  "
        f"{'viol(mean/p90)':>15s}  "
        f"{'backtrack_p10':>13s}"
    )

    print(f"\nGradients: n={n_gradients}, steps={steps}, seed={seed}")
    print("Eval metric: CIEDE2000 (ΔE00) on XYZ")
    print(header)
    print("-" * len(header))

    for name, _space in methods:
        ms = per_method[name]
        u_mean, u_p90 = summarize([m.uniformity_cv for m in ms])
        s_mean, s_p90 = summarize([m.smoothness for m in ms])
        v_mean, v_p90 = summarize([m.monotonic_violations for m in ms])
        back_p10 = float(np.nanpercentile([m.max_backtrack for m in ms], 10.0))
        print(
            f"{name:24s}  "
            f"{u_mean:7.3f}/{u_p90:7.3f}  "
            f"{s_mean:7.3f}/{s_p90:7.3f}  "
            f"{v_mean:6.2f}/{v_p90:6.2f}  "
            f"{back_p10:13.4f}"
        )


def main() -> None:
    run_benchmark()


if __name__ == "__main__":
    main()
