"""
BADAS Metrics — clean, modular metric API.

Two families:

- :class:`ClassificationMetric` — stateless, works for both single-window and
  sliding inference. Call :meth:`compute` with scores + labels.

- :class:`TemporalMetric` — stateful accumulator, sliding inference only.
  Call :meth:`update` once per video, then :meth:`compute` to aggregate,
  :meth:`reset` to start fresh.

Convenience::

    from badas.validation.metrics import (
        compute_classification_metrics,
        default_temporal_metrics,
        run_temporal_metrics,
    )

    # Classification
    result = compute_classification_metrics(scores, labels, threshold=0.75)

    # Temporal (sliding mode)
    metrics = default_temporal_metrics(threshold=0.75, max_gap_frames=8)
    result  = run_temporal_metrics(group_df, predictions, metrics,
                                   id_col="id", toe_col="time_of_event")
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from enum import Enum
from typing import Dict, List, Optional

import numpy as np
import pandas as pd


# =============================================================================
# SCOPE
# =============================================================================

class MetricScope(Enum):
    CLASSIFICATION = "classification"
    TEMPORAL       = "temporal"


# =============================================================================
# BASE CLASSES
# =============================================================================

class ClassificationMetric(ABC):
    """Stateless classification metric.

    Works for both *single-window* and *sliding* inference modes because it
    only needs a list of per-video scores and ground-truth labels.
    """

    scope: MetricScope = MetricScope.CLASSIFICATION

    @property
    @abstractmethod
    def keys(self) -> List[str]:
        """Names of the dict keys returned by :meth:`compute`."""

    @abstractmethod
    def compute(
        self,
        scores: List[float],
        labels: List[int],
        threshold: float,
    ) -> Dict[str, float]:
        """Compute metrics from per-video scores and binary labels.

        Args:
            scores:    Per-video prediction scores (e.g. max probability).
            labels:    Per-video ground-truth labels (0 or 1).
            threshold: Decision threshold for binary predictions.

        Returns:
            Dict mapping metric name → float value (NaN when not computable).
        """


class TemporalMetric(ABC):
    """Stateful temporal metric — sliding inference only.

    Usage::

        m = SomeTemporalMetric(...)
        for video in group:
            m.update(predictions, target, time_of_event)
        result = m.compute()
        m.reset()  # ready for next group
    """

    scope: MetricScope = MetricScope.TEMPORAL

    @property
    @abstractmethod
    def keys(self) -> List[str]:
        """Names of the dict keys returned by :meth:`compute`."""

    @abstractmethod
    def update(
        self,
        predictions: List[Dict],
        target: int,
        time_of_event: Optional[float],
    ) -> None:
        """Accumulate one video's data.

        Args:
            predictions:   List of ``{"timestamp": float, "probability": float,
                           "raw_probability": float}`` dicts from :meth:`predict_video`.
            target:        Ground-truth label (0 or 1).
            time_of_event: Timestamp of the event in seconds, or None if unknown.
        """

    @abstractmethod
    def compute(self) -> Dict[str, float]:
        """Aggregate accumulated data and return metric values."""

    @abstractmethod
    def reset(self) -> None:
        """Clear accumulated state so the metric can be reused for a new group."""


# =============================================================================
# PRIVATE HELPERS
# =============================================================================

def _find_sustained_alert(
    predictions: List[Dict],
    event_time: float,
    threshold: float,
    max_gap_frames: int,
) -> Optional[float]:
    """Return the timestamp of the first sustained alert before ``event_time``.

    A *sustained alert* starts when ``probability >= threshold`` and is allowed
    to dip below for up to ``max_gap_frames`` consecutive frames before being
    considered broken.

    Args:
        predictions:    Sliding-window predictions sorted by timestamp.
        event_time:     Event timestamp; only predictions before this are used.
        threshold:      Alert threshold.
        max_gap_frames: Max consecutive sub-threshold frames before reset.

    Returns:
        Timestamp of alert start, or ``None`` if no alert found.
    """
    alert_start = None
    gap_count = 0
    for pred in predictions:
        if pred["timestamp"] >= event_time:
            break
        if pred["probability"] >= threshold:
            if alert_start is None:
                alert_start = pred["timestamp"]
            gap_count = 0
        elif alert_start is not None:
            gap_count += 1
            if gap_count > max_gap_frames:
                alert_start = None
                gap_count = 0
    return alert_start


# =============================================================================
# CLASSIFICATION METRICS
# =============================================================================

class ClassificationMetrics(ClassificationMetric):
    """All standard binary classification metrics computed in one pass.

    Returns: ``auc_roc``, ``ap``, ``precision``, ``recall``, ``f1``,
    ``accuracy``, ``fpr``.  All NaN when only one class is present.
    """

    _KEYS = ["auc_roc", "ap", "precision", "recall", "f1", "accuracy", "fpr"]

    @property
    def keys(self) -> List[str]:
        return list(self._KEYS)

    def compute(
        self,
        scores: List[float],
        labels: List[int],
        threshold: float,
    ) -> Dict[str, float]:
        from sklearn.metrics import (
            roc_auc_score, average_precision_score, confusion_matrix,
        )

        nan = float("nan")
        if not scores:
            return {k: nan for k in self._KEYS}

        labels_arr = np.array(labels)
        scores_arr = np.array(scores)
        preds = (scores_arr >= threshold).astype(int)
        has_both = len(np.unique(labels_arr)) > 1

        auc_roc  = round(float(roc_auc_score(labels_arr, scores_arr)), 4) if has_both else nan
        ap       = round(float(average_precision_score(labels_arr, scores_arr)), 4) if has_both else nan
        accuracy = round(float(np.mean(preds == labels_arr)), 4)

        if has_both:
            cm = confusion_matrix(labels_arr, preds, labels=[0, 1])
            tn, fp, fn, tp = cm.ravel()
            precision = round(tp / (tp + fp), 4) if (tp + fp) > 0 else nan
            recall    = round(tp / (tp + fn), 4) if (tp + fn) > 0 else nan
            fpr       = round(fp / (fp + tn), 4) if (fp + tn) > 0 else nan
            f1        = round(2 * tp / (2 * tp + fp + fn), 4) if (2 * tp + fp + fn) > 0 else nan
        else:
            precision = recall = fpr = f1 = nan

        return {
            "auc_roc":   auc_roc,
            "ap":        ap,
            "precision": precision,
            "recall":    recall,
            "f1":        f1,
            "accuracy":  accuracy,
            "fpr":       fpr,
        }


# =============================================================================
# TEMPORAL METRICS
# =============================================================================

class TimeToAlertMetric(TemporalMetric):
    """Mean Time To Alert (MTTA) and Early Warning Recall (EWR).

    For each positive video with a known time-of-event, finds the first
    *sustained* alert before the event and computes the lead time
    ``time_of_event − alert_timestamp``.

    Returns: ``mtta``, ``mtta_std``, ``early_warning_recall``, ``n_with_alert``.
    """

    _KEYS = ["mtta", "mtta_std", "early_warning_recall", "n_with_alert"]

    def __init__(self, threshold: float = 0.75, max_gap_frames: int = 8):
        self.threshold = threshold
        self.max_gap_frames = max_gap_frames
        self._lead_times: List[float] = []
        self._total_positives: int = 0

    @property
    def keys(self) -> List[str]:
        return list(self._KEYS)

    def reset(self) -> None:
        self._lead_times = []
        self._total_positives = 0

    def update(
        self,
        predictions: List[Dict],
        target: int,
        time_of_event: Optional[float],
    ) -> None:
        if not (target and time_of_event is not None
                and np.isfinite(time_of_event) and time_of_event > 0):
            return
        if not predictions:
            return
        self._total_positives += 1
        alert_time = _find_sustained_alert(
            predictions, time_of_event, self.threshold, self.max_gap_frames
        )
        if alert_time is not None:
            lead_time = time_of_event - alert_time
            if lead_time > 0:
                self._lead_times.append(lead_time)

    def compute(self) -> Dict[str, float]:
        nan = float("nan")
        lt = self._lead_times
        mtta     = round(float(np.mean(lt)), 2)  if lt else nan
        mtta_std = round(float(np.std(lt)), 2)   if lt else nan
        ewr      = round(len(lt) / self._total_positives, 4) if self._total_positives > 0 else nan
        return {
            "mtta":                 mtta,
            "mtta_std":             mtta_std,
            "early_warning_recall": ewr,
            "n_with_alert":         len(lt),
        }


class BrierScoreMetric(TemporalMetric):
    """Binned Brier score averaged across videos.

    Each video is split into 1-second bins; bins near the event get target=1,
    all others get target=0.  The per-video Brier score is the mean squared
    error over valid bins.

    Returns: ``brier``.
    """

    _KEYS = ["brier"]

    def __init__(
        self,
        bin_duration: float = 1.0,
        event_before: float = 2.0,
        event_after:  float = 1.0,
    ):
        self.bin_duration = bin_duration
        self.event_before = event_before
        self.event_after  = event_after
        self._scores: List[float] = []

    @property
    def keys(self) -> List[str]:
        return list(self._KEYS)

    def reset(self) -> None:
        self._scores = []

    def update(
        self,
        predictions: List[Dict],
        target: int,
        time_of_event: Optional[float],
    ) -> None:
        score = self._compute_video(predictions, bool(target), time_of_event)
        if score is not None:
            self._scores.append(score)

    def _compute_video(
        self,
        predictions: List[Dict],
        target: bool,
        time_of_event: Optional[float],
    ) -> Optional[float]:
        if not predictions:
            return None
        if target and (time_of_event is None
                       or not np.isfinite(time_of_event)
                       or time_of_event <= 0):
            return None

        timestamps = [p["timestamp"] for p in predictions]
        min_t, max_t = min(timestamps), max(timestamps)
        bin_edges = np.arange(min_t, max_t + self.bin_duration, self.bin_duration)
        n_bins = len(bin_edges) - 1
        if n_bins <= 0:
            return None

        event_start = (time_of_event - self.event_before) if target else None
        event_end   = (time_of_event + self.event_after)  if target else None

        brier_sum, n_valid = 0.0, 0
        for i in range(n_bins):
            bin_probs = [p["probability"] for p in predictions
                         if bin_edges[i] <= p["timestamp"] < bin_edges[i + 1]]
            if not bin_probs:
                continue
            bin_center = (bin_edges[i] + bin_edges[i + 1]) / 2.0
            bin_target = (1.0 if (target and event_start <= bin_center <= event_end)
                          else 0.0)
            brier_sum += (float(np.mean(bin_probs)) - bin_target) ** 2
            n_valid += 1

        return brier_sum / n_valid if n_valid > 0 else None

    def compute(self) -> Dict[str, float]:
        val = round(float(np.mean(self._scores)), 4) if self._scores else float("nan")
        return {"brier": val}


class SignalNoiseMetric(TemporalMetric):
    """Signal noise metrics: smoothness and stability of the probability signal.

    Per-video scalars: mean-abs-change/s (delta) and std for both raw and
    smoothed probability signals, plus threshold-crossing count.
    Group-level result: median of each scalar across videos.

    Returns: ``noise_delta_raw_median``, ``noise_delta_smooth_median``,
    ``noise_std_raw_median``, ``noise_std_smooth_median``,
    ``noise_crossings_median``.
    """

    _KEYS = [
        "noise_delta_raw_median",
        "noise_delta_smooth_median",
        "noise_std_raw_median",
        "noise_std_smooth_median",
        "noise_crossings_median",
    ]

    def __init__(self, crossing_threshold: float = 0.5):
        self.crossing_threshold = crossing_threshold
        self._records: List[Dict[str, float]] = []

    @property
    def keys(self) -> List[str]:
        return list(self._KEYS)

    def reset(self) -> None:
        self._records = []

    def update(
        self,
        predictions: List[Dict],
        target: int,
        time_of_event: Optional[float],
    ) -> None:
        record = self._compute_video(predictions)
        if record is not None:
            self._records.append(record)

    def _compute_video(self, predictions: List[Dict]) -> Optional[Dict[str, float]]:
        if not predictions:
            return None
        ts     = np.array([p["timestamp"]      for p in predictions], dtype=float)
        raw    = np.array([p["raw_probability"] for p in predictions], dtype=float)
        smooth = np.array([p["probability"]     for p in predictions], dtype=float)
        n = len(ts)
        thr = self.crossing_threshold
        crossings = int(np.sum(np.diff((smooth >= thr).astype(int)) != 0)) if n > 1 else 0
        std_raw    = float(np.std(raw))
        std_smooth = float(np.std(smooth))
        if n < 2:
            delta_raw = delta_smooth = float("nan")
        else:
            dt = np.maximum(np.diff(ts), 0.01)
            delta_raw    = float(np.mean(np.abs(np.diff(raw))    / dt))
            delta_smooth = float(np.mean(np.abs(np.diff(smooth)) / dt))
        return {
            "delta_raw":    delta_raw,
            "delta_smooth": delta_smooth,
            "std_raw":      std_raw,
            "std_smooth":   std_smooth,
            "crossings":    float(crossings),
        }

    def compute(self) -> Dict[str, float]:
        if not self._records:
            return {}
        df = pd.DataFrame(self._records)
        return {
            f"noise_{col}_median": round(float(df[col].dropna().median()), 4)
            for col in df.columns
            if len(df[col].dropna()) > 0
        }


# =============================================================================
# FACTORIES
# =============================================================================

def default_classification_metrics() -> List[ClassificationMetric]:
    """Return the default set of classification metrics."""
    return [ClassificationMetrics()]


def default_temporal_metrics(
    threshold: float = 0.75,
    max_gap_frames: int = 8,
    brier_bin_duration: float = 1.0,
    brier_event_before: float = 2.0,
    brier_event_after: float = 1.0,
    noise_crossing_threshold: float = 0.5,
) -> List[TemporalMetric]:
    """Return the default set of temporal metrics, configured from evaluator params."""
    return [
        TimeToAlertMetric(threshold=threshold, max_gap_frames=max_gap_frames),
        BrierScoreMetric(
            bin_duration=brier_bin_duration,
            event_before=brier_event_before,
            event_after=brier_event_after,
        ),
        SignalNoiseMetric(crossing_threshold=noise_crossing_threshold),
    ]


# =============================================================================
# CONVENIENCE FUNCTIONS
# =============================================================================

def compute_classification_metrics(
    scores: List[float],
    labels: List[int],
    threshold: float,
) -> Dict[str, float]:
    """Stateless convenience wrapper around :class:`ClassificationMetrics`.

    Equivalent to ``ClassificationMetrics().compute(scores, labels, threshold)``.
    """
    return ClassificationMetrics().compute(scores, labels, threshold)


def run_temporal_metrics(
    group_df: "pd.DataFrame",
    predictions: Dict[str, List[Dict]],
    metrics: List[TemporalMetric],
    id_col: str,
    toe_col: str,
) -> Dict[str, float]:
    """Reset, update per-video, then aggregate all temporal metrics.

    Args:
        group_df:    DataFrame subset for one evaluation group; must have
                     columns ``id_col``, ``target``, and ``toe_col``.
        predictions: ``{video_id: [prediction_dict, ...]}`` map.
        metrics:     List of :class:`TemporalMetric` instances to run.
        id_col:      Column name for video IDs.
        toe_col:     Column name for time-of-event values.

    Returns:
        Merged dict of all metric outputs.
    """
    for m in metrics:
        m.reset()

    for _, row in group_df.iterrows():
        preds = predictions.get(str(row[id_col]), [])
        toe_raw = row.get(toe_col)
        toe = (float(toe_raw)
               if toe_raw is not None and pd.notna(toe_raw)
               else None)
        for m in metrics:
            m.update(preds, int(bool(row["target"])), toe)

    result: Dict[str, float] = {}
    for m in metrics:
        result.update(m.compute())
    return result
