"""Small Pareto-style search for near-neutral stabilization trade-offs.

Objective (UI-focused):
  - minimize COMBVD STRESS (global psychophysics)
  - minimize achromatic_chromatic zone STRESS (local feedback subset)
Constraint example:
  - COMBVD <= 24.0

This script does a small grid search over a couple of *invertible* stabilization
knobs (no re-training):
  - chroma_power_c0 (gate chroma-power nonlinearity near neutral)
  - Lh_c0 (gate hue-dependent L correction near neutral)

Metric parameters are kept fixed to the current NN-improved candidate. If you
want more exhaustive search, expand grids or add random sampling.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from helmlab.data.combvd import load_combvd
from helmlab.diagnostics.near_neutral import load_feedback_pairs
from helmlab.metrics.stress import stress
from helmlab.spaces.analytical import AnalyticalParams, AnalyticalSpace


def _pearson_r(x: np.ndarray, y: np.ndarray) -> float:
    x = np.asarray(x, dtype=np.float64).ravel()
    y = np.asarray(y, dtype=np.float64).ravel()
    if x.size < 2:
        return float("nan")
    sx = float(np.std(x))
    sy = float(np.std(y))
    if sx == 0.0 or sy == 0.0:
        return float("nan")
    return float(np.corrcoef(x, y)[0, 1])


def _metrics(DV: np.ndarray, DE: np.ndarray) -> tuple[float, float]:
    DV = np.asarray(DV, dtype=np.float64).ravel()
    DE = np.asarray(DE, dtype=np.float64).ravel()
    return float(stress(DV, DE)), _pearson_r(DE, DV)


@dataclass(frozen=True)
class Row:
    chroma_power_c0: float
    Lh_c0: float
    combvd_stress: float
    ac_stress: float
    ac_r: float


def _pareto_front(rows: list[Row]) -> list[Row]:
    """Return non-dominated rows (minimize combvd_stress, ac_stress)."""
    front: list[Row] = []
    for r in rows:
        dominated = False
        for q in rows:
            if (q.combvd_stress <= r.combvd_stress and q.ac_stress <= r.ac_stress) and (
                q.combvd_stress < r.combvd_stress or q.ac_stress < r.ac_stress
            ):
                dominated = True
                break
        if not dominated:
            front.append(r)
    front.sort(key=lambda x: (x.combvd_stress, x.ac_stress))
    return front


def main() -> None:
    combvd = load_combvd()
    fb = load_feedback_pairs()
    mask_ac = fb.mask_zone("achromatic_chromatic")

    base_params = AnalyticalParams.load("src/helmlab/data/analytical_params.json")

    # Fixed metric (NN-improved candidate)
    params = AnalyticalParams.from_dict(base_params.to_dict())
    params.dist_power = 1.1321138755705196
    params.dist_wC = 0.8887074032076554
    params.dist_compress = 10.0
    params.dist_post_power = 0.8930521602495525

    gate_kind = "rational2"
    combvd_max = 24.0

    cp_grid = np.array([0.00, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07], dtype=np.float64)
    Lh_grid = np.array([0.00, 0.05, 0.10, 0.20, 0.30, 0.50, 1.00, 2.00], dtype=np.float64)

    rows: list[Row] = []
    for cp_c0 in cp_grid:
        for Lh_c0 in Lh_grid:
            space = AnalyticalSpace(
                params,
                neutral_correction=True,
                ab_rotate_deg=-28.2,
                chroma_power_c0=float(cp_c0),
                chroma_power_gate_source="pre_cs",
                Lh_c0=float(Lh_c0),
                near_neutral_gate=gate_kind,
            )
            DE_c = space.distance(combvd["XYZ_1"], combvd["XYZ_2"])
            s_c, _ = _metrics(combvd["DV"], DE_c)
            DE_f = space.distance(fb.XYZ_1, fb.XYZ_2)
            s_ac, r_ac = _metrics(fb.DV[mask_ac], DE_f[mask_ac])
            rows.append(
                Row(
                    chroma_power_c0=float(cp_c0),
                    Lh_c0=float(Lh_c0),
                    combvd_stress=float(s_c),
                    ac_stress=float(s_ac),
                    ac_r=float(r_ac),
                )
            )

    feasible = [r for r in rows if r.combvd_stress <= combvd_max]
    feasible.sort(key=lambda r: r.ac_stress)
    front = _pareto_front(feasible)

    print("\n== Search setup ==")
    print(f"gate={gate_kind!r}, chroma_power_gate_source='pre_cs', COMBVD_max={combvd_max}")
    print(f"cp_grid={cp_grid.tolist()}")
    print(f"Lh_grid={Lh_grid.tolist()}")

    print("\n== Best feasible rows (sorted by AC STRESS) ==")
    print(f"{'cp_c0':>6s}  {'Lh_c0':>6s}  {'COMBVD':>8s}  {'AC':>8s}  {'r_ac':>7s}")
    print("-" * 50)
    for r in feasible[:12]:
        print(f"{r.chroma_power_c0:6.2f}  {r.Lh_c0:6.2f}  {r.combvd_stress:8.3f}  {r.ac_stress:8.3f}  {r.ac_r:+7.3f}")

    print("\n== Pareto front (non-dominated) ==")
    for r in front[:20]:
        print(f"cp_c0={r.chroma_power_c0:.2f}, Lh_c0={r.Lh_c0:.2f} -> COMBVD={r.combvd_stress:.3f}, AC={r.ac_stress:.3f}")
    if len(front) > 20:
        print(f"... ({len(front)} total)")


if __name__ == "__main__":
    main()
