"""
Tests for FederatedCalibrator — v0.30.0

Validates:
  1. Local fit + export round-trip
  2. DP noise (epsilon) reduces |y_break| by ≤ expected magnitude
  3. Federated merge ≈ centralised calibration (within ±5pp)
  4. threshold_at_fpr gives conservative bound
  5. Privacy: no labels or raw scores in exported params
  6. Edge cases: single site, all-correct, min_samples enforcement
"""
import sys, numpy as np, pytest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.federated_calibrator import FederatedCalibrator, CalibParams


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

def _make_scores_labels(n=100, wrong_rate=0.4, seed=0):
    rng = np.random.RandomState(seed)
    labels = (rng.rand(n) < wrong_rate).astype(int)
    scores = np.where(labels == 1,
                      rng.beta(5, 2, n),   # wrong chains: high risk
                      rng.beta(2, 5, n))   # correct chains: low risk
    return scores.tolist(), labels.tolist()


# ── Local fit & export ────────────────────────────────────────────────────────

class TestLocalFitExport:

    def test_fit_succeeds(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        assert cal.is_fitted
        assert cal.n_samples == 100

    def test_fit_requires_min_samples(self):
        sc, lb = _make_scores_labels(5)
        cal = FederatedCalibrator(min_samples=20)
        with pytest.raises(ValueError, match="≥ 20"):
            cal.fit(sc, lb)

    def test_export_no_noise(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        params = cal.export_params(epsilon=None, site_id="site_a")
        assert params.epsilon is None
        assert params.site_id == "site_a"
        assert params.n_samples == 100
        assert len(params.x_breaks) == len(params.y_breaks)
        assert all(0.0 <= y <= 1.0 for y in params.y_breaks)
        # x_breaks must be monotone increasing
        assert all(params.x_breaks[i] <= params.x_breaks[i+1]
                   for i in range(len(params.x_breaks)-1))

    def test_export_with_dp_noise(self):
        sc, lb = _make_scores_labels(200)
        cal = FederatedCalibrator(seed=42)
        cal.fit(sc, lb)
        params_clean = cal.export_params(epsilon=None)
        params_noisy = cal.export_params(epsilon=1.0)
        # y_breaks should differ (DP noise applied)
        diff = [abs(a-b) for a,b in zip(params_clean.y_breaks, params_noisy.y_breaks)]
        assert max(diff) > 0
        # But all y_breaks still in [0, 1]
        assert all(0.0 <= y <= 1.0 for y in params_noisy.y_breaks)
        assert params_noisy.epsilon == pytest.approx(1.0)

    def test_dp_noise_magnitude(self):
        """DP noise for epsilon=1.0, n=200 should have mean |noise| ≈ Δf/ε = 1/(200*1)=0.005."""
        sc, lb = _make_scores_labels(200)
        cal = FederatedCalibrator(seed=7)
        cal.fit(sc, lb)
        p_clean = cal.export_params(epsilon=None)
        p_noisy = cal.export_params(epsilon=1.0)
        mean_noise = np.mean([abs(a-b) for a,b in zip(p_clean.y_breaks, p_noisy.y_breaks)])
        # Expected mean |Laplace noise| = Δf/ε = 0.005; allow 10x margin for small samples
        assert mean_noise < 0.10  # sanity upper bound

    def test_serialisation_round_trip(self):
        sc, lb = _make_scores_labels(80)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        params = cal.export_params()
        # to/from dict
        d = params.to_dict()
        restored = CalibParams.from_dict(d)
        assert restored.x_breaks == pytest.approx(params.x_breaks)
        assert restored.y_breaks == pytest.approx(params.y_breaks)
        assert restored.n_samples == params.n_samples
        # to/from JSON
        j = params.to_json()
        restored_j = CalibParams.from_json(j)
        assert restored_j.x_breaks == pytest.approx(params.x_breaks)


# ── Inference ─────────────────────────────────────────────────────────────────

class TestInference:

    def test_predict_range(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        scores = [0.0, 0.25, 0.5, 0.75, 1.0]
        out = cal.predict(scores)
        assert all(0.0 <= v <= 1.0 for v in out)

    def test_predict_one(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        r = cal.predict_one(0.6)
        assert isinstance(r, float)
        assert 0.0 <= r <= 1.0

    def test_predict_monotone(self):
        """Calibrated output should be monotone (isotonic regression property)."""
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        xs = np.linspace(0.0, 1.0, 50).tolist()
        ys = cal.predict(xs)
        # Should be non-decreasing
        assert all(ys[i] <= ys[i+1] + 1e-9 for i in range(len(ys)-1))

    def test_predict_before_fit_raises(self):
        cal = FederatedCalibrator()
        with pytest.raises(RuntimeError):
            cal.predict([0.5])

    def test_threshold_at_fpr(self):
        sc, lb = _make_scores_labels(200)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        t = cal.threshold_at_fpr(0.10)
        assert 0.0 <= t <= 1.0

    def test_threshold_stricter_for_lower_fpr(self):
        sc, lb = _make_scores_labels(200)
        cal = FederatedCalibrator()
        cal.fit(sc, lb)
        t10 = cal.threshold_at_fpr(0.10)
        t05 = cal.threshold_at_fpr(0.05)
        # Stricter FPR → higher threshold
        assert t05 >= t10 - 1e-9


# ── Federated merge ───────────────────────────────────────────────────────────

class TestFederatedMerge:

    def test_merge_two_sites(self):
        sc_a, lb_a = _make_scores_labels(100, seed=0)
        sc_b, lb_b = _make_scores_labels(100, seed=1)
        cal_a = FederatedCalibrator().fit(sc_a, lb_a)
        cal_b = FederatedCalibrator().fit(sc_b, lb_b)
        params_a = cal_a.export_params(epsilon=None)
        params_b = cal_b.export_params(epsilon=None)
        merged = FederatedCalibrator.merge([params_a, params_b])
        assert merged.is_fitted
        assert merged.n_samples == 200
        out = merged.predict([0.3, 0.5, 0.7])
        assert all(0.0 <= v <= 1.0 for v in out)

    def test_merge_vs_centralised(self):
        """
        Federated merge should be within ±5pp of centralised calibration
        on the same combined dataset.  This validates the main correctness claim.
        """
        np.random.seed(0)
        sc_a, lb_a = _make_scores_labels(150, seed=10)
        sc_b, lb_b = _make_scores_labels(150, seed=11)

        # Centralised: fit on all data together
        central = FederatedCalibrator(min_samples=20)
        central.fit(sc_a + sc_b, lb_a + lb_b)

        # Federated: fit separately, merge
        cal_a = FederatedCalibrator(min_samples=20).fit(sc_a, lb_a)
        cal_b = FederatedCalibrator(min_samples=20).fit(sc_b, lb_b)
        merged = FederatedCalibrator.merge(
            [cal_a.export_params(), cal_b.export_params()]
        )

        # Compare on a test grid
        x_test = np.linspace(0.1, 0.9, 50).tolist()
        y_central = central.predict(x_test)
        y_merged  = merged.predict(x_test)
        mae = float(np.mean(np.abs(y_central - y_merged)))
        assert mae < 0.05, f"Federated vs centralised MAE = {mae:.4f} > 0.05"

    def test_merge_with_dp_noise(self):
        """DP-noised merge should still be within ±8pp of centralised (looser bound)."""
        sc_a, lb_a = _make_scores_labels(200, seed=20)
        sc_b, lb_b = _make_scores_labels(200, seed=21)
        sc_c, lb_c = _make_scores_labels(200, seed=22)

        central = FederatedCalibrator().fit(sc_a + sc_b + sc_c, lb_a + lb_b + lb_c)
        params_list = []
        for sc, lb, s in [(sc_a, lb_a, 0), (sc_b, lb_b, 1), (sc_c, lb_c, 2)]:
            cal = FederatedCalibrator(seed=s).fit(sc, lb)
            params_list.append(cal.export_params(epsilon=1.0))
        merged = FederatedCalibrator.merge(params_list)

        x_test = np.linspace(0.1, 0.9, 50).tolist()
        mae = float(np.mean(np.abs(central.predict(x_test) - merged.predict(x_test))))
        assert mae < 0.08, f"DP-federated MAE = {mae:.4f} > 0.08"

    def test_merge_weighted(self):
        """Larger-site weight should pull merged calibration toward that site."""
        sc_a, lb_a = _make_scores_labels(100, wrong_rate=0.2, seed=30)
        sc_b, lb_b = _make_scores_labels(100, wrong_rate=0.8, seed=31)
        cal_a = FederatedCalibrator().fit(sc_a, lb_a)
        cal_b = FederatedCalibrator().fit(sc_b, lb_b)
        p_a, p_b = cal_a.export_params(), cal_b.export_params()

        merged_a_heavy = FederatedCalibrator.merge([p_a, p_b], weights=[10.0, 1.0])
        merged_b_heavy = FederatedCalibrator.merge([p_a, p_b], weights=[1.0, 10.0])
        # A-heavy should have lower calibrated output (lower wrong rate)
        x_mid = [0.5]
        assert merged_a_heavy.predict(x_mid)[0] < merged_b_heavy.predict(x_mid)[0]

    def test_merge_single_site(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator().fit(sc, lb)
        merged = FederatedCalibrator.merge([cal.export_params()])
        # Single-site merge should predict the same as original
        x_test = np.linspace(0.1, 0.9, 20).tolist()
        # Grid interpolation introduces rounding at breakpoints — allow 10pp tolerance
        mae = float(np.mean(np.abs(merged.predict(x_test) - cal.predict(x_test))))
        assert mae < 0.10, f"Single-site merge MAE = {mae:.4f} > 0.10"

    def test_merge_empty_raises(self):
        with pytest.raises(ValueError):
            FederatedCalibrator.merge([])

    def test_merge_output_monotone(self):
        sc_a, lb_a = _make_scores_labels(100, seed=40)
        sc_b, lb_b = _make_scores_labels(100, seed=41)
        params = [FederatedCalibrator().fit(sc_a, lb_a).export_params(),
                  FederatedCalibrator().fit(sc_b, lb_b).export_params()]
        merged = FederatedCalibrator.merge(params)
        xs = np.linspace(0.0, 1.0, 100).tolist()
        ys = merged.predict(xs)
        assert all(ys[i] <= ys[i+1] + 1e-9 for i in range(len(ys)-1))


# ── Privacy: no raw data in export ───────────────────────────────────────────

class TestPrivacy:

    def test_params_contain_no_raw_scores(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator().fit(sc, lb)
        params = cal.export_params(epsilon=1.0)
        d = params.to_dict()
        # Only allowed keys
        assert set(d.keys()) == {"x_breaks", "y_breaks", "n_samples", "epsilon", "site_id"}
        # Breakpoints are aggregated — there are far fewer than n_samples
        assert len(params.x_breaks) < len(sc)

    def test_no_labels_in_params(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator().fit(sc, lb)
        params = cal.export_params()
        d = params.to_dict()
        # No labels key
        assert "labels" not in d
        assert "label" not in d

    def test_exported_from_package(self):
        from llm_guard import FederatedCalibrator as FC, CalibParams as CP
        assert FC is not None
        assert CP is not None


# ── Edge cases ────────────────────────────────────────────────────────────────

class TestEdgeCases:

    def test_all_correct_warns(self):
        sc = [0.1, 0.2, 0.15, 0.12, 0.18, 0.22, 0.09, 0.11, 0.17, 0.21,
              0.10, 0.19, 0.14, 0.16, 0.13, 0.20, 0.08, 0.23, 0.07, 0.25]
        lb = [0] * 20  # all correct
        cal = FederatedCalibrator(min_samples=20)
        with pytest.warns(UserWarning, match="one class"):
            cal.fit(sc, lb)
        assert cal.is_fitted

    def test_score_outside_training_range(self):
        sc, lb = _make_scores_labels(100)
        cal = FederatedCalibrator().fit(sc, lb)
        # IsotonicRegression clips out-of-bounds; should not raise
        out = cal.predict([-0.5, 1.5])
        assert all(0.0 <= v <= 1.0 for v in out)

    def test_five_sites_merge(self):
        params_list = []
        for seed in range(5):
            sc, lb = _make_scores_labels(60, seed=seed)
            cal = FederatedCalibrator(min_samples=20).fit(sc, lb)
            params_list.append(cal.export_params(epsilon=0.5, site_id=f"site_{seed}"))
        merged = FederatedCalibrator.merge(params_list)
        assert merged.n_samples == 300
        out = merged.predict([0.3, 0.6])
        assert all(0.0 <= v <= 1.0 for v in out)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
