"""Compare near-neutral gate functions for Helmlab stabilization knobs.

This is a lightweight, reproducible diagnostic. It prints COMBVD + local
human-feedback metrics for a few configurations while varying:
  - near_neutral_gate: rational2 | rational1 | smoothstep | exponential | logistic

The goal is to answer: "Which gate shape gives the best trade-off?"
"""

from __future__ import annotations

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)


def _print_header(title: str) -> None:
    print(f"\n== {title} ==")
    print(f"{'gate':11s}  {'COMBVD':>8s}  {'FB':>8s}  {'AC':>8s}  {'r_ac':>7s}")
    print("-" * 50)


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")

    gate_kinds = ["rational2", "rational1", "smoothstep", "exponential", "logistic"]

    # Scenario A: isolate hue-lightness gating impact.
    _print_header("Scenario A: hue_lightness_c0=0.1 (baseline metric)")
    for gate in gate_kinds:
        space = AnalyticalSpace(
            base_params,
            neutral_correction=True,
            ab_rotate_deg=-28.2,
            hue_lightness_c0=0.1,
            near_neutral_gate=gate,
        )
        DE_c = space.distance(combvd["XYZ_1"], combvd["XYZ_2"])
        DE_f = space.distance(fb.XYZ_1, fb.XYZ_2)
        s_c, _ = _metrics(combvd["DV"], DE_c)
        s_f, _ = _metrics(fb.DV, DE_f)
        s_ac, r_ac = _metrics(fb.DV[mask_ac], DE_f[mask_ac])
        print(f"{gate:11s}  {s_c:8.3f}  {s_f:8.3f}  {s_ac:8.3f}  {r_ac:+7.3f}")

    # Scenario B: NN-improved candidate (metric retuned + cp/Lh near-neutral gates).
    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

    _print_header("Scenario B: NN-improved (cp pre_cs + Lh gate + retuned metric)")
    for gate in gate_kinds:
        space = 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,
            near_neutral_gate=gate,
        )
        DE_c = space.distance(combvd["XYZ_1"], combvd["XYZ_2"])
        DE_f = space.distance(fb.XYZ_1, fb.XYZ_2)
        s_c, _ = _metrics(combvd["DV"], DE_c)
        s_f, _ = _metrics(fb.DV, DE_f)
        s_ac, r_ac = _metrics(fb.DV[mask_ac], DE_f[mask_ac])
        print(f"{gate:11s}  {s_c:8.3f}  {s_f:8.3f}  {s_ac:8.3f}  {r_ac:+7.3f}")


if __name__ == "__main__":
    main()
