"""Integration tests: full pipeline from ClaimsHistory to posterior premium."""

import numpy as np
import pytest
import polars as pl
from insurance_experience import (
    ClaimsHistory,
    StaticCredibilityModel,
    DynamicPoissonGammaModel,
    SurrogateModel,
    balance_calibrate,
    calibrated_predict_fn,
    balance_report,
)


def make_realistic_portfolio(
    n: int,
    rng: np.random.Generator,
) -> list[ClaimsHistory]:
    """Generate a portfolio with heterogeneous risks and varying history lengths."""
    histories = []
    # Two risk tiers: clean (lambda=0.4) and standard (lambda=1.2)
    for i in range(n // 2):
        n_periods = rng.integers(2, 6)
        counts = rng.poisson(0.4, size=n_periods).tolist()
        histories.append(
            ClaimsHistory(
                f"CLEAN_{i}",
                list(range(1, n_periods + 1)),
                counts,
                exposures=[rng.uniform(0.5, 1.5) for _ in range(n_periods)],
                prior_premium=500.0,
            )
        )
    for i in range(n // 2):
        n_periods = rng.integers(2, 6)
        counts = rng.poisson(1.2, size=n_periods).tolist()
        histories.append(
            ClaimsHistory(
                f"STD_{i}",
                list(range(1, n_periods + 1)),
                counts,
                exposures=[rng.uniform(0.5, 1.5) for _ in range(n_periods)],
                prior_premium=600.0,
            )
        )
    return histories


class TestFullPipelineStatic:
    def test_full_pipeline(self):
        rng = np.random.default_rng(500)
        histories = make_realistic_portfolio(60, rng)

        model = StaticCredibilityModel()
        model.fit(histories)

        cal = balance_calibrate(model.predict, histories)
        calibrated = calibrated_predict_fn(model.predict, cal)

        for h in histories:
            cf = calibrated(h)
            posterior = h.prior_premium * cf
            assert posterior > 0, f"Non-positive posterior: {posterior}"
            assert np.isfinite(posterior), f"Non-finite posterior: {posterior}"

    def test_clean_risks_lower_posterior(self):
        """Clean risks should have lower posterior than standard risks."""
        rng = np.random.default_rng(501)
        histories = make_realistic_portfolio(60, rng)

        model = StaticCredibilityModel()
        model.fit(histories)

        clean_cfs = [
            model.predict(h) for h in histories if h.policy_id.startswith("CLEAN")
        ]
        std_cfs = [
            model.predict(h) for h in histories if h.policy_id.startswith("STD")
        ]

        assert np.mean(clean_cfs) < np.mean(std_cfs), (
            f"Mean CF clean={np.mean(clean_cfs):.3f} not < std={np.mean(std_cfs):.3f}"
        )


class TestFullPipelineDynamic:
    def test_full_pipeline(self):
        rng = np.random.default_rng(502)
        histories = make_realistic_portfolio(60, rng)

        model = DynamicPoissonGammaModel()
        model.fit(histories)

        cal = balance_calibrate(model.predict, histories)
        calibrated = calibrated_predict_fn(model.predict, cal)

        for h in histories:
            cf = calibrated(h)
            posterior = h.prior_premium * cf
            assert posterior > 0
            assert np.isfinite(posterior)

    def test_dynamic_and_static_broadly_agree(self):
        """Static and dynamic models should broadly rank policies similarly."""
        rng = np.random.default_rng(503)
        histories = make_realistic_portfolio(80, rng)

        static = StaticCredibilityModel()
        static.fit(histories)

        dynamic = DynamicPoissonGammaModel()
        dynamic.fit(histories)

        static_cfs = [static.predict(h) for h in histories]
        dynamic_cfs = [dynamic.predict(h) for h in histories]

        # Spearman rank correlation should be positive
        from scipy.stats import spearmanr
        corr, _ = spearmanr(static_cfs, dynamic_cfs)
        assert corr > 0.0, f"Expected positive rank correlation, got {corr}"


class TestFullPipelineSurrogate:
    def test_full_pipeline(self):
        rng = np.random.default_rng(504)
        histories = make_realistic_portfolio(60, rng)

        model = SurrogateModel(n_is_samples=300, random_state=42)
        model.fit(histories)

        for h in histories[:10]:
            cf = model.predict(h)
            posterior = h.prior_premium * cf
            assert posterior > 0
            assert np.isfinite(posterior)


class TestBalancePropertyEndToEnd:
    def test_balance_enforced_after_calibration_static(self):
        rng = np.random.default_rng(505)
        histories = make_realistic_portfolio(80, rng)
        model = StaticCredibilityModel()
        model.fit(histories)
        cal = balance_calibrate(model.predict, histories)
        calibrated = calibrated_predict_fn(model.predict, cal)

        sum_actual = sum(h.claim_frequency * h.total_exposure for h in histories)
        sum_posterior = sum(
            h.prior_premium * calibrated(h) * h.total_exposure for h in histories
        )
        assert sum_posterior == pytest.approx(sum_actual, rel=1e-4)

    def test_balance_enforced_after_calibration_dynamic(self):
        rng = np.random.default_rng(506)
        histories = make_realistic_portfolio(80, rng)
        model = DynamicPoissonGammaModel()
        model.fit(histories)
        cal = balance_calibrate(model.predict, histories)
        calibrated = calibrated_predict_fn(model.predict, cal)

        sum_actual = sum(h.claim_frequency * h.total_exposure for h in histories)
        sum_posterior = sum(
            h.prior_premium * calibrated(h) * h.total_exposure for h in histories
        )
        assert sum_posterior == pytest.approx(sum_actual, rel=1e-4)


class TestEdgeCasesIntegration:
    def test_single_period_histories_work(self):
        histories = [
            ClaimsHistory(f"P{i}", [1], [i % 3], prior_premium=1.0)
            for i in range(20)
        ]
        model = StaticCredibilityModel()
        model.fit(histories)
        for h in histories:
            cf = model.predict(h)
            assert np.isfinite(cf)
            assert cf >= 0.0

    def test_all_zero_claims_works(self):
        histories = [
            ClaimsHistory(f"P{i}", [1, 2, 3], [0, 0, 0], prior_premium=1.5)
            for i in range(20)
        ]
        model = StaticCredibilityModel()
        model.fit(histories)
        cf = model.predict(histories[0])
        assert cf == pytest.approx(0.0, abs=1e-6)

    def test_very_high_claim_counts(self):
        """Model should not crash or produce NaN for extreme claim counts."""
        rng = np.random.default_rng(507)
        histories = [
            ClaimsHistory(f"P{i}", [1, 2, 3], [1, 1, 1], prior_premium=1.0)
            for i in range(20)
        ]
        extreme = ClaimsHistory("EXTREME", [1, 2, 3], [100, 0, 100], prior_premium=1.0)

        model = StaticCredibilityModel()
        model.fit(histories)
        cf = model.predict(extreme)
        assert np.isfinite(cf)

    def test_predict_batch_polars_schema(self):
        """Verify DataFrame column types are sensible."""
        rng = np.random.default_rng(508)
        histories = [
            ClaimsHistory(f"P{i}", [1, 2, 3], rng.poisson(1, 3).tolist(), prior_premium=1.0)
            for i in range(20)
        ]
        model = StaticCredibilityModel()
        model.fit(histories)
        df = model.predict_batch(histories)

        assert df["policy_id"].dtype == pl.Utf8 or df["policy_id"].dtype == pl.String
        assert df["prior_premium"].dtype in (pl.Float64, pl.Float32)
        assert df["credibility_factor"].dtype in (pl.Float64, pl.Float32)
        assert df["posterior_premium"].dtype in (pl.Float64, pl.Float32)

    def test_balance_report_full_run(self):
        rng = np.random.default_rng(509)
        histories = [
            ClaimsHistory(f"P{i}", [1, 2, 3, 4], rng.poisson(0.8, 4).tolist(), prior_premium=1.0)
            for i in range(30)
        ]
        model = StaticCredibilityModel()
        model.fit(histories)
        df = balance_report(model.predict, histories, by_n_periods=False)
        assert len(df) == 30

        df_grouped = balance_report(model.predict, histories, by_n_periods=True)
        # All policies have 4 periods, so 1 group
        assert len(df_grouped) == 1
