from __future__ import annotations

import math
from pathlib import Path
import unittest

from analyze_transport_calibration import (
    apply_calibrator,
    fit_logit_calibrator,
    nested_crossfit,
    read_frozen_primary_contrast,
)


HERE = Path(__file__).resolve().parent


class TransportCalibrationTests(unittest.TestCase):
    def test_calibrator_returns_finite_probabilities(self) -> None:
        probabilities = [0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 0.55, 0.45]
        labels = [0, 0, 0, 1, 0, 1, 1, 1, 1, 0]
        intercept, slope = fit_logit_calibrator(probabilities, labels)
        calibrated = [apply_calibrator(value, intercept, slope) for value in probabilities]
        self.assertTrue(all(math.isfinite(value) and 0 < value < 1 for value in calibrated))

    def test_nested_crossfit_excludes_both_held_axes(self) -> None:
        rows = []
        for method in ("trace", "semantic"):
            for model_index, model in enumerate(("a", "b", "c", "d")):
                for fold in range(2):
                    for repeat in range(6):
                        rows.append(
                            {
                                "row_id": f"{model}:{fold}:{repeat}",
                                "group_id": f"g:{fold}:{repeat}",
                                "model_id": model,
                                "domain": "d",
                                "task_fold": fold,
                                "success": (model_index + fold + repeat) % 2,
                                "method": method,
                                "probability": 0.3 + 0.1 * ((model_index + repeat) % 5),
                            }
                        )
        predictions, parameters = nested_crossfit(rows)
        self.assertEqual(len(predictions), len(rows))
        self.assertTrue(all(int(row["training_model_overlap"]) == 0 for row in parameters))
        self.assertTrue(all(int(row["training_task_fold_overlap"]) == 0 for row in parameters))
        self.assertEqual({(row["row_id"], row["method"]) for row in predictions}, {(row["row_id"], row["method"]) for row in rows})

    def test_frozen_interval_comes_from_primary_contrast(self) -> None:
        difference, lower, upper = read_frozen_primary_contrast(
            HERE / "results" / "paired_contrasts.csv"
        )
        self.assertAlmostEqual(difference, 0.005708725458335169)
        self.assertAlmostEqual(lower, -0.0002401924037327331)
        self.assertAlmostEqual(upper, 0.011695577345967162)


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