"""
BADAS Evaluator - Modular, multi-GPU evaluation pipeline.

Three-step pipeline:
1. Inference  → predictions.json  (multi-GPU parallel, batched for single_window)
2. Metrics    → metrics_by_group.csv
   - single_window: classification only (AUC-ROC, AP, precision/recall/F1/FPR @ threshold)
   - sliding:       classification + temporal (MTTA, EWR, Brier, noise)
3. Visualizations → plots/

Usage:
    from badas.validation.evaluator import Evaluator, EvaluatorConfig

    config = EvaluatorConfig(
        checkpoint_path="model.ckpt",
        labels_csv="test.csv",
        videos_dir="/path/to/videos",
        num_gpus=-1,   # use all GPUs
        batch_size=32, # videos per forward pass (single_window only)
    )
    evaluator = Evaluator(config)
    evaluator.run()

    # Or run steps individually:
    evaluator.run_inference()
    evaluator.compute_metrics()
    evaluator.generate_visualizations()
"""

import json
import os
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple

import numpy as np
import pandas as pd

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


# =============================================================================
# MODULE-LEVEL HELPERS  (picklable — safe for multiprocessing workers)
# =============================================================================

def _load_frames_by_original_index(video_path: str, orig_indices: List[int]) -> List[np.ndarray]:
    """Load specific frames by original index. Returns BGR uint8 numpy arrays."""
    try:
        from decord import VideoReader, cpu
        import cv2
        vr = VideoReader(video_path, ctx=cpu(0))
        batch = vr.get_batch(orig_indices).asnumpy()  # decord returns RGB
        return [cv2.cvtColor(f, cv2.COLOR_RGB2BGR) for f in batch]
    except Exception:
        import cv2
        cap = cv2.VideoCapture(video_path)
        frames = []
        try:
            for idx in orig_indices:
                cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
                ret, frame = cap.read()
                if ret:
                    frames.append(frame)
        finally:
            cap.release()
        return frames


def _load_single_window_frames(
    video_path: str,
    toe: Optional[float],
    window_size: int,
    target_fps: float,
) -> Tuple[Optional[List[np.ndarray]], Optional[float]]:
    """Load window_size BGR frames anchored at toe (or last frame if None).

    Returns:
        (frames, timestamp) or (None, None) on failure.
    """
    try:
        from badas.inference.frame_sampling import get_video_info
        info = get_video_info(video_path)
        source_fps = info["fps"]
        total_frames = info["total_frames"]

        if total_frames == 0 or source_fps == 0:
            return None, None

        frame_interval = source_fps / target_fps
        anchor_frame = (
            min(int(toe * source_fps), total_frames - 1)
            if toe and toe > 0
            else total_frames - 1
        )

        window_span = (window_size - 1) * frame_interval
        start_frame = max(0.0, anchor_frame - window_span)
        orig_indices = [
            min(int(start_frame + i * frame_interval), total_frames - 1)
            for i in range(window_size)
        ]

        while len(orig_indices) < window_size:
            orig_indices.insert(0, orig_indices[0])

        frames = _load_frames_by_original_index(video_path, orig_indices)
        if len(frames) < window_size:
            return None, None

        return frames[:window_size], anchor_frame / source_fps
    except Exception:
        return None, None


def _run_single_window_chunk(
    video_items: List[Tuple],
    predictor,
    batch_size: int,
    num_io_threads: int,
) -> Dict[str, List[Dict]]:
    """Process videos in single-window mode: parallel I/O + batched GPU inference.

    Args:
        video_items: list of (video_id, video_path, toe)
        predictor:   BADAS instance
        batch_size:  max videos per forward pass
        num_io_threads: parallel frame decoders

    Returns:
        predictions dict {video_id: [{timestamp, probability, raw_probability}]}
    """
    def _load(item):
        vid, path, toe = item
        return vid, *_load_single_window_frames(path, toe, predictor.window_size, predictor.target_fps)

    # Parallel I/O: decode all videos concurrently
    with ThreadPoolExecutor(max_workers=num_io_threads) as exe:
        loaded = list(exe.map(_load, video_items))

    predictions = {}
    batch_vids, batch_frames, batch_ts = [], [], []

    def _flush():
        if not batch_vids:
            return
        probs = predictor.predict_batch(batch_frames)
        for vid, prob, ts in zip(batch_vids, probs, batch_ts):
            predictions[vid] = [{"timestamp": ts, "probability": prob, "raw_probability": prob}]
        batch_vids.clear(); batch_frames.clear(); batch_ts.clear()

    for vid, frames, ts in loaded:
        if frames is None:
            predictions[vid] = []
            continue
        batch_vids.append(vid)
        batch_frames.append(frames)
        batch_ts.append(ts)
        if len(batch_vids) >= batch_size:
            _flush()

    _flush()
    return predictions


