"""Tests for calibrators."""

import warnings

import pytest
import torch

from rankcal import (
    BaseCalibrator,
    IsotonicCalibrator,
    MonotonicNNCalibrator,
    PiecewiseLinearCalibrator,
    SplineCalibrator,
    TemperatureScaling,
)
from rankcal.utils import generate_calibrated_data, generate_miscalibrated_data


class TestTemperatureScaling:
    def test_init(self):
        cal = TemperatureScaling()
        assert cal.temperature.item() == 1.0
        assert not cal.fitted

    def test_init_custom_temperature(self):
        cal = TemperatureScaling(init_temperature=2.0)
        assert cal.temperature.item() == 2.0

    def test_fit_returns_self(self):
        cal = TemperatureScaling()
        scores, labels = generate_calibrated_data(100, seed=42)
        result = cal.fit(scores, labels)
        assert result is cal
        assert cal.fitted

    def test_forward_requires_fit(self):
        cal = TemperatureScaling()
        scores = torch.rand(10)
        with pytest.raises(RuntimeError, match="must be fitted"):
            cal(scores)

    def test_forward_preserves_shape(self):
        cal = TemperatureScaling()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels)

        test_scores = torch.rand(50)
        calibrated = cal(test_scores)
        assert calibrated.shape == test_scores.shape

    def test_calibration_improves_miscalibrated(self):
        """Temperature scaling should reduce miscalibration."""
        scores, labels = generate_miscalibrated_data(500, temperature=2.0, seed=42)

        cal = TemperatureScaling()
        cal.fit(scores, labels)

        # After fitting, temperature should be > 1 to counter the miscalibration
        assert cal.temperature.item() > 1.0

    def test_output_in_valid_range(self):
        cal = TemperatureScaling()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels)

        calibrated = cal(scores)
        assert (calibrated >= 0).all()
        assert (calibrated <= 1).all()


class TestIsotonicCalibrator:
    def test_init(self):
        cal = IsotonicCalibrator()
        assert not cal.fitted

    def test_fit_returns_self(self):
        cal = IsotonicCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        result = cal.fit(scores, labels)
        assert result is cal
        assert cal.fitted

    def test_forward_requires_fit(self):
        cal = IsotonicCalibrator()
        scores = torch.rand(10)
        with pytest.raises(RuntimeError, match="must be fitted"):
            cal(scores)

    def test_forward_preserves_shape(self):
        cal = IsotonicCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels)

        test_scores = torch.rand(50)
        calibrated = cal(test_scores)
        assert calibrated.shape == test_scores.shape

    def test_output_is_monotonic(self):
        """Isotonic calibrator should produce monotonically increasing outputs."""
        cal = IsotonicCalibrator()
        scores, labels = generate_miscalibrated_data(200, seed=42)
        cal.fit(scores, labels)

        # Test on sorted inputs
        test_scores = torch.linspace(0.01, 0.99, 100)
        calibrated = cal(test_scores)

        # Check monotonicity (allowing for floating point tolerance)
        diffs = calibrated[1:] - calibrated[:-1]
        assert (diffs >= -1e-6).all()

    def test_output_in_valid_range(self):
        cal = IsotonicCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels)

        calibrated = cal(scores)
        assert (calibrated >= 0).all()
        assert (calibrated <= 1).all()


class TestPiecewiseLinearCalibrator:
    def test_init(self):
        cal = PiecewiseLinearCalibrator()
        assert cal.n_knots == 10
        assert not cal.fitted

    def test_init_custom_knots(self):
        cal = PiecewiseLinearCalibrator(n_knots=5)
        assert cal.n_knots == 5

    def test_fit_returns_self(self):
        cal = PiecewiseLinearCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        result = cal.fit(scores, labels, max_iter=50)
        assert result is cal
        assert cal.fitted

    def test_forward_requires_fit(self):
        cal = PiecewiseLinearCalibrator()
        scores = torch.rand(10)
        with pytest.raises(RuntimeError, match="must be fitted"):
            cal(scores)

    def test_forward_preserves_shape(self):
        cal = PiecewiseLinearCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels, max_iter=50)

        test_scores = torch.rand(50)
        calibrated = cal(test_scores)
        assert calibrated.shape == test_scores.shape

    def test_output_is_monotonic(self):
        """Piecewise linear calibrator should produce monotonically increasing outputs."""
        cal = PiecewiseLinearCalibrator()
        scores, labels = generate_calibrated_data(200, seed=42)
        cal.fit(scores, labels, max_iter=100)

        test_scores = torch.linspace(0.01, 0.99, 100)
        calibrated = cal(test_scores)

        diffs = calibrated[1:] - calibrated[:-1]
        assert (diffs >= -1e-5).all()

    def test_output_in_valid_range(self):
        cal = PiecewiseLinearCalibrator()
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels, max_iter=50)

        calibrated = cal(scores)
        assert (calibrated >= 0).all()
        assert (calibrated <= 1).all()


