"""tests/test_calibration_energy_aware.py — Unit tests for llm_guard.calibration.energy_aware_calibrator."""
import pytest
import numpy as np

pytestmark = pytest.mark.unit


class TestImport:
    def test_import(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        assert EnergyAwareCalibrator is not None

    def test_import_result(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyCalibrationResult
        assert EnergyCalibrationResult is not None


class TestInstantiation:
    def test_default_construction(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        assert cal.high_threshold == 0.80
        assert cal.low_threshold == 0.50

    def test_invalid_thresholds_raise(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        with pytest.raises(ValueError):
            EnergyAwareCalibrator(low_threshold=0.8, high_threshold=0.5)
        with pytest.raises(ValueError):
            EnergyAwareCalibrator(low_threshold=0.0, high_threshold=0.8)

    def test_not_fitted_initially(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        assert not cal.is_fitted


class TestThreshold:
    def test_high_energy_gives_high_threshold(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator(high_threshold=0.80, low_threshold=0.50)
        assert abs(cal.threshold(1.0) - 0.80) < 1e-9

    def test_low_energy_gives_low_threshold(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator(high_threshold=0.80, low_threshold=0.50)
        assert abs(cal.threshold(0.0) - 0.50) < 1e-9

    def test_medium_energy_interpolates(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator(high_threshold=0.80, low_threshold=0.50)
        t = cal.threshold(0.5)
        assert 0.50 < t < 0.80

    def test_energy_clipped(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        t1 = cal.threshold(1.5)  # clipped to 1.0
        t2 = cal.threshold(-0.5)  # clipped to 0.0
        assert t1 == cal.threshold(1.0)
        assert t2 == cal.threshold(0.0)


class TestCalibrate:
    def test_calibrate_returns_result(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        result = cal.calibrate(0.6, energy=0.7)
        assert hasattr(result, "alert")
        assert hasattr(result, "threshold")
        assert hasattr(result, "energy_regime")

    def test_energy_regime_high(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        result = cal.calibrate(0.5, energy=0.9)
        assert result.energy_regime == "high"

    def test_energy_regime_medium(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        result = cal.calibrate(0.5, energy=0.5)
        assert result.energy_regime == "medium"

    def test_energy_regime_low(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        result = cal.calibrate(0.5, energy=0.1)
        assert result.energy_regime == "low"

    def test_calibrate_batch(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        results = cal.calibrate_batch([0.4, 0.6, 0.8], [0.9, 0.5, 0.1])
        assert len(results) == 3

    def test_optimal_energy_for_recall(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        e = cal.optimal_energy_for_recall(1.0)
        assert e == 0.0  # recall=1 → energy=0


class TestFit:
    def test_fit_marks_fitted(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        rng = np.random.default_rng(42)
        scores = rng.uniform(0, 1, 20).tolist()
        labels = [int(s > 0.5) for s in scores]
        cal.fit(scores, labels)
        assert cal.is_fitted

    def test_calibrate_after_fit(self):
        from llm_guard.calibration.energy_aware_calibrator import EnergyAwareCalibrator
        cal = EnergyAwareCalibrator()
        rng = np.random.default_rng(42)
        scores = rng.uniform(0, 1, 20).tolist()
        labels = [int(s > 0.5) for s in scores]
        cal.fit(scores, labels)
        result = cal.calibrate(0.7, energy=0.5)
        assert 0.0 <= result.calibrated_risk <= 1.0
