"""Tests for ModelTransferCalibrator — Platt scaling transfer between judge models."""
import sys
import numpy as np
import pytest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.calibration.model_transfer_calibrator import ModelTransferCalibrator
from llm_guard.calibration.federated_calibrator import CalibParams


# ── Fixtures ──────────────────────────────────────────────────────────────────

def make_source_params():
    """Simple identity-like isotonic mapping."""
    x = [0.0, 0.25, 0.5, 0.75, 1.0]
    y = [0.05, 0.2, 0.5, 0.8, 0.95]
    return CalibParams(x_breaks=x, y_breaks=y, n_samples=100, epsilon=None, site_id="source")


def make_chains(n=5):
    return [{"question": f"q{i}", "steps": [], "risk_score": 0.3 + i * 0.1} for i in range(n)]


# ── Fit contract ──────────────────────────────────────────────────────────────

def test_fit_requires_2_chains():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    chains = make_chains(1)
    with pytest.raises(ValueError, match="at least 2"):
        cal.fit(chains, [1])


def test_fit_succeeds_with_2_chains():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    raw_scores = [0.2, 0.8]
    chains = [{"raw_score": s} for s in raw_scores]
    labels = [1, 0]
    cal.fit(chains, labels, raw_scores=raw_scores)
    assert cal._T is not None
    assert cal._b is not None


def test_fit_with_5_chains():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    raw_scores = [0.1, 0.2, 0.7, 0.8, 0.9]
    chains = [{"raw_score": s} for s in raw_scores]
    labels = [1, 1, 0, 0, 0]
    cal.fit(chains, labels, raw_scores=raw_scores)
    assert cal._T is not None


# ── predict_one ───────────────────────────────────────────────────────────────

def test_predict_one_in_range():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    raw_scores = [0.2, 0.8]
    cal.fit([{}] * 2, [1, 0], raw_scores=raw_scores)
    for x in [0.0, 0.1, 0.5, 0.9, 1.0]:
        out = cal.predict_one(x)
        assert 0.0 <= out <= 1.0, f"predict_one({x})={out} out of [0,1]"


def test_predict_one_monotone():
    """Higher raw risk → higher calibrated risk."""
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    raw_scores = [0.1, 0.9]
    cal.fit([{}] * 2, [1, 0], raw_scores=raw_scores)
    assert cal.predict_one(0.8) > cal.predict_one(0.2)


def test_predict_one_before_fit_raises():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    with pytest.raises(RuntimeError, match="fit"):
        cal.predict_one(0.5)


# ── export_params ─────────────────────────────────────────────────────────────

def test_export_params_returns_calib_params():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * 2, [1, 0], raw_scores=[0.2, 0.8])
    params = cal.export_params()
    assert isinstance(params, CalibParams)


def test_export_params_x_monotone():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * 2, [1, 0], raw_scores=[0.2, 0.8])
    params = cal.export_params()
    assert all(params.x_breaks[i] < params.x_breaks[i+1] for i in range(len(params.x_breaks)-1))


def test_export_params_n_samples_equals_new_chains():
    n = 3
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * n, [1, 0, 1], raw_scores=[0.2, 0.8, 0.3])
    params = cal.export_params()
    assert params.n_samples == n


def test_export_params_site_id_transfer():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * 2, [1, 0], raw_scores=[0.2, 0.8])
    params = cal.export_params()
    assert params.site_id == "transfer"


def test_export_params_epsilon_none():
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * 2, [1, 0], raw_scores=[0.2, 0.8])
    params = cal.export_params()
    assert params.epsilon is None


def test_export_params_compatible_with_federated_merge():
    """CalibParams from export_params() can be passed to FederatedCalibrator.merge()."""
    from llm_guard.calibration.federated_calibrator import FederatedCalibrator
    cal = ModelTransferCalibrator(make_source_params(), "claude-sonnet-3-5")
    cal.fit([{}] * 2, [1, 0], raw_scores=[0.2, 0.8])
    params = cal.export_params()
    # Should not raise
    merged = FederatedCalibrator.merge([params, make_source_params()])
    assert merged is not None
