"""Unit tests for ``memory_core.calibration`` (W2-I, v11).

Pure numerics — no LLM, no DB. Covers the L-BFGS path, the
hand-rolled gradient-descent fallback path (via monkeypatched scipy),
ECE before/after fit, save/load round-trip, and degenerate inputs.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pytest


SRC = Path(__file__).parent.parent / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from memory_core import calibration as calmod  # noqa: E402
from memory_core.calibration import (  # noqa: E402
    PlattCalibrator,
    apply,
    expected_calibration_error,
    fit_platt,
    load,
    save,
)


_FIXTURE = Path(__file__).parent / "fixtures" / "calibration_validation.json"


# ────────────────────────────────────────────────────────────────────
# Fixture helpers
# ────────────────────────────────────────────────────────────────────


@pytest.fixture(scope="module")
def fixture_arrays() -> tuple[np.ndarray, np.ndarray]:
    """Load (scores, labels) from the deterministic validation fixture.

    The fixture is generated by ``scripts/calibrate.py
    --generate-fixture``; we don't regenerate inside the test (the spec
    says the script must regenerate it, not the tests). We do skip
    cleanly with a clear message if the fixture is absent.
    """
    if not _FIXTURE.exists():
        pytest.skip(
            f"calibration validation fixture missing at {_FIXTURE}; "
            f"run: python scripts/calibrate.py --generate-fixture"
        )
    raw = json.loads(_FIXTURE.read_text())
    examples = raw["examples"]
    scores = np.array(
        [float(ex["raw_score"]) for ex in examples], dtype=np.float64
    )
    labels = np.array(
        [int(ex["label"]) for ex in examples], dtype=np.int64
    )
    return scores, labels


# ────────────────────────────────────────────────────────────────────
# Sigmoid math + apply
# ────────────────────────────────────────────────────────────────────


def test_apply_returns_in_unit_interval() -> None:
    """``apply`` always returns a value in [0, 1]."""
    cal = PlattCalibrator(a=-3.0, b=1.0)
    for s in (-100.0, -1.0, 0.0, 0.5, 1.0, 100.0):
        p = apply(cal, s)
        assert 0.0 <= p <= 1.0


def test_apply_clamps_to_unit() -> None:
    """Even with extreme params and inputs we stay in [0, 1]."""
    cal = PlattCalibrator(a=-50.0, b=-50.0)  # ~saturates at 1.0 for s>0
    assert apply(cal, 100.0) == pytest.approx(1.0, abs=1e-6)
    cal2 = PlattCalibrator(a=50.0, b=50.0)   # ~saturates at 0.0
    assert apply(cal2, 100.0) == pytest.approx(0.0, abs=1e-6)


def test_apply_handles_non_finite_score() -> None:
    """NaN/Inf raw score → 0.0 (defensive default)."""
    cal = PlattCalibrator(a=-2.0, b=0.0)
    assert apply(cal, float("nan")) == 0.0
    assert apply(cal, float("inf")) == 0.0
    assert apply(cal, float("-inf")) == 0.0


def test_apply_handles_garbage_score() -> None:
    """Non-numeric raw score → 0.0 instead of raising."""
    cal = PlattCalibrator(a=-2.0, b=0.0)
    assert apply(cal, "not a number") == 0.0  # type: ignore[arg-type]


# ────────────────────────────────────────────────────────────────────
# Fit — happy path
# ────────────────────────────────────────────────────────────────────


def test_fit_monotonic(fixture_arrays: tuple[np.ndarray, np.ndarray]) -> None:
    """After fit, p(score) is monotonic in score.

    Property of any sigmoid with fixed (a, b): increasing s either
    monotonically increases or monotonically decreases the output.
    On positively-correlated data we expect *increasing*.
    """
    scores, labels = fixture_arrays
    cal = fit_platt(scores, labels.astype(np.float64))
    grid = np.linspace(0.0, 1.0, 101)
    probs = np.array([apply(cal, float(s)) for s in grid])
    diffs = np.diff(probs)
    # All diffs same sign (within tiny FP slack). Most data → positive.
    assert np.all(diffs >= -1e-9) or np.all(diffs <= 1e-9), (
        "calibrated probability should be monotonic in raw score"
    )


def test_ece_drops_after_fit(fixture_arrays: tuple[np.ndarray, np.ndarray]) -> None:
    """ECE on calibrated probs is lower than ECE on raw scores."""
    scores, labels = fixture_arrays
    y = labels.astype(np.float64)
    ece_before = expected_calibration_error(scores, y)
    cal = fit_platt(scores, y)
    probs = np.array([apply(cal, float(s)) for s in scores])
    ece_after = expected_calibration_error(probs, y)
    assert ece_after < ece_before, (
        f"expected ECE to drop, got before={ece_before:.4f} "
        f"after={ece_after:.4f}"
    )


def test_ece_under_005(fixture_arrays: tuple[np.ndarray, np.ndarray]) -> None:
    """Calibrated ECE on the validation fixture stays under 0.05."""
    scores, labels = fixture_arrays
    y = labels.astype(np.float64)
    cal = fit_platt(scores, y)
    probs = np.array([apply(cal, float(s)) for s in scores])
    ece = expected_calibration_error(probs, y)
    assert ece <= 0.05, f"post-fit ECE {ece:.4f} exceeds 0.05 budget"


def test_fit_returns_finite_params(fixture_arrays: tuple[np.ndarray, np.ndarray]) -> None:
    """``a`` and ``b`` are finite real numbers after fit."""
    scores, labels = fixture_arrays
    cal = fit_platt(scores, labels.astype(np.float64))
    assert np.isfinite(cal.a)
    assert np.isfinite(cal.b)


def test_fit_a_is_negative_on_positive_correlation(
    fixture_arrays: tuple[np.ndarray, np.ndarray]
) -> None:
    """Higher score ↔ higher label probability → ``a`` should be < 0.

    With p = sigmoid(-(a*s + b)), increasing s must drop the exponent
    for p to grow, which requires a<0.
    """
    scores, labels = fixture_arrays
    cal = fit_platt(scores, labels.astype(np.float64))
    assert cal.a < 0.0


# ────────────────────────────────────────────────────────────────────
# Fit — fallback path (no scipy)
# ────────────────────────────────────────────────────────────────────


def test_no_scipy_path(
    monkeypatch: pytest.MonkeyPatch,
    fixture_arrays: tuple[np.ndarray, np.ndarray],
) -> None:
    """When scipy is unavailable the gradient-descent fallback runs.

    We monkeypatch the module-level flag and minimize handle so the
    fit code takes the fallback branch. The fitted calibrator must
    still hit reasonable ECE on the validation fixture.
    """
    scores, labels = fixture_arrays
    y = labels.astype(np.float64)
    monkeypatch.setattr(calmod, "HAS_SCIPY", False)
    monkeypatch.setattr(calmod, "_scipy_minimize", None)

    cal = fit_platt(scores, y)
    assert np.isfinite(cal.a)
    assert np.isfinite(cal.b)
    probs = np.array([apply(cal, float(s)) for s in scores])
    ece = expected_calibration_error(probs, y)
    assert ece <= 0.06, (
        f"fallback path ECE {ece:.4f} above 0.06 — gradient descent "
        f"is not converging"
    )


def test_fallback_matches_scipy_within_tolerance(
    monkeypatch: pytest.MonkeyPatch,
    fixture_arrays: tuple[np.ndarray, np.ndarray],
) -> None:
    """Both paths converge to similar parameters on the same data."""
    scores, labels = fixture_arrays
    y = labels.astype(np.float64)

    cal_scipy = fit_platt(scores, y)
    monkeypatch.setattr(calmod, "HAS_SCIPY", False)
    monkeypatch.setattr(calmod, "_scipy_minimize", None)
    cal_fallback = fit_platt(scores, y)

    # Compare via predictions — direct param equality is too strict
    # because L-BFGS and GD reach different local optima at machine
    # precision. Predictions on the training set should agree closely.
    p_scipy = np.array([apply(cal_scipy, float(s)) for s in scores])
    p_fb = np.array([apply(cal_fallback, float(s)) for s in scores])
    assert np.max(np.abs(p_scipy - p_fb)) < 0.05


# ────────────────────────────────────────────────────────────────────
# Degenerate inputs
# ────────────────────────────────────────────────────────────────────


def test_fit_handles_degenerate_input_all_positive() -> None:
    """All-positive labels → fit returns a finite calibrator (no NaN).

    Platt smoothing is the whole reason this works — without it the
    BCE gradient blows up.
    """
    scores = np.linspace(0.1, 0.9, 50)
    labels = np.ones(50, dtype=np.float64)
    cal = fit_platt(scores, labels)
    assert np.isfinite(cal.a)
    assert np.isfinite(cal.b)
    # All predictions should lean toward 1, but stay in [0, 1].
    probs = np.array([apply(cal, float(s)) for s in scores])
    assert np.all(probs >= 0.0) and np.all(probs <= 1.0)
    assert float(np.mean(probs)) > 0.6


def test_fit_handles_degenerate_input_all_negative() -> None:
    """All-negative labels → finite calibrator, predictions skewed low."""
    scores = np.linspace(0.1, 0.9, 50)
    labels = np.zeros(50, dtype=np.float64)
    cal = fit_platt(scores, labels)
    assert np.isfinite(cal.a)
    assert np.isfinite(cal.b)
    probs = np.array([apply(cal, float(s)) for s in scores])
    assert np.all(probs >= 0.0) and np.all(probs <= 1.0)
    assert float(np.mean(probs)) < 0.4


def test_fit_rejects_empty_input() -> None:
    """Empty arrays → ValueError, not silent degenerate result."""
    with pytest.raises(ValueError):
        fit_platt(np.array([]), np.array([]))


def test_fit_rejects_shape_mismatch() -> None:
    """Mismatched score/label shapes → ValueError."""
    with pytest.raises(ValueError):
        fit_platt(np.array([0.1, 0.2, 0.3]), np.array([1, 0]))


# ────────────────────────────────────────────────────────────────────
# ECE
# ────────────────────────────────────────────────────────────────────


def test_ece_perfect_calibration_is_zero() -> None:
    """When predictions exactly match labels, ECE == 0."""
    probs = np.array([0.0, 0.0, 1.0, 1.0])
    labels = np.array([0.0, 0.0, 1.0, 1.0])
    assert expected_calibration_error(probs, labels) == pytest.approx(0.0)


def test_ece_uniform_random_is_around_half() -> None:
    """All probs at 0.5 with balanced labels → ECE = 0 (mean prob == mean label)."""
    probs = np.full(100, 0.5)
    labels = np.array([0.0] * 50 + [1.0] * 50)
    assert expected_calibration_error(probs, labels) == pytest.approx(0.0, abs=1e-9)


def test_ece_extreme_miscalibration() -> None:
    """All probs at 1.0 but all labels 0 → ECE = 1.0."""
    probs = np.full(50, 1.0)
    labels = np.zeros(50)
    assert expected_calibration_error(probs, labels) == pytest.approx(1.0)


def test_ece_handles_empty() -> None:
    """Empty arrays return 0 (nothing to be miscalibrated about)."""
    assert expected_calibration_error(np.array([]), np.array([])) == 0.0


# ────────────────────────────────────────────────────────────────────
# Persistence
# ────────────────────────────────────────────────────────────────────


def test_save_load_roundtrip(tmp_path: Path) -> None:
    """Save and reload preserves parameters bit-exact (via JSON repr)."""
    cal = PlattCalibrator(a=-2.71828, b=1.41421)
    out = tmp_path / "cal.json"
    save(cal, out)
    cal2 = load(out)
    assert cal2.a == pytest.approx(cal.a, rel=1e-12)
    assert cal2.b == pytest.approx(cal.b, rel=1e-12)


def test_load_rejects_wrong_kind(tmp_path: Path) -> None:
    """Loading a JSON with kind != 'platt' raises ValueError."""
    p = tmp_path / "bad.json"
    p.write_text(json.dumps({"kind": "isotonic", "a": 0.0, "b": 0.0}))
    with pytest.raises(ValueError):
        load(p)


def test_load_rejects_non_finite(tmp_path: Path) -> None:
    """Loading a JSON with non-finite params raises ValueError.

    JSON does not natively support NaN; we write a bare string and
    check that ``float()`` failure surfaces a ValueError.
    """
    p = tmp_path / "nan.json"
    p.write_text(json.dumps({"kind": "platt", "a": "not a number", "b": 0.0}))
    with pytest.raises(ValueError):
        load(p)


def test_save_creates_parent_directories(tmp_path: Path) -> None:
    """Save creates intermediate directories on demand."""
    cal = PlattCalibrator(a=-1.0, b=0.5)
    out = tmp_path / "deep" / "nested" / "cal.json"
    save(cal, out)
    assert out.exists()
    cal2 = load(out)
    assert cal2.a == pytest.approx(-1.0)
    assert cal2.b == pytest.approx(0.5)
