"""Near-neutral debugging utilities.

This module is intentionally lightweight and reproducible:
- Loads COMBVD and local human feedback pairs
- Computes near-neutral masks based on *raw* chroma (pre-hue stages)
- Runs stage-by-stage ablations and prints STRESS/correlation summaries
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import numpy as np

from helmlab.data.combvd import load_combvd
from helmlab.metrics.stress import stress
from helmlab.spaces.analytical import AnalyticalParams, AnalyticalSpace
from helmlab.utils.srgb_convert import hex_to_srgb, sRGB_to_XYZ


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 _raw_lab(params: AnalyticalParams, XYZ: np.ndarray) -> np.ndarray:
    """Stage 1-3 only: XYZ -> Lab_raw (before hue correction / enrichment)."""
    XYZ = np.asarray(XYZ, dtype=np.float64)
    LMS = XYZ @ params.M1.T
    LMS_c = np.sign(LMS) * np.abs(LMS) ** params.gamma
    return LMS_c @ params.M2.T


def _raw_chroma(params: AnalyticalParams, XYZ: np.ndarray) -> np.ndarray:
    Lab = _raw_lab(params, XYZ)
    return np.sqrt(Lab[..., 1] ** 2 + Lab[..., 2] ** 2)


@dataclass(frozen=True)
class FeedbackData:
    XYZ_1: np.ndarray
    XYZ_2: np.ndarray
    DV: np.ndarray
    zone: list[str]

    def mask_zone(self, zone_name: str) -> np.ndarray:
        return np.array([z == zone_name for z in self.zone], dtype=bool)


def load_feedback_pairs(path: str | Path = "data/human_feedback.json") -> FeedbackData:
    path = Path(path)
    with path.open() as f:
        raw = json.load(f)

    judgements = raw.get("judgements", [])

    # Average DV per unique pair; keep zone label.
    pair_info: dict[tuple[str, str], dict] = {}
    for j in judgements:
        key = (j["hex1"], j["hex2"])
        info = pair_info.setdefault(key, {"dvs": [], "zone": j.get("metadata", {}).get("zone", "unknown")})
        info["dvs"].append(float(j["perceived_dv"]))

    pairs = list(pair_info.items())
    XYZ_1 = np.array([sRGB_to_XYZ(hex_to_srgb(h1)) for (h1, _), _info in pairs], dtype=np.float64)
    XYZ_2 = np.array([sRGB_to_XYZ(hex_to_srgb(h2)) for (_, h2), _info in pairs], dtype=np.float64)
    DV = np.array([float(np.mean(_info["dvs"])) for _k, _info in pairs], dtype=np.float64)
    zone = [_info["zone"] for _k, _info in pairs]
    return FeedbackData(XYZ_1=XYZ_1, XYZ_2=XYZ_2, DV=DV, zone=zone)


def near_neutral_mask_raw(
    params: AnalyticalParams,
    XYZ_1: np.ndarray,
    XYZ_2: np.ndarray,
    c_thresh: float = 0.02,
    mode: str = "avg",
) -> np.ndarray:
    """Mask pairs considered near-neutral based on *raw* chroma.

    mode:
      - "avg": (C1+C2)/2 < thresh
      - "max": max(C1,C2) < thresh
    """
    C1 = _raw_chroma(params, XYZ_1)
    C2 = _raw_chroma(params, XYZ_2)
    if mode == "avg":
        C = 0.5 * (C1 + C2)
    elif mode == "max":
        C = np.maximum(C1, C2)
    else:
        raise ValueError("mode must be 'avg' or 'max'")
    return C < float(c_thresh)


def _metrics(DV: np.ndarray, DE: np.ndarray, mask: np.ndarray | None = None) -> dict:
    DV = np.asarray(DV, dtype=np.float64).ravel()
    DE = np.asarray(DE, dtype=np.float64).ravel()
    if mask is not None:
        mask = np.asarray(mask, dtype=bool).ravel()
        DV = DV[mask]
        DE = DE[mask]
    n = int(DV.size)
    if n == 0:
        return {"n": 0, "stress": float("nan"), "r": float("nan")}
    return {"n": n, "stress": float(stress(DV, DE)), "r": _pearson_r(DE, DV)}


def _clone_params(p: AnalyticalParams) -> AnalyticalParams:
    # AnalyticalParams mostly stores numpy arrays; do a conservative deep-ish copy.
    d = p.to_dict()
    return AnalyticalParams.from_dict(d)


def _ablate_params(params: AnalyticalParams, stage: str) -> AnalyticalParams:
    p = _clone_params(params)
    if stage == "hue_correction":
        for k in ("hue_cos1", "hue_sin1", "hue_cos2", "hue_sin2", "hue_cos3", "hue_sin3", "hue_cos4", "hue_sin4"):
            setattr(p, k, 0.0)
        return p
    if stage == "hk":
        for k in ("hk_weight", "hk_power", "hk_hue_mod", "hk_sin1", "hk_cos2", "hk_sin2", "hk_weight_S", "hk_power_S", "hk_hue_S"):
            setattr(p, k, 0.0)
        return p
    if stage == "L_corr":
        for k in ("L_corr_p1", "L_corr_p2", "L_corr_p3", "Lh_cos1", "Lh_sin1"):
            setattr(p, k, 0.0)
        return p
    if stage == "dark_L":
        for k in ("lp_dark", "lp_dark_hcos", "lp_dark_hsin", "lp_dark_S", "lp_dark_S2"):
            setattr(p, k, 0.0)
        return p
    if stage == "chroma_scale":
        for k in ("cs_cos1", "cs_sin1", "cs_cos2", "cs_sin2", "cs_cos3", "cs_sin3", "cs_cos4", "cs_sin4", "cs_S_lin", "cs_S_quad"):
            setattr(p, k, 0.0)
        return p
    if stage == "chroma_power":
        for k in ("cp_cos1", "cp_sin1", "cp_cos2", "cp_sin2"):
            setattr(p, k, 0.0)
        return p
    if stage == "L_chroma":
        for k in ("lc1", "lc2", "lc_S_lin", "lc_S_quad"):
            setattr(p, k, 0.0)
        return p
    if stage == "HLC":
        for k in ("hlc_cos1", "hlc_sin1", "hlc_cos2", "hlc_sin2"):
            setattr(p, k, 0.0)
        return p
    if stage == "hue_lightness":
        for k in ("hl_cos1", "hl_sin1", "hl_cos2", "hl_sin2", "hl_S_lin"):
            setattr(p, k, 0.0)
        return p
    if stage == "metric":
        # Reset to Euclidean.
        for k in ("dist_power", "dist_wC", "dist_nl", "dist_sat", "dist_compress", "dist_linear", "dist_post_power",
                  "dist_sl", "dist_sc",
                  "dist_sl_hcos1", "dist_sl_hsin1", "dist_sl_hcos2", "dist_sl_hsin2",
                  "dist_sc_hcos1", "dist_sc_hsin1", "dist_sc_hcos2", "dist_sc_hsin2"):
            setattr(p, k, 0.0)
        p.dist_power = 1.0
        p.dist_wC = 1.0
        p.dist_post_power = 1.0
        return p
    raise ValueError(f"Unknown stage: {stage!r}")


@dataclass(frozen=True)
class EvalConfig:
    c_thresh: float = 0.02
    raw_mode: str = "avg"
    neutral_correction: bool = True
    ab_rotate_deg: float = -28.2

    # Gating knobs (all default 0 => disabled)
    near_neutral_gate: str = "rational2"
    hue_correction_c0: float = 0.0
    hk_hue_c0: float = 0.0
    Lh_c0: float = 0.0
    dark_L_hue_c0: float = 0.0
    chroma_power_c0: float = 0.0
    chroma_power_gate_source: str = "post"
    hue_lightness_c0: float = 0.0


def evaluate_variants(
    params: AnalyticalParams,
    stages: list[str],
    cfg: EvalConfig | None = None,
) -> dict[str, dict]:
    cfg = cfg or EvalConfig()
    combvd = load_combvd()
    fb = load_feedback_pairs()

    mask_c_nn = near_neutral_mask_raw(params, combvd["XYZ_1"], combvd["XYZ_2"], cfg.c_thresh, mode=cfg.raw_mode)
    mask_f_nn = near_neutral_mask_raw(params, fb.XYZ_1, fb.XYZ_2, cfg.c_thresh, mode=cfg.raw_mode)
    mask_ac = fb.mask_zone("achromatic_chromatic")
    mask_ac_nn = mask_ac & mask_f_nn

    results: dict[str, dict] = {}

    def make_space(p: AnalyticalParams) -> AnalyticalSpace:
        return AnalyticalSpace(
            p,
            neutral_correction=cfg.neutral_correction,
            ab_rotate_deg=cfg.ab_rotate_deg,
            chroma_power_c0=cfg.chroma_power_c0,
            chroma_power_gate_source=cfg.chroma_power_gate_source,
            hue_lightness_c0=cfg.hue_lightness_c0,
            near_neutral_gate=cfg.near_neutral_gate,
            hue_correction_c0=cfg.hue_correction_c0,
            hk_hue_c0=cfg.hk_hue_c0,
            Lh_c0=cfg.Lh_c0,
            dark_L_hue_c0=cfg.dark_L_hue_c0,
        )

    def eval_space(label: str, space: AnalyticalSpace):
        DE_c = space.distance(combvd["XYZ_1"], combvd["XYZ_2"])
        DE_f = space.distance(fb.XYZ_1, fb.XYZ_2)
        results[label] = {
            "COMBVD": _metrics(combvd["DV"], DE_c),
            "COMBVD_near_neutral": _metrics(combvd["DV"], DE_c, mask=mask_c_nn),
            "FB": _metrics(fb.DV, DE_f),
            "FB_near_neutral": _metrics(fb.DV, DE_f, mask=mask_f_nn),
            "AC": _metrics(fb.DV, DE_f, mask=mask_ac),
            "AC_near_neutral": _metrics(fb.DV, DE_f, mask=mask_ac_nn),
        }

    # Baseline
    eval_space("baseline", make_space(params))

    # Stage ablations
    for stage in stages:
        p_ab = _ablate_params(params, stage)
        eval_space(f"abl_{stage}", make_space(p_ab))

    return results


def print_report(results: dict[str, dict]) -> None:
    def fmt(m: dict) -> str:
        if m["n"] == 0 or not np.isfinite(m["stress"]):
            return "n=0  STRESS=nan  r=nan"
        return f"n={m['n']:4d}  STRESS={m['stress']:6.2f}  r={m['r']:+.3f}"

    blocks = ["COMBVD", "COMBVD_near_neutral", "FB", "FB_near_neutral", "AC", "AC_near_neutral"]
    for label, row in results.items():
        print(f"\n== {label} ==")
        for b in blocks:
            print(f"  {b:18s}  {fmt(row[b])}")


def main() -> None:
    params = AnalyticalParams.load(Path(__file__).resolve().parents[1] / "data" / "analytical_params.json")
    stages = [
        "hue_correction",
        "hk",
        "L_corr",
        "dark_L",
        "chroma_scale",
        "chroma_power",
        "L_chroma",
        "HLC",
        "hue_lightness",
        "metric",
    ]
    cfg = EvalConfig(c_thresh=0.02, raw_mode="avg")
    results = evaluate_variants(params, stages, cfg=cfg)
    print_report(results)


if __name__ == "__main__":
    main()
