"""Tests for ForgettingTracker and forgetting metrics."""
import json
import math
import tempfile
from pathlib import Path

import pytest

from llm_refresh.metrics import ForgettingTracker
from llm_refresh.models import AnchorPoint, EvalResult, RefreshResult


def make_eval(ppl: float, loss: float = 1.0, n: int = 100, label: str = "") -> EvalResult:
    return EvalResult(
        timestamp="2026-01-01T00:00:00+00:00",
        perplexity=ppl,
        loss=loss,
        dataset_size=n,
        label=label,
    )


def make_result(
    bwt: float = 0.0,
    fwt: float = 0.0,
    krs: float = 100.0,
    rolled_back: bool = False,
    pre_ppl: float = 5.0,
    post_ppl: float = 5.0,
) -> RefreshResult:
    return RefreshResult.new(
        examples_trained=100,
        replay_examples=30,
        new_examples=70,
        pre_eval=make_eval(pre_ppl),
        post_eval=make_eval(post_ppl),
        bwt=bwt,
        fwt=fwt,
        krs=krs,
        rolled_back=rolled_back,
        rollback_reason="",
    )


class TestForgettingTracker:
    def test_initial_state(self):
        t = ForgettingTracker()
        assert t.compute_bwt() == 0.0
        assert t.compute_fwt() == 0.0
        assert t.compute_krs() == 100.0
        assert not t.is_forgetting()

    def test_record_and_compute_bwt(self):
        t = ForgettingTracker()
        # Simulate 3 cycles with varying BWT
        t.record_result(make_result(bwt=-0.5))  # forgetting
        t.record_result(make_result(bwt=0.2))   # improvement
        t.record_result(make_result(bwt=-0.3))  # forgetting

        avg_bwt = t.compute_bwt()
        expected = (-0.5 + 0.2 + -0.3) / 3
        assert abs(avg_bwt - expected) < 1e-9

    def test_krs_no_forgetting(self):
        t = ForgettingTracker(krs_lambda=0.01)
        t.record_result(make_result(bwt=0.1))   # positive BWT = no forgetting
        t.record_result(make_result(bwt=0.05))
        # cumulative_forgetting = 0
        assert t.compute_krs() == pytest.approx(100.0, abs=0.01)

    def test_krs_with_forgetting(self):
        t = ForgettingTracker(krs_lambda=0.01)
        # Two forgetting events
        t.record_result(make_result(bwt=-1.0))
        t.record_result(make_result(bwt=-2.0))
        # cumulative = 3.0
        expected = 100.0 * math.exp(-0.01 * 3.0)
        assert t.compute_krs() == pytest.approx(expected, rel=1e-6)

    def test_krs_full_forgetting(self):
        t = ForgettingTracker(krs_lambda=1.0)
        # Extreme forgetting
        t.record_result(make_result(bwt=-100.0))
        assert t.compute_krs() < 1.0  # nearly zero

    def test_is_forgetting_true(self):
        t = ForgettingTracker()
        t.record_result(make_result(bwt=-0.5))  # well below default -0.1
        assert t.is_forgetting(threshold=-0.1)

    def test_is_forgetting_false(self):
        t = ForgettingTracker()
        t.record_result(make_result(bwt=0.2))  # positive BWT
        assert not t.is_forgetting(threshold=-0.1)

    def test_is_forgetting_borderline(self):
        t = ForgettingTracker()
        t.record_result(make_result(bwt=-0.05))  # above threshold -0.1
        assert not t.is_forgetting(threshold=-0.1)

    def test_fwt_requires_two_cycles(self):
        t = ForgettingTracker()
        t.record_result(make_result(fwt=0.3))
        # Only one cycle — FWT undefined
        assert t.compute_fwt() == 0.0

    def test_fwt_multiple_cycles(self):
        t = ForgettingTracker()
        t.record_result(make_result(fwt=0.0))  # first cycle, no FWT
        t.record_result(make_result(fwt=0.4))
        t.record_result(make_result(fwt=0.2))
        avg_fwt = t.compute_fwt()
        # Average of cycles 2+ (fwt=0.4, fwt=0.2); excludes zeros
        assert abs(avg_fwt - 0.3) < 1e-9

    def test_summary(self):
        t = ForgettingTracker()
        t.record_result(make_result(bwt=-0.3, rolled_back=True))
        t.record_result(make_result(bwt=0.1))

        s = t.summary()
        assert s["cycles"] == 2
        assert s["total_forgetting_events"] == 1
        assert s["total_rollbacks"] == 1
        assert s["last_bwt"] == pytest.approx(0.1, abs=1e-6)

    def test_history_property(self):
        t = ForgettingTracker()
        r1 = make_result(bwt=-0.2)
        r2 = make_result(bwt=0.1)
        t.record_result(r1)
        t.record_result(r2)
        history = t.history
        assert len(history) == 2
        # Returns a copy
        history.append(make_result())
        assert len(t.history) == 2

    def test_record_pre_post_eval(self):
        t = ForgettingTracker()
        t.record_pre_eval(make_eval(10.0))
        t.record_post_eval(make_eval(8.0))
        # Should have 2 anchor points
        assert len(t._anchor_history) == 2


class TestSerialization:
    def test_to_from_dict(self):
        t = ForgettingTracker(krs_lambda=0.05)
        t.record_pre_eval(make_eval(10.0))
        t.record_post_eval(make_eval(9.0))
        t.record_result(make_result(bwt=-0.3, fwt=0.1, krs=97.0))

        d = t.to_dict()
        t2 = ForgettingTracker.from_dict(d)

        assert t2.krs_lambda == 0.05
        assert len(t2.history) == 1
        assert t2.history[0].bwt == pytest.approx(-0.3)
        assert len(t2._anchor_history) == 2

    def test_save_load(self, tmp_path):
        t = ForgettingTracker()
        t.record_result(make_result(bwt=-0.5, krs=95.0))
        path = tmp_path / "metrics.json"
        t.save(path)

        t2 = ForgettingTracker.load(path)
        assert len(t2.history) == 1
        assert t2.history[0].bwt == pytest.approx(-0.5)

    def test_eval_result_serialization(self):
        er = make_eval(7.5, loss=2.1, n=200, label="anchor")
        d = er.to_dict()
        er2 = EvalResult.from_dict(d)
        assert er2.perplexity == 7.5
        assert er2.loss == 2.1
        assert er2.dataset_size == 200
        assert er2.label == "anchor"

    def test_refresh_result_serialization(self):
        r = make_result(bwt=-0.3, fwt=0.2, krs=97.0, rolled_back=True)
        d = r.to_dict()
        r2 = RefreshResult.from_dict(d)
        assert r2.bwt == pytest.approx(-0.3)
        assert r2.fwt == pytest.approx(0.2)
        assert r2.rolled_back is True
        assert r2.pre_eval is not None
        assert r2.post_eval is not None
