"""
Starvex Drift Monitor - Track model drift over time.

Tracks:
- Daily/weekly F1 score trends
- Confidence distribution shifts
- Automatic retraining triggers
"""

from __future__ import annotations

import math
from collections import defaultdict, deque
from dataclasses import dataclass
from datetime import datetime, timedelta, date
from typing import Any, Deque, Dict, Iterable, List, Optional, Tuple


@dataclass
class PredictionRecord:
    timestamp: datetime
    y_true: bool
    y_pred: bool
    confidence: float


@dataclass
class DriftConfig:
    window_days: int = 7
    baseline_days: int = 28
    min_samples: int = 50
    min_f1: float = 0.85
    f1_drop_threshold: float = 0.08
    confidence_shift_threshold: float = 0.12
    histogram_bins: int = 10


class DriftMonitor:
    """
    In-memory drift monitor for binary classification.

    Provide ground-truth labels to compute accuracy metrics.
    """

    def __init__(self, config: Optional[DriftConfig] = None) -> None:
        self.config = config or DriftConfig()
        self._records: Deque[PredictionRecord] = deque()

    def record(
        self,
        y_true: bool,
        y_pred: bool,
        confidence: float,
        timestamp: Optional[datetime] = None,
    ) -> None:
        self._records.append(
            PredictionRecord(
                timestamp=timestamp or datetime.utcnow(),
                y_true=bool(y_true),
                y_pred=bool(y_pred),
                confidence=float(confidence),
            )
        )
        self._prune()

    def daily_f1_trend(self, days: int = 30) -> List[Tuple[date, float]]:
        records = list(self._records)
        cutoff = datetime.utcnow() - timedelta(days=days)
        by_day: Dict[date, List[PredictionRecord]] = defaultdict(list)
        for rec in records:
            if rec.timestamp >= cutoff:
                by_day[rec.timestamp.date()].append(rec)

        trend: List[Tuple[date, float]] = []
        for day in sorted(by_day.keys()):
            metrics = _compute_metrics(by_day[day])
            trend.append((day, metrics["f1"]))
        return trend

    def weekly_f1_trend(self, weeks: int = 8) -> List[Tuple[str, float]]:
        records = list(self._records)
        cutoff = datetime.utcnow() - timedelta(days=weeks * 7)
        by_week: Dict[str, List[PredictionRecord]] = defaultdict(list)
        for rec in records:
            if rec.timestamp >= cutoff:
                year, week, _ = rec.timestamp.isocalendar()
                key = f"{year}-W{week:02d}"
                by_week[key].append(rec)

        trend: List[Tuple[str, float]] = []
        for key in sorted(by_week.keys()):
            metrics = _compute_metrics(by_week[key])
            trend.append((key, metrics["f1"]))
        return trend

    def confidence_shift(self) -> float:
        baseline, recent = self._split_baseline_recent()
        if not baseline or not recent:
            return 0.0
        base_hist = _confidence_histogram(baseline, self.config.histogram_bins)
        recent_hist = _confidence_histogram(recent, self.config.histogram_bins)
        return _jensen_shannon_divergence(base_hist, recent_hist)

    def should_retrain(self) -> Dict[str, Any]:
        baseline, recent = self._split_baseline_recent()
        reasons: List[str] = []

        if len(recent) < self.config.min_samples:
            return {
                "should_retrain": False,
                "reason": "insufficient_samples",
                "details": {
                    "recent_samples": len(recent),
                    "required": self.config.min_samples,
                },
            }

        recent_metrics = _compute_metrics(recent)
        baseline_metrics = _compute_metrics(baseline) if baseline else None

        if recent_metrics["f1"] < self.config.min_f1:
            reasons.append("f1_below_min")

        if baseline_metrics:
            if (baseline_metrics["f1"] - recent_metrics["f1"]) > self.config.f1_drop_threshold:
                reasons.append("f1_drop")

        shift = self.confidence_shift()
        if shift > self.config.confidence_shift_threshold:
            reasons.append("confidence_shift")

        return {
            "should_retrain": bool(reasons),
            "reason": ",".join(reasons) if reasons else None,
            "details": {
                "recent_f1": recent_metrics["f1"],
                "baseline_f1": baseline_metrics["f1"] if baseline_metrics else None,
                "confidence_shift": shift,
                "recent_samples": len(recent),
                "baseline_samples": len(baseline),
            },
        }

    # -- internals -------------------------------------------------------

    def _split_baseline_recent(self) -> Tuple[List[PredictionRecord], List[PredictionRecord]]:
        now = datetime.utcnow()
        recent_start = now - timedelta(days=self.config.window_days)
        baseline_start = now - timedelta(days=self.config.baseline_days)

        baseline: List[PredictionRecord] = []
        recent: List[PredictionRecord] = []
        for rec in self._records:
            if rec.timestamp >= recent_start:
                recent.append(rec)
            elif rec.timestamp >= baseline_start:
                baseline.append(rec)
        return baseline, recent

    def _prune(self) -> None:
        cutoff = datetime.utcnow() - timedelta(days=self.config.baseline_days)
        while self._records and self._records[0].timestamp < cutoff:
            self._records.popleft()


def _compute_metrics(records: Iterable[PredictionRecord]) -> Dict[str, float]:
    records_list = list(records)
    if not records_list:
        return {"precision": 0.0, "recall": 0.0, "f1": 0.0, "accuracy": 0.0}

    tp = sum(1 for r in records_list if r.y_true and r.y_pred)
    fp = sum(1 for r in records_list if not r.y_true and r.y_pred)
    tn = sum(1 for r in records_list if not r.y_true and not r.y_pred)
    fn = sum(1 for r in records_list if r.y_true and not r.y_pred)

    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
    accuracy = (tp + tn) / len(records_list) if records_list else 0.0

    return {
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "accuracy": accuracy,
    }


def _confidence_histogram(records: Iterable[PredictionRecord], bins: int) -> List[float]:
    hist = [0.0 for _ in range(bins)]
    records_list = list(records)
    if not records_list:
        return hist
    for rec in records_list:
        idx = min(bins - 1, max(0, int(rec.confidence * bins)))
        hist[idx] += 1.0
    total = sum(hist) or 1.0
    return [h / total for h in hist]


def _jensen_shannon_divergence(p: List[float], q: List[float]) -> float:
    m = [(pi + qi) / 2.0 for pi, qi in zip(p, q)]
    return 0.5 * (_kl_divergence(p, m) + _kl_divergence(q, m))


def _kl_divergence(p: List[float], q: List[float]) -> float:
    eps = 1e-12
    total = 0.0
    for pi, qi in zip(p, q):
        if pi <= 0:
            continue
        total += pi * math.log((pi + eps) / (qi + eps))
    return total