class TestSplineCalibratorDeprecation:
    """Tests for the deprecated SplineCalibrator alias."""

    def test_deprecation_warning(self):
        """SplineCalibrator should emit a deprecation warning."""
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            cal = SplineCalibrator()
            assert len(w) == 1
            assert issubclass(w[0].category, DeprecationWarning)
            assert "PiecewiseLinearCalibrator" in str(w[0].message)

    def test_still_works(self):
        """SplineCalibrator should still work despite deprecation."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            cal = SplineCalibrator()
            scores, labels = generate_calibrated_data(100, seed=42)
            cal.fit(scores, labels, max_iter=50)
            assert cal.fitted


class TestMonotonicNNCalibrator:
    def test_init(self):
        cal = MonotonicNNCalibrator()
        assert cal.hidden_dims == (16, 16)
        assert not cal.fitted

    def test_init_custom_dims(self):
        cal = MonotonicNNCalibrator(hidden_dims=(8, 8, 8))
        assert cal.hidden_dims == (8, 8, 8)

    def test_fit_returns_self(self):
        cal = MonotonicNNCalibrator(hidden_dims=(8,))
        scores, labels = generate_calibrated_data(100, seed=42)
        result = cal.fit(scores, labels, max_iter=50)
        assert result is cal
        assert cal.fitted

    def test_forward_requires_fit(self):
        cal = MonotonicNNCalibrator()
        scores = torch.rand(10)
        with pytest.raises(RuntimeError, match="must be fitted"):
            cal(scores)

    def test_forward_preserves_shape(self):
        cal = MonotonicNNCalibrator(hidden_dims=(8,))
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels, max_iter=50)

        test_scores = torch.rand(50)
        calibrated = cal(test_scores)
        assert calibrated.shape == test_scores.shape

    def test_output_is_monotonic(self):
        """Monotonic NN should produce monotonically increasing outputs."""
        cal = MonotonicNNCalibrator(hidden_dims=(16,))
        scores, labels = generate_calibrated_data(200, seed=42)
        cal.fit(scores, labels, max_iter=200)

        test_scores = torch.linspace(0.01, 0.99, 100)
        calibrated = cal(test_scores)

        diffs = calibrated[1:] - calibrated[:-1]
        assert (diffs >= -1e-5).all()

    def test_output_in_valid_range(self):
        cal = MonotonicNNCalibrator(hidden_dims=(8,))
        scores, labels = generate_calibrated_data(100, seed=42)
        cal.fit(scores, labels, max_iter=50)

        calibrated = cal(scores)
        assert (calibrated >= 0).all()
        assert (calibrated <= 1).all()


class TestInputValidation:
    """Test input validation across calibrators."""

    @pytest.mark.parametrize(
        "calibrator_cls",
        [TemperatureScaling, IsotonicCalibrator, PiecewiseLinearCalibrator],
    )
    def test_shape_mismatch_raises(self, calibrator_cls):
        cal = calibrator_cls()
        scores = torch.rand(10)
        labels = torch.rand(5)
        with pytest.raises(ValueError, match="same shape"):
            cal.fit(scores, labels)

    @pytest.mark.parametrize(
        "calibrator_cls",
        [TemperatureScaling, IsotonicCalibrator, PiecewiseLinearCalibrator],
    )
    def test_accepts_numpy_arrays(self, calibrator_cls):
        import numpy as np

        cal = calibrator_cls()
        scores = np.random.rand(100)
        labels = np.random.randint(0, 2, 100).astype(float)

        if calibrator_cls == PiecewiseLinearCalibrator:
            cal.fit(scores, labels, max_iter=10)
        else:
            cal.fit(scores, labels)

        assert cal.fitted


class TestSerialization:
    """Tests for save/load functionality."""

    @pytest.fixture
    def sample_data(self):
        return generate_miscalibrated_data(500, seed=42)

    def test_save_load_temperature_scaling(self, sample_data, tmp_path):
        scores, labels = sample_data
        calibrator = TemperatureScaling()
        calibrator.fit(scores, labels)

        # Save
        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))

        # Load
        loaded = TemperatureScaling.load(str(path))

        # Verify
        assert loaded.fitted
        assert torch.allclose(loaded.temperature, calibrator.temperature)
        assert torch.allclose(loaded(scores), calibrator(scores))

    def test_save_load_isotonic(self, sample_data, tmp_path):
        scores, labels = sample_data
        calibrator = IsotonicCalibrator()
        calibrator.fit(scores, labels)

        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))
        loaded = IsotonicCalibrator.load(str(path))

        assert loaded.fitted
        assert torch.allclose(loaded(scores), calibrator(scores))

    def test_save_load_piecewise_linear(self, sample_data, tmp_path):
        scores, labels = sample_data
        calibrator = PiecewiseLinearCalibrator()
        calibrator.fit(scores, labels)

        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))
        loaded = PiecewiseLinearCalibrator.load(str(path))

        assert loaded.fitted
        assert torch.allclose(loaded(scores), calibrator(scores))

    def test_save_load_monotonic_nn(self, sample_data, tmp_path):
        scores, labels = sample_data
        calibrator = MonotonicNNCalibrator()
        calibrator.fit(scores, labels)

        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))
        loaded = MonotonicNNCalibrator.load(str(path))

        assert loaded.fitted
        assert torch.allclose(loaded(scores), calibrator(scores), atol=1e-5)

    def test_save_unfitted_raises(self, tmp_path):
        calibrator = TemperatureScaling()
        path = tmp_path / "calibrator.pt"

        with pytest.raises(RuntimeError, match="must be fitted"):
            calibrator.save(str(path))

    def test_load_wrong_class_raises(self, sample_data, tmp_path):
        scores, labels = sample_data
        calibrator = TemperatureScaling()
        calibrator.fit(scores, labels)

        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))

        with pytest.raises(ValueError, match="Saved calibrator is"):
            IsotonicCalibrator.load(str(path))

    def test_base_calibrator_load(self, sample_data, tmp_path):
        """Test loading via BaseCalibrator.load() auto-detects class."""
        scores, labels = sample_data
        calibrator = TemperatureScaling()
        calibrator.fit(scores, labels)

        path = tmp_path / "calibrator.pt"
        calibrator.save(str(path))

        loaded = BaseCalibrator.load(str(path))
        assert isinstance(loaded, TemperatureScaling)
        assert torch.allclose(loaded(scores), calibrator(scores))