def _run_sliding_chunk(
    video_items: List[Tuple],
    predictor,
    stride: int,
    gpu_id: int = 0,
    log_every: int = 10,
) -> Dict[str, List[Dict]]:
    """Process videos in sliding-window mode: sequential predict_video per video."""
    predictions = {}
    n = len(video_items)
    t_chunk = datetime.now()
    for i, (vid, path, _) in enumerate(video_items):
        try:
            preds = predictor.predict_video(path, stride=stride, verbose=False, quiet=True)
            predictions[vid] = [
                {"timestamp": p["timestamp"], "probability": p["probability"], "raw_probability": p["raw_probability"]}
                for p in preds
            ]
        except Exception:
            predictions[vid] = []
        if (i + 1) % log_every == 0 or (i + 1) == n:
            elapsed = (datetime.now() - t_chunk).total_seconds()
            rate = (i + 1) / elapsed if elapsed > 0 else float("inf")
            eta = (n - i - 1) / rate if rate > 0 else 0
            print(
                f"[{datetime.now().strftime('%H:%M:%S')}] GPU {gpu_id}  {i+1}/{n} videos"
                f"  ({rate:.1f} vid/s  ETA {eta/60:.1f}min)",
                flush=True,
            )
    return predictions


def _inference_worker(args) -> int:
    """Module-level worker for parallel multi-GPU inference.

    Spawned as a subprocess by _run_parallel. Loads model on the assigned GPU,
    processes its video chunk, and writes predictions to output_path.
    """
    rank, video_items, checkpoint_path, gpu_id, mode, batch_size, num_io_threads, stride, smoothing, num_frames, model_fps, output_path = args

    import os
    import torch
    from badas.inference import BADAS, BADASConfig, SmoothingConfig

    # Must be set before vLLM initializes (CUDA_VISIBLE_DEVICES restricts visible GPUs)
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    t0 = datetime.now()
    # compile is worthwhile for sliding mode (hundreds of fwd passes per video × many videos)
    # but adds ~30-60s startup overhead with no benefit for single_window batched mode
    use_compile = (mode == "sliding")
    badas_kwargs = dict(
        smoothing=SmoothingConfig(enabled=smoothing),
        use_compile=use_compile,
        num_frames=num_frames,
        quiet=True, show_progress=False,
    )
    if model_fps is not None:
        badas_kwargs["model_fps"] = model_fps
    predictor = BADAS(checkpoint_path, config=BADASConfig(**badas_kwargs))
    print(f"[{datetime.now().strftime('%H:%M:%S')}] GPU {gpu_id} worker ready, {len(video_items)} videos", flush=True)

    if mode == "single_window":
        predictions = _run_single_window_chunk(video_items, predictor, batch_size, num_io_threads)
    else:
        predictions = _run_sliding_chunk(video_items, predictor, stride, gpu_id=gpu_id)

    elapsed = (datetime.now() - t0).total_seconds()
    print(f"[{datetime.now().strftime('%H:%M:%S')}] GPU {gpu_id} done: {len(predictions)} videos in {elapsed:.1f}s", flush=True)

    with open(output_path, 'w') as f:
        json.dump(predictions, f)

    return len(predictions)


# =============================================================================
# EVALUATOR CONFIG
# =============================================================================

