import math
import unittest

import numpy as np
import pandas as pd

from evaluate_calibration import (
    aggregate_metric_rows,
    compute_ipcw_cell,
    fit_weighted_logistic_calibration,
)


class IPCWCalibrationMetricTests(unittest.TestCase):
    def test_no_censoring_matches_binary_metrics(self):
        result = compute_ipcw_cell(
            probabilities=np.asarray([0.2, 0.8]),
            event_times=np.asarray([np.inf, 0.5]),
            censor_times=np.asarray([2.0, 2.0]),
            horizon=1.0,
            min_cases=1,
            min_controls=1,
            max_ipcw_weight=0.0,
        )

        self.assertIsNotNone(result)
        row, arrays = result
        self.assertEqual(row["n_events"], 1)
        self.assertEqual(row["n_controls"], 1)
        self.assertAlmostEqual(row["brier_ipcw"], 0.04)
        self.assertAlmostEqual(row["nll_ipcw"], -math.log(0.8))
        self.assertAlmostEqual(row["predicted_mean"], 0.5)
        self.assertAlmostEqual(row["observed_rate_ipcw"], 0.5)
        np.testing.assert_allclose(arrays["metric_weights"], [1.0, 1.0])

    def test_censored_before_horizon_gets_zero_outcome_weight(self):
        result = compute_ipcw_cell(
            probabilities=np.asarray([0.8, 0.2, 0.4]),
            event_times=np.asarray([0.5, np.inf, np.inf]),
            censor_times=np.asarray([2.0, 2.0, 0.5]),
            horizon=1.0,
            min_cases=1,
            min_controls=1,
            max_ipcw_weight=0.0,
        )

        self.assertIsNotNone(result)
        row, arrays = result
        self.assertEqual(row["n_censored_before_horizon"], 1)
        self.assertAlmostEqual(row["known_fraction"], 2.0 / 3.0)
        np.testing.assert_allclose(
            arrays["metric_weights"],
            [1.0, 1.5, 0.0],
        )
        self.assertAlmostEqual(row["brier_ipcw"], 0.1 / 3.0)
        self.assertAlmostEqual(
            row["nll_ipcw"],
            -2.5 * math.log(0.8) / 3.0,
        )
        self.assertAlmostEqual(row["observed_rate_ipcw"], 1.0 / 3.0)

    def test_calibration_intercept_and_slope_recover_identity(self):
        probabilities = np.repeat([0.1, 0.3, 0.7, 0.9], 100)
        outcomes = np.concatenate(
            [
                np.r_[np.ones(10), np.zeros(90)],
                np.r_[np.ones(30), np.zeros(70)],
                np.r_[np.ones(70), np.zeros(30)],
                np.r_[np.ones(90), np.zeros(10)],
            ]
        )
        weights = np.ones_like(probabilities)

        calibration_in_large, intercept, slope = (
            fit_weighted_logistic_calibration(
                probabilities,
                outcomes,
                weights,
            )
        )

        self.assertAlmostEqual(calibration_in_large, 0.0, places=7)
        self.assertAlmostEqual(intercept, 0.0, places=7)
        self.assertAlmostEqual(slope, 1.0, places=7)

    def test_metric_aggregation_uses_contribution_sums(self):
        metrics = pd.DataFrame(
            [
                {
                    "outcome": "Disease",
                    "sex": "Female",
                    "horizon": 5.0,
                    "n_at_risk": 10,
                    "n_events": 2,
                    "n_controls": 7,
                    "n_censored_before_horizon": 1,
                    "prediction_sum": 2.0,
                    "event_weight_sum": 2.0,
                    "brier_ipcw_sum": 1.0,
                    "nll_ipcw_sum": 3.0,
                    "calibration_in_the_large": 0.1,
                    "calibration_intercept": 0.2,
                    "calibration_slope": 0.9,
                    "ipcw_weight_max": 1.2,
                    "ipcw_weights_clipped": 0,
                },
                {
                    "outcome": "Disease",
                    "sex": "Female",
                    "horizon": 5.0,
                    "n_at_risk": 10,
                    "n_events": 3,
                    "n_controls": 6,
                    "n_censored_before_horizon": 1,
                    "prediction_sum": 3.0,
                    "event_weight_sum": 3.0,
                    "brier_ipcw_sum": 2.0,
                    "nll_ipcw_sum": 4.0,
                    "calibration_in_the_large": -0.1,
                    "calibration_intercept": -0.2,
                    "calibration_slope": 1.1,
                    "ipcw_weight_max": 1.4,
                    "ipcw_weights_clipped": 1,
                },
            ]
        )

        aggregated = aggregate_metric_rows(
            metrics,
            group_columns=["outcome", "sex", "horizon"],
        ).iloc[0]

        self.assertEqual(aggregated["n_at_risk"], 20)
        self.assertAlmostEqual(aggregated["predicted_mean"], 0.25)
        self.assertAlmostEqual(aggregated["observed_rate_ipcw"], 0.25)
        self.assertAlmostEqual(aggregated["brier_ipcw"], 0.15)
        self.assertAlmostEqual(aggregated["nll_ipcw"], 0.35)
        self.assertAlmostEqual(aggregated["calibration_slope_median"], 1.0)
        self.assertEqual(aggregated["ipcw_weights_clipped"], 1)


if __name__ == "__main__":
    unittest.main()