"""Forgetting metrics: BWT, FWT, KRS."""
from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Optional

from llm_refresh.models import AnchorPoint, EvalResult, RefreshResult


class ForgettingTracker:
    """
    Tracks Backward Transfer (BWT), Forward Transfer (FWT), and
    Knowledge Retention Score (KRS) across continual learning cycles.

    Math:
        BWT_t  = PPL_anchor_before_t - PPL_anchor_after_t
                 (negative = perplexity went up = forgetting)

        FWT_t  = PPL_pre_(t-1) - PPL_pre_t
                 (positive = prior training helped on new data)

        KRS    = 100 * exp(-lambda * cumulative_forgetting)
                 where cumulative_forgetting = sum(|BWT_t| for BWT_t < 0)
    """

    def __init__(self, krs_lambda: float = 0.01):
        self.krs_lambda = krs_lambda
        self._anchor_history: list[AnchorPoint] = []
        self._refresh_history: list[RefreshResult] = []
        self._pending_pre: Optional[EvalResult] = None
        self._cycle_step: int = 0

    # ------------------------------------------------------------------
    # Recording
    # ------------------------------------------------------------------

    def record_pre_eval(self, result: EvalResult) -> None:
        """Call before each training cycle."""
        self._pending_pre = result

    def record_post_eval(self, result: EvalResult) -> None:
        """Call after each training cycle. Pairs with the last pre_eval."""
        self._cycle_step += 1
        if self._pending_pre is not None:
            self._anchor_history.append(
                AnchorPoint(step=self._cycle_step, eval_result=self._pending_pre)
            )
        self._anchor_history.append(
            AnchorPoint(step=self._cycle_step, eval_result=result)
        )

    def record_result(self, result: RefreshResult) -> None:
        self._refresh_history.append(result)

    # ------------------------------------------------------------------
    # Metrics
    # ------------------------------------------------------------------

    def compute_bwt(self) -> float:
        """
        Average backward transfer over all cycles.
        BWT_t = PPL_before - PPL_after
        Negative = perplexity increased = forgetting occurred.
        """
        if len(self._refresh_history) == 0:
            return 0.0
        bwts = [r.bwt for r in self._refresh_history if r.bwt != 0.0]
        if not bwts:
            return 0.0
        return sum(bwts) / len(bwts)

    def compute_fwt(self) -> float:
        """
        Average forward transfer: did earlier training help on subsequent data?
        Positive = subsequent training started at lower perplexity.
        """
        if len(self._refresh_history) < 2:
            return 0.0
        fwts = [r.fwt for r in self._refresh_history[1:] if r.fwt != 0.0]
        if not fwts:
            return 0.0
        return sum(fwts) / len(fwts)

    def compute_krs(self) -> float:
        """
        Knowledge Retention Score [0–100]. 100 = no forgetting ever.
        KRS = 100 * exp(-lambda * cumulative_forgetting)
        """
        cumulative = sum(
            abs(r.bwt) for r in self._refresh_history if r.bwt < 0
        )
        return 100.0 * math.exp(-self.krs_lambda * cumulative)

    def is_forgetting(self, threshold: float = -0.1) -> bool:
        """
        Returns True if the most recent BWT is below `threshold`.
        threshold=-0.1 means perplexity rose by more than 0.1.
        """
        if not self._refresh_history:
            return False
        return self._refresh_history[-1].bwt < threshold

    def summary(self) -> dict[str, Any]:
        return {
            "cycles": len(self._refresh_history),
            "bwt": round(self.compute_bwt(), 4),
            "fwt": round(self.compute_fwt(), 4),
            "krs": round(self.compute_krs(), 2),
            "total_forgetting_events": sum(
                1 for r in self._refresh_history if r.bwt < 0
            ),
            "total_rollbacks": sum(1 for r in self._refresh_history if r.rolled_back),
            "last_bwt": round(self._refresh_history[-1].bwt, 4) if self._refresh_history else None,
        }

    @property
    def history(self) -> list[RefreshResult]:
        return list(self._refresh_history)

    # ------------------------------------------------------------------
    # Serialization
    # ------------------------------------------------------------------

    def to_dict(self) -> dict[str, Any]:
        return {
            "krs_lambda": self.krs_lambda,
            "cycle_step": self._cycle_step,
            "anchor_history": [a.to_dict() for a in self._anchor_history],
            "refresh_history": [r.to_dict() for r in self._refresh_history],
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "ForgettingTracker":
        tracker = cls(krs_lambda=d.get("krs_lambda", 0.01))
        tracker._cycle_step = d.get("cycle_step", 0)
        tracker._anchor_history = [AnchorPoint.from_dict(a) for a in d.get("anchor_history", [])]
        tracker._refresh_history = [RefreshResult.from_dict(r) for r in d.get("refresh_history", [])]
        return tracker

    def save(self, path: Path | str) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2))

    @classmethod
    def load(cls, path: Path | str) -> "ForgettingTracker":
        return cls.from_dict(json.loads(Path(path).read_text()))