@dataclass
class EvaluatorConfig:
    """Configuration for the Evaluator."""

    # Required
    checkpoint_path: str = None
    labels_csv: str = None

    # Video location
    videos_dir: str = None
    video_id_column: str = "id"
    video_extension: str = ".mov"

    # Inference settings
    stride: int = 1
    smoothing: bool = True
    device: str = None  # deprecated — use gpu_ids instead

    # Metrics settings
    threshold: float = 0.75
    max_gap_frames: int = 8       # sustained-alert tolerance (sliding only)
    group_column: str = "test_group"
    time_of_event_column: str = "time_of_event"

    # Brier score (sliding only)
    brier_bin_duration: float = 1.0
    brier_event_before: float = 2.0
    brier_event_after: float = 1.0

    # Noise metrics (sliding only)
    noise_crossing_threshold: float = 0.5

    # Output
    output_dir: str = None

    # Inference mode
    # "sliding"       — full sliding-window over the video
    # "single_window" — one window per video anchored at time_of_event
    inference_mode: str = "sliding"

    # Dataset format
    # "testset" — internal test set: CSV has id, target, time_of_event, group
    # "kaggle"  — Nexar competition: CSV has id, target, group (0/1/2), Usage
    dataset_format: str = "testset"

    # Kaggle: save submission CSV
    submission_path: Optional[str] = None

    # Override frame count (None = use checkpoint value)
    num_frames: Optional[int] = None

    # Override model FPS (None = use checkpoint/default value; set to 4 for VLM models)
    model_fps: Optional[int] = None

    # Performance: multi-GPU + batched inference
    num_gpus: int = -1          # -1 = auto-detect all available GPUs
    gpu_ids: Optional[List[int]] = None  # explicit list; overrides num_gpus
    batch_size: int = 32        # videos per forward pass (single_window only)
    num_io_threads: int = 16    # parallel frame decoders per worker

    # Logging
    quiet: bool = False
    show_progress: bool = True

    def __post_init__(self):
        if self.output_dir is None:
            self.output_dir = f"evaluation_{datetime.now().strftime('%Y%m%d_%H%M%S')}"


# =============================================================================
# EVALUATOR
# =============================================================================

class Evaluator:
    """
    Modular evaluation pipeline for BADAS models.

    Step 1 — Inference:   run model on videos → predictions.json
    Step 2 — Metrics:     compute metrics from predictions + labels → CSV
    Step 3 — Visualize:   generate plots → plots/

    Metric types by inference_mode:
        single_window → classification only (AUC-ROC, AP, precision/recall/F1/FPR)
        sliding       → classification + temporal (MTTA, EWR, Brier, noise)
    """

    def __init__(self, config: EvaluatorConfig):
        self.config = config
        self._predictions: Optional[Dict[str, List[Dict]]] = None
        self._labels_df: Optional[pd.DataFrame] = None
        self._metrics_df: Optional[pd.DataFrame] = None

        self.output_dir = Path(config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.predictions_path = self.output_dir / "predictions.json"
        self.metrics_path = self.output_dir / "metrics_by_group.csv"
        self.plots_dir = self.output_dir / "plots"

        self._temporal_metrics = default_temporal_metrics(
            threshold=config.threshold,
            max_gap_frames=config.max_gap_frames,
            brier_bin_duration=config.brier_bin_duration,
            brier_event_before=config.brier_event_before,
            brier_event_after=config.brier_event_after,
            noise_crossing_threshold=config.noise_crossing_threshold,
        )

    def _log(self, msg: str):
        if not self.config.quiet:
            print(msg, flush=True)

    def _load_labels_df(self) -> pd.DataFrame:
        df = pd.read_csv(self.config.labels_csv)
        if self.config.dataset_format == "kaggle":
            col = self.config.video_id_column
            df[col] = df[col].apply(lambda x: f"{int(x):05d}")
        # Normalize target to int 0/1 — handles "True"/"False" strings, booleans, floats
        if "target" in df.columns:
            def _to_label(v):
                if isinstance(v, bool):
                    return int(v)
                if isinstance(v, str):
                    return 1 if v.strip().lower() in ("1", "true", "yes") else 0
                return int(bool(v))
            df["target"] = df["target"].apply(_to_label)
        return df

    def _detect_gpu_ids(self) -> List[int]:
        """Return list of GPU IDs to use, based on config and availability."""
        import torch
        if not torch.cuda.is_available():
            return []
        n = torch.cuda.device_count()
        if self.config.gpu_ids:
            return [g for g in self.config.gpu_ids if g < n]
        if self.config.num_gpus == -1:
            return list(range(n))
        return list(range(min(self.config.num_gpus, n)))

    # =========================================================================
    # STEP 1: INFERENCE
    # =========================================================================

    def run_inference(self, video_ids: Optional[List[str]] = None) -> Dict[str, List[Dict]]:
        """Run inference on all videos, save predictions to JSON.

        Uses all available GPUs in parallel (multi-GPU) with batched forward
        passes for single_window mode. Falls back to single-GPU in-process
        when only one GPU is available.
        """
        from badas.inference import BADAS, BADASConfig, SmoothingConfig

        self._log("=" * 60)
        self._log("STEP 1: RUNNING INFERENCE")
        self._log("=" * 60)

        labels_df = self._load_labels_df()
        self._labels_df = labels_df

        if video_ids is not None:
            labels_df = labels_df[labels_df[self.config.video_id_column].isin(video_ids)]

        # Build (video_id, video_path, toe) triples
        video_items = []
        for _, row in labels_df.iterrows():
            vid = row[self.config.video_id_column]
            path = self._get_video_path(str(vid))
            if path is None or not Path(path).exists():
                continue
            toe_raw = row.get(self.config.time_of_event_column)
            toe = float(toe_raw) if toe_raw is not None and pd.notna(toe_raw) and float(toe_raw) > 0 else None
            video_items.append((str(vid), path, toe))

        skipped = len(labels_df) - len(video_items)
        self._log(f"Videos to process: {len(video_items)}  (skipped {skipped} missing files)")
        self._log(f"Checkpoint: {self.config.checkpoint_path}")
        self._log(f"Mode: {self.config.inference_mode}  |  batch_size={self.config.batch_size}  |  io_threads={self.config.num_io_threads}")

        gpu_ids = self._detect_gpu_ids()
        self._log(f"GPUs: {gpu_ids if gpu_ids else 'CPU'}")

        if len(gpu_ids) > 1:
            predictions = self._run_parallel(video_items, gpu_ids)
        else:
            # Single-GPU or CPU: load model in-process
            import torch
            if gpu_ids:
                torch.cuda.set_device(gpu_ids[0])

            self._log("Loading model...")
            use_compile = (self.config.inference_mode == "sliding")
            badas_kwargs = dict(
                smoothing=SmoothingConfig(enabled=self.config.smoothing),
                use_compile=use_compile,
                num_frames=self.config.num_frames,
                quiet=True, show_progress=False,
            )
            if self.config.model_fps is not None:
                badas_kwargs["model_fps"] = self.config.model_fps
            predictor = BADAS(self.config.checkpoint_path, config=BADASConfig(**badas_kwargs))
            self._log("Model loaded")

            if self.config.inference_mode == "single_window":
                predictions = _run_single_window_chunk(
                    video_items, predictor,
                    self.config.batch_size, self.config.num_io_threads,
                )
            else:
                predictions = _run_sliding_chunk(
                    video_items, predictor, self.config.stride,
                    gpu_id=gpu_ids[0] if gpu_ids else 0,
                )

            self._log(f"Completed: {len(predictions)} processed")

        self._predictions = predictions
        self._save_predictions(predictions)
        return predictions

    def _run_parallel(self, video_items: List[Tuple], gpu_ids: List[int]) -> Dict[str, List[Dict]]:
        """Spawn one subprocess per GPU, merge results."""
        import multiprocessing as mp
        from concurrent.futures import ProcessPoolExecutor

        n = len(gpu_ids)
        # Round-robin split to balance load
        chunks = [video_items[i::n] for i in range(n)]
        part_paths = [str(self.output_dir / f"predictions_part{i}.json") for i in range(n)]

        worker_args = [
            (
                rank, chunks[rank], self.config.checkpoint_path, gpu_ids[rank],
                self.config.inference_mode, self.config.batch_size, self.config.num_io_threads,
                self.config.stride, self.config.smoothing, self.config.num_frames,
                self.config.model_fps, part_paths[rank],
            )
            for rank in range(n)
        ]

        self._log(f"Spawning {n} GPU workers...")
        ctx = mp.get_context("spawn")
        with ProcessPoolExecutor(max_workers=n, mp_context=ctx) as pool:
            counts = list(pool.map(_inference_worker, worker_args))

        self._log(f"Parallel inference done: {sum(counts)} total videos")

        # Merge partial results
        predictions: Dict[str, List[Dict]] = {}
        for path in part_paths:
            with open(path) as f:
                predictions.update(json.load(f))
            Path(path).unlink()  # clean up partials

        return predictions

    def _get_video_path(self, video_id: str) -> Optional[str]:
        if self.config.videos_dir is None:
            return None
        base = Path(self.config.videos_dir) / video_id
        for ext in [".mov", ".mp4", ".avi", ".mkv", ""]:
            p = base.with_suffix(ext) if ext else base
            if p.exists():
                return str(p)
            p2 = Path(self.config.videos_dir) / f"{video_id}{ext}"
            if p2.exists():
                return str(p2)
        return None

    def _save_predictions(self, predictions: Dict[str, List[Dict]]):
        with open(self.predictions_path, "w") as f:
            json.dump(predictions, f, indent=2)
        self._log(f"Saved predictions to: {self.predictions_path}")

    def load_predictions(self) -> Dict[str, List[Dict]]:
        with open(self.predictions_path) as f:
            self._predictions = json.load(f)
        return self._predictions

    # =========================================================================
    # STEP 2: METRICS
    # =========================================================================

    def compute_metrics(self) -> pd.DataFrame:
        """Compute metrics and save to CSV.

        Dispatches by dataset_format:
          testset → groups by group_column, mode-aware metrics
          kaggle  → groups by Usage/TTA, AP per split
        """
        self._log("\n" + "=" * 60)
        self._log("STEP 2: COMPUTING METRICS")
        self._log("=" * 60)

        if self._predictions is None:
            self.load_predictions()
        if self._labels_df is None:
            self._labels_df = self._load_labels_df()

        if self.config.dataset_format == "kaggle":
            return self._compute_kaggle_metrics()

        labels_df = self._labels_df
        predictions = self._predictions

        self._log(f"Labels: {len(labels_df)}  |  Predictions: {len(predictions)}")
        self._log(f"Mode: {self.config.inference_mode}  |  Threshold: {self.config.threshold}")

        groups = labels_df[self.config.group_column].dropna().unique()
        self._log(f"Groups: {list(groups)}")

        results = []
        for group in list(groups) + ["__overall__"]:
            if group == "__overall__":
                group_df, group_name = labels_df, "OVERALL"
            else:
                group_df = labels_df[labels_df[self.config.group_column] == group]
                group_name = group

            metrics = self._compute_group_metrics(group_df, predictions)
            metrics["group"] = group_name
            results.append(metrics)

        metrics_df = pd.DataFrame(results)

        # Column ordering
        base_cols = ["group", "n_total", "n_positive", "n_negative",
                     "auc_roc", "ap", "precision", "recall", "f1", "accuracy", "fpr"]
        temporal_cols = ["brier", "mtta", "mtta_std", "early_warning_recall", "n_with_alert",
                         "noise_delta_raw_median", "noise_delta_smooth_median",
                         "noise_std_raw_median", "noise_std_smooth_median", "noise_crossings_median"]
        ordered = base_cols + temporal_cols
        metrics_df = metrics_df[[c for c in ordered if c in metrics_df.columns]]

        self._metrics_df = metrics_df
        metrics_df.to_csv(self.metrics_path, index=False)
        self._log(f"\nSaved metrics to: {self.metrics_path}")
        self._print_metrics_summary(metrics_df)

        return metrics_df

    def _compute_group_metrics(
        self,
        group_df: pd.DataFrame,
        predictions: Dict[str, List[Dict]],
    ) -> Dict[str, Any]:
        """Mode-aware metrics for one group.

        Always computes: classification (AUC-ROC, AP, P/R/F1/FPR @ threshold).
        Additionally for sliding mode: temporal (MTTA, EWR, Brier, noise).
        """
        scores, labels = [], []
        for _, row in group_df.iterrows():
            vid = str(row[self.config.video_id_column])
            preds = predictions.get(vid, [])
            labels.append(int(bool(row["target"])))
            scores.append(max(p["probability"] for p in preds) if preds else 0.0)

        n_total = len(labels)
        n_positive = sum(labels)

        metrics: Dict[str, Any] = {
            "n_total":    n_total,
            "n_positive": n_positive,
            "n_negative": n_total - n_positive,
            **compute_classification_metrics(scores, labels, self.config.threshold),
        }

        if self.config.inference_mode == "sliding":
            metrics.update(run_temporal_metrics(
                group_df, predictions, self._temporal_metrics,
                id_col=self.config.video_id_column,
                toe_col=self.config.time_of_event_column,
            ))

        return metrics


    def _compute_kaggle_metrics(self) -> pd.DataFrame:
        """Nexar/Kaggle metrics: AP per TTA group (0.5s/1.0s/1.5s) × Usage split."""
        from sklearn.metrics import average_precision_score

        labels_df = self._labels_df
        predictions = self._predictions
        id_col = self.config.video_id_column
        group_to_tta = {0: "0.5s", 1: "1.0s", 2: "1.5s"}

        self._log(f"Labels: {len(labels_df)}  |  Predictions: {len(predictions)}")

        scores = {
            str(vid): max(p["probability"] for p in preds)
            for vid, preds in predictions.items() if preds
        }

        if self.config.submission_path:
            sub = [{"id": vid, "target": score} for vid, score in sorted(scores.items())]
            pd.DataFrame(sub).to_csv(self.config.submission_path, index=False)
            self._log(f"Saved Kaggle submission to: {self.config.submission_path}")

        df = labels_df.copy()
        df["_score"] = df[id_col].astype(str).map(scores)
        missing = df["_score"].isna().sum()
        if missing > 0:
            self._log(f"  Warning: {missing} videos have no predictions")
        df = df.dropna(subset=["_score"])

        splits = (
            [("Public", df[df["Usage"] == "Public"]),
             ("Private", df[df["Usage"] == "Private"]),
             ("OVERALL", df)]
            if "Usage" in df.columns
            else [("OVERALL", df)]
        )

        results = []
        for split_name, split_df in splits:
            if len(split_df) == 0:
                continue

            row: Dict[str, Any] = {
                "group":      split_name,
                "n_total":    len(split_df),
                "n_positive": int(split_df["target"].sum()),
                "n_negative": int((split_df["target"] == 0).sum()),
            }

            # Generic classification metrics for the whole split
            row.update(compute_classification_metrics(
                split_df["_score"].tolist(),
                split_df["target"].tolist(),
                self.config.threshold,
            ))

            # Kaggle-specific: AP per TTA group
            tta_aps = {}
            if "group" in split_df.columns:
                for group_int, tta_label in group_to_tta.items():
                    g_df = split_df[split_df["group"] == group_int]
                    if len(g_df) > 0 and g_df["target"].nunique() > 1:
                        tta_aps[tta_label] = round(
                            float(average_precision_score(g_df["target"], g_df["_score"])), 4
                        )
                    else:
                        tta_aps[tta_label] = float('nan')

                row["mean_ap"] = round(float(np.nanmean(list(tta_aps.values()))), 4)
                for lbl, ap in tta_aps.items():
                    row[f"ap_{lbl}"] = ap

            results.append(row)

        metrics_df = pd.DataFrame(results)
        self._metrics_df = metrics_df
        metrics_df.to_csv(self.metrics_path, index=False)
        self._log(f"\nSaved metrics to: {self.metrics_path}")

        self._log("\n" + "-" * 60)
        self._log("KAGGLE METRICS SUMMARY")
        self._log("-" * 60)
        tta_cols = [c for c in metrics_df.columns if c.startswith("ap_")]
        for _, row in metrics_df.iterrows():
            tta_str = "  ".join(f"{c.replace('ap_','')}={row[c]:.4f}" for c in tta_cols if c in row)
            mean_ap = row.get("mean_ap", float('nan'))
            self._log(
                f"  {str(row['group']):<10s} | n={row['n_total']:4d} | "
                f"AUC={row['auc_roc']:.4f}  AP={row['ap']:.4f}  mAP={mean_ap:.4f}  [{tta_str}]"
            )

        return metrics_df

    # =========================================================================
    # METRICS SUMMARY
    # =========================================================================

    def _print_metrics_summary(self, metrics_df: pd.DataFrame):
        def _v(r, col):
            val = r.get(col, float("nan"))
            return None if (isinstance(val, float) and np.isnan(val)) else float(val)

        self._log("\n" + "-" * 60)
        self._log("METRICS SUMMARY")
        self._log("-" * 60)

        # ── Classification columns present in this run ─────────────────────
        clf_cols = [c for c in ["auc_roc", "ap", "precision", "recall", "f1", "accuracy", "fpr"]
                    if c in metrics_df.columns]
        # ── Temporal columns ───────────────────────────────────────────────
        tmp_cols = [c for c in ["mtta", "mtta_std", "early_warning_recall", "n_with_alert", "brier"]
                    if c in metrics_df.columns]
        # ── Noise columns ──────────────────────────────────────────────────
        noise_cols = [c for c in metrics_df.columns if c.startswith("noise_")]

        tmp_labels = {
            "mtta":                 "MTTA (s)",
            "mtta_std":             "MTTA std",
            "early_warning_recall": "EWR",
            "n_with_alert":         "n_alert",
            "brier":                "Brier",
        }
        noise_labels = {c: c.replace("noise_", "").replace("_median", "") for c in noise_cols}

        # ── OVERALL block ──────────────────────────────────────────────────
        overall = metrics_df[metrics_df["group"] == "OVERALL"]
        if len(overall) > 0:
            r = overall.iloc[0]
            self._log(f"\nOVERALL  (n={int(r['n_total'])}  pos={int(r['n_positive'])}  neg={int(r['n_negative'])})")
            if clf_cols:
                self._log("  Classification:")
                for col in clf_cols:
                    v = _v(r, col)
                    if v is not None:
                        self._log(f"    {col:<12s}: {v:.4f}")
            if tmp_cols:
                self._log("  Temporal:")
                for col in tmp_cols:
                    v = _v(r, col)
                    if v is not None:
                        lbl = tmp_labels.get(col, col)
                        fmt = ".0f" if col == "n_with_alert" else ".4f"
                        self._log(f"    {lbl:<12s}: {v:{fmt}}")
            if noise_cols:
                self._log("  Noise (median across videos):")
                for col in noise_cols:
                    v = _v(r, col)
                    if v is not None:
                        self._log(f"    {noise_labels[col]:<20s}: {v:.4f}")

        # ── Per-group table ────────────────────────────────────────────────
        group_rows = metrics_df[metrics_df["group"] != "OVERALL"].copy()
        if len(group_rows) == 0:
            return

        sort_col = next((c for c in ["auc_roc", "f1"] if c in group_rows.columns), None)
        if sort_col:
            group_rows = group_rows.sort_values(sort_col, ascending=False)

        # Build header
        per_group_show = (
            [c for c in ["auc_roc", "ap", "recall", "precision", "f1", "fpr"] if c in group_rows.columns]
            + [c for c in ["mtta", "early_warning_recall", "brier"] if c in group_rows.columns]
            + noise_cols
        )
        col_labels = {**{c: c for c in per_group_show}, **tmp_labels, **noise_labels}

        self._log(f"\nPER GROUP:")
        hdr_parts = ["group                    ", "  n  "] + [f"{col_labels.get(c,c):>8}" for c in per_group_show]
        self._log("  " + "  ".join(hdr_parts))
        self._log("  " + "-" * (len("  ".join(hdr_parts))))

        for _, r in group_rows.iterrows():
            vals = []
            for col in per_group_show:
                v = _v(r, col)
                if v is None:
                    vals.append(f"{'—':>8}")
                elif col in ("n_with_alert",):
                    vals.append(f"{int(v):>8d}")
                else:
                    vals.append(f"{v:>8.3f}")
            self._log(f"  {str(r['group']):<25s}  {int(r['n_total']):>4d}  " + "  ".join(vals))

    # =========================================================================
    # STEP 3: VISUALIZATIONS
    # =========================================================================

    def generate_visualizations(self):
        """Generate plots from metrics and predictions."""
        import matplotlib
        matplotlib.use('Agg')

        self._log("\n" + "=" * 60)
        self._log("STEP 3: GENERATING VISUALIZATIONS")
        self._log("=" * 60)

        if self._metrics_df is None:
            self._metrics_df = pd.read_csv(self.metrics_path)
        if self._predictions is None:
            self.load_predictions()

        self.plots_dir.mkdir(exist_ok=True)

        if self.config.dataset_format == "kaggle":
            self._plot_kaggle_metrics()
        else:
            self._plot_metrics_by_group()
            if self.config.inference_mode == "sliding":
                self._plot_mtta_by_group()
            self._plot_sample_timeseries()

        self._log(f"\nSaved plots to: {self.plots_dir}")

    def _plot_kaggle_metrics(self):
        import matplotlib.pyplot as plt
        df = self._metrics_df[self._metrics_df["group"] != "OVERALL"].copy()
        ap_cols = [c for c in df.columns if c.startswith("ap_")]
        if not ap_cols:
            return
        x = np.arange(len(df))
        width = 0.25
        fig, ax = plt.subplots(figsize=(10, 5))
        for i, col in enumerate(ap_cols):
            ax.bar(x + i * width, df[col], width, label=col)
        ax.set_xticks(x + width); ax.set_xticklabels(df["group"])
        ax.set_ylabel("Average Precision"); ax.set_title("Kaggle AP by Split")
        ax.set_ylim(0, 1); ax.legend()
        plt.tight_layout()
        plt.savefig(self.plots_dir / "kaggle_ap_by_split.png", dpi=150, bbox_inches="tight")
        plt.close()

    def _plot_metrics_by_group(self):
        import matplotlib.pyplot as plt
        df = self._metrics_df[self._metrics_df["group"] != "OVERALL"].copy()
        plot_cols = [c for c in ["auc_roc", "ap", "f1", "recall", "fpr"] if c in df.columns]
        if not plot_cols:
            return
        df = df.sort_values(plot_cols[0], ascending=True)
        n_plots = len(plot_cols)
        fig, axes = plt.subplots(1, n_plots, figsize=(5 * n_plots, 6))
        if n_plots == 1:
            axes = [axes]
        colors = ["#3498db", "#2ecc71", "#e67e22", "#e74c3c", "#9b59b6"]
        for ax, col, color in zip(axes, plot_cols, colors):
            ax.barh(df["group"], df[col], color=color)
            ax.set_xlabel(col.upper()); ax.set_title(f"{col.upper()} by Group"); ax.set_xlim(0, 1)
        plt.tight_layout()
        plt.savefig(self.plots_dir / "metrics_by_group.png", dpi=150, bbox_inches="tight")
        plt.close()

    def _plot_mtta_by_group(self):
        import matplotlib.pyplot as plt
        df = self._metrics_df[self._metrics_df["group"] != "OVERALL"].copy()
        df = df[df["mtta"].notna() & (df["mtta"] > 0)].sort_values("mtta", ascending=True)
        if len(df) == 0:
            return
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.barh(df["group"], df["mtta"], color="#9b59b6")
        ax.errorbar(df["mtta"], df["group"], xerr=df["mtta_std"],
                    fmt='none', color='black', capsize=3)
        ax.set_xlabel("Mean Time To Alert (seconds)")
        ax.set_title("MTTA by Group (higher = earlier warning)")
        plt.tight_layout()
        plt.savefig(self.plots_dir / "mtta_by_group.png", dpi=150, bbox_inches="tight")
        plt.close()

    def _plot_sample_timeseries(self, n_samples: int = 6):
        import matplotlib.pyplot as plt
        if self._labels_df is None:
            self._labels_df = self._load_labels_df()
        positive_df = self._labels_df[self._labels_df["target"].astype(bool)]
        samples = [
            (str(row[self.config.video_id_column]), row)
            for _, row in positive_df.iterrows()
            if str(row[self.config.video_id_column]) in self._predictions
            and len(self._predictions[str(row[self.config.video_id_column])]) > 0
        ][:n_samples]
        if not samples:
            return
        n_cols = min(3, len(samples))
        n_rows = (len(samples) + n_cols - 1) // n_cols
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))
        axes = [axes] if len(samples) == 1 else axes.flatten()
        for idx, (vid, row) in enumerate(samples):
            ax = axes[idx]
            preds = self._predictions[vid]
            ax.plot([p["timestamp"] for p in preds], [p["probability"] for p in preds],
                    color="#3498db", linewidth=1.5)
            ax.axhline(y=self.config.threshold, color="#e74c3c", linestyle="--",
                       alpha=0.7, label=f"Threshold ({self.config.threshold})")
            toe = row.get(self.config.time_of_event_column)
            if toe is not None and pd.notna(toe):
                ax.axvline(x=float(toe), color="#8e44ad", linewidth=2, label="Event")
            ax.set_ylim(0, 1); ax.set_xlabel("Time (s)"); ax.set_ylabel("Probability")
            ax.set_title(f"{vid[:20]}..."); ax.legend(fontsize=8); ax.grid(True, alpha=0.3)
        for idx in range(len(samples), len(axes)):
            axes[idx].set_visible(False)
        plt.tight_layout()
        plt.savefig(self.plots_dir / "sample_timeseries.png", dpi=150, bbox_inches="tight")
        plt.close()

    # =========================================================================
    # MAIN ENTRY POINT
    # =========================================================================

    def run(self, skip_inference: bool = False) -> pd.DataFrame:
        """Run the full evaluation pipeline (inference → metrics → visualizations).

        Args:
            skip_inference: Load existing predictions.json instead of running inference.
        """
        self._log("\n" + "=" * 60)
        self._log("BADAS EVALUATOR")
        self._log("=" * 60)
        self._log(f"Output directory: {self.output_dir}")

        if skip_inference and self.predictions_path.exists():
            self._log("Skipping inference, loading existing predictions...")
            self.load_predictions()
        else:
            self.run_inference()

        self.compute_metrics()
        self.generate_visualizations()

        self._log("\n" + "=" * 60)
        self._log("EVALUATION COMPLETE")
        self._log("=" * 60)
        self._log(f"Results saved to: {self.output_dir}")

        return self._metrics_df
