"""
Regression tests for the unified evaluation code.

These tests verify that:
1. Kaggle metrics (AP @ 0.5s/1.0s/1.5s, accuracy) produced by the unified
   Evaluator match what the old benchmark_kaggle.compute_competition_metrics()
   produced for the same inputs.
2. Testset metrics (F1, recall, FPR, MTTA, Brier) are unchanged after the
   Kaggle-support additions.

No model or real videos are needed — tests operate entirely on synthetic
prediction dicts and label DataFrames, exercising only the metrics layer.
"""
import math
import tempfile
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from badas.validation.evaluator import Evaluator, EvaluatorConfig


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_preds_dict(video_ids, probs):
    """Build a predictions dict {vid: [{timestamp, probability, raw_probability}]}."""
    return {
        str(vid): [{"timestamp": 1.0, "probability": float(p), "raw_probability": float(p)}]
        for vid, p in zip(video_ids, probs)
    }


def _make_evaluator(labels_df, predictions, dataset_format="testset", **kwargs):
    """
    Create an Evaluator with pre-loaded labels + predictions (no CSV / video files needed).
    Uses a temp dir for output so tests don't litter the workspace.
    """
    tmpdir = tempfile.mkdtemp()
    # Write a minimal CSV so EvaluatorConfig has a path (even though we override _labels_df).
    csv_path = str(Path(tmpdir) / "labels.csv")
    labels_df.to_csv(csv_path, index=False)

    cfg = EvaluatorConfig(
        checkpoint_path="dummy",
        labels_csv=csv_path,
        output_dir=tmpdir,
        dataset_format=dataset_format,
        quiet=True,
        **kwargs,
    )
    ev = Evaluator(cfg)
    # Inject pre-computed data so no inference or CSV normalisation is needed.
    ev._labels_df = labels_df.copy()
    ev._predictions = predictions
    return ev


# ---------------------------------------------------------------------------
# Reference implementation (copied verbatim from scripts/benchmark_kaggle.py)
# to serve as the ground-truth oracle for regression.
# ---------------------------------------------------------------------------

def _reference_kaggle_metrics(solution_df: pd.DataFrame, predictions: dict) -> dict:
    """
    Identical to compute_competition_metrics() in scripts/benchmark_kaggle.py.
    Returns nested dict: {split_lower: {accuracy, ap_0.5s, ap_1.0s, ap_1.5s, mean_ap, n_samples}}
    """
    from sklearn.metrics import average_precision_score

    group_to_tta = {0: 0.5, 1: 1.0, 2: 1.5}
    results = {}

    for split in ["Public", "Private"]:
        split_df = solution_df[solution_df["Usage"] == split].copy()
        split_df["prob"] = split_df["id"].map(predictions)
        split_df = split_df.dropna(subset=["prob"])

        split_results = {}
        for group, tta in group_to_tta.items():
            group_df = split_df[split_df["group"] == group]
            if len(group_df) > 0:
                ap = average_precision_score(group_df["target"], group_df["prob"])
                split_results[tta] = ap

        acc = ((split_df["prob"] > 0.5).astype(int) == split_df["target"]).mean()
        mean_ap = np.mean(list(split_results.values())) if split_results else 0.0

        results[split.lower()] = {
            "accuracy":  acc,
            "ap_0.5s":   split_results.get(0.5, 0),
            "ap_1.0s":   split_results.get(1.0, 0),
            "ap_1.5s":   split_results.get(1.5, 0),
            "mean_ap":   mean_ap,
            "n_samples": len(split_df),
        }

    return results


# ---------------------------------------------------------------------------
# Kaggle regression tests
# ---------------------------------------------------------------------------

def _make_kaggle_fixture(seed=0):
    """
    Create a synthetic Kaggle solution DataFrame + predictions dict.
    Mirrors the Kaggle CSV format: id (5-digit str), target, group (0/1/2), Usage.
    """
    rng = np.random.default_rng(seed)
    n = 120  # 40 per TTA group, 60 public / 60 private

    ids = [f"{i+1:05d}" for i in range(n)]
    targets = rng.integers(0, 2, n).tolist()
    groups = [i % 3 for i in range(n)]
    usages = ["Public" if i < n // 2 else "Private" for i in range(n)]

    solution_df = pd.DataFrame({
        "id": ids,
        "target": targets,
        "group": groups,
        "Usage": usages,
    })

    # Predictions: slightly noisy scores correlated with targets.
    raw_probs = np.clip(rng.normal(0.5, 0.3, n) + 0.2 * np.array(targets), 0.01, 0.99)
    predictions = {vid: float(p) for vid, p in zip(ids, raw_probs)}

    return solution_df, predictions


def test_kaggle_ap_matches_reference():
    """AP scores from unified Evaluator must exactly match the reference implementation."""
    solution_df, flat_predictions = _make_kaggle_fixture(seed=42)

    # Reference (old benchmark_kaggle.py logic)
    ref = _reference_kaggle_metrics(solution_df, flat_predictions)

    # New unified Evaluator — predictions need the list-of-dicts format
    preds_dict = _make_preds_dict(flat_predictions.keys(), flat_predictions.values())
    ev = _make_evaluator(solution_df, preds_dict, dataset_format="kaggle")
    metrics_df = ev.compute_metrics()

    for split_name, ref_metrics in ref.items():
        row = metrics_df[metrics_df["group"] == split_name.capitalize()]
        assert len(row) == 1, f"Missing split row: {split_name}"
        row = row.iloc[0]

        # The evaluator stores metrics rounded to 4 decimal places; compare at that precision.
        tol = 5e-5
        assert abs(row["mean_ap"] - ref_metrics["mean_ap"]) < tol, (
            f"{split_name} mean_ap: got {row['mean_ap']}, expected {ref_metrics['mean_ap']}"
        )
        for tta_col, ref_val in [
            ("ap_0.5s", ref_metrics["ap_0.5s"]),
            ("ap_1.0s", ref_metrics["ap_1.0s"]),
            ("ap_1.5s", ref_metrics["ap_1.5s"]),
        ]:
            if tta_col in row.index:
                assert abs(row[tta_col] - ref_val) < tol, (
                    f"{split_name} {tta_col}: got {row[tta_col]}, expected {ref_val}"
                )

        assert abs(row["accuracy"] - ref_metrics["accuracy"]) < tol, (
            f"{split_name} accuracy mismatch"
        )


def test_kaggle_id_normalisation():
    """Integer IDs in the Kaggle CSV must be normalised to 5-digit strings."""
    solution_df, flat_predictions = _make_kaggle_fixture(seed=7)

    # Use integer IDs in the DataFrame (as they appear in the raw CSV).
    solution_df_int = solution_df.copy()
    solution_df_int["id"] = solution_df_int["id"].astype(int)

    tmpdir = tempfile.mkdtemp()
    csv_path = str(Path(tmpdir) / "solution.csv")
    solution_df_int.to_csv(csv_path, index=False)

    preds_dict = _make_preds_dict(flat_predictions.keys(), flat_predictions.values())

    cfg = EvaluatorConfig(
        checkpoint_path="dummy",
        labels_csv=csv_path,
        output_dir=tmpdir,
        dataset_format="kaggle",
        quiet=True,
    )
    ev = Evaluator(cfg)
    # _load_labels_df() should zero-pad IDs; inject predictions directly.
    ev._labels_df = ev._load_labels_df()
    ev._predictions = preds_dict

    metrics_df = ev.compute_metrics()
    # If normalisation is correct, all videos are matched and n_total is correct.
    overall = metrics_df[metrics_df["group"] == "OVERALL"].iloc[0]
    assert overall["n_total"] == len(solution_df), (
        f"ID normalisation failed: expected {len(solution_df)} videos, got {overall['n_total']}"
    )


def test_kaggle_submission_csv(tmp_path):
    """Submission CSV must be written with correct id/target columns."""
    solution_df, flat_predictions = _make_kaggle_fixture(seed=3)
    preds_dict = _make_preds_dict(flat_predictions.keys(), flat_predictions.values())
    submission_path = str(tmp_path / "submission.csv")

    ev = _make_evaluator(
        solution_df, preds_dict,
        dataset_format="kaggle",
        submission_path=submission_path,
    )
    ev.compute_metrics()

    assert Path(submission_path).exists(), "Submission CSV was not created"
    sub_df = pd.read_csv(submission_path)
    assert set(sub_df.columns) >= {"id", "target"}, "Submission CSV missing required columns"
    assert len(sub_df) == len(flat_predictions), "Submission CSV row count mismatch"


def test_kaggle_multi_window_uses_max():
    """When multiple windows exist per video, Kaggle score must be max(probability)."""
    # Two videos: one positive, one negative.
    solution_df = pd.DataFrame({
        "id": ["00001", "00002"],
        "target": [1, 0],
        "group": [0, 0],
        "Usage": ["Public", "Public"],
    })

    # Video 00001 has a high peak then falls; 00002 has a low peak.
    predictions = {
        "00001": [
            {"timestamp": 0.5, "probability": 0.9, "raw_probability": 0.9},
            {"timestamp": 1.0, "probability": 0.2, "raw_probability": 0.2},
        ],
        "00002": [
            {"timestamp": 0.5, "probability": 0.3, "raw_probability": 0.3},
            {"timestamp": 1.0, "probability": 0.1, "raw_probability": 0.1},
        ],
    }

    ev = _make_evaluator(solution_df, predictions, dataset_format="kaggle")
    metrics_df = ev.compute_metrics()

    # With max-score aggregation, 00001 score=0.9 and 00002 score=0.3,
    # both correctly classified at threshold 0.5.
    overall = metrics_df[metrics_df["group"] == "OVERALL"].iloc[0]
    assert overall["accuracy"] == 1.0, "Both videos should be correctly classified using max score"


# ---------------------------------------------------------------------------
# Testset regression tests
# ---------------------------------------------------------------------------

def _make_testset_fixture(seed=0):
    """Create synthetic testset labels + predictions."""
    rng = np.random.default_rng(seed)
    n = 40
    ids = [f"vid{i:04d}" for i in range(n)]
    targets = (rng.random(n) > 0.5).astype(int).tolist()
    groups = [["animal", "pedestrian"][i % 2] for i in range(n)]
    events = [rng.uniform(2.0, 8.0) if t else float("nan") for t in targets]

    labels_df = pd.DataFrame({
        "id": ids,
        "target": targets,
        "test_group": groups,
        "time_of_event": events,
    })

    # Build sliding-window style predictions (multiple timestamps per video).
    predictions = {}
    for vid, t, toe in zip(ids, targets, events):
        n_windows = 8
        timestamps = [float(j) * 0.5 for j in range(n_windows)]
        base = 0.8 if t else 0.1
        probs = [float(np.clip(base + rng.normal(0, 0.1), 0.0, 1.0)) for _ in timestamps]
        predictions[vid] = [
            {"timestamp": ts, "probability": p, "raw_probability": p}
            for ts, p in zip(timestamps, probs)
        ]

    return labels_df, predictions


def test_testset_metrics_stable():
    """
    Testset metrics must be deterministic and self-consistent given fixed inputs.
    Run twice and verify identical outputs.
    """
    labels_df, predictions = _make_testset_fixture(seed=1)

    ev1 = _make_evaluator(labels_df, predictions, dataset_format="testset")
    df1 = ev1.compute_metrics()

    ev2 = _make_evaluator(labels_df, predictions, dataset_format="testset")
    df2 = ev2.compute_metrics()

    pd.testing.assert_frame_equal(df1, df2, check_exact=True)


def test_testset_metrics_sanity():
    """Perfect classifier should yield precision=recall=f1=1, fpr=0."""
    n = 20
    ids = [f"v{i}" for i in range(n)]
    targets = [i % 2 for i in range(n)]
    labels_df = pd.DataFrame({
        "id": ids,
        "target": targets,
        "test_group": ["all"] * n,
        "time_of_event": [5.0 if t else float("nan") for t in targets],
    })

    # Perfect scores: positive → 1.0, negative → 0.0
    threshold = 0.75
    predictions = {
        vid: [{"timestamp": 1.0, "probability": float(t), "raw_probability": float(t)}]
        for vid, t in zip(ids, targets)
    }

    ev = _make_evaluator(labels_df, predictions, dataset_format="testset", threshold=threshold)
    df = ev.compute_metrics()
    overall = df[df["group"] == "OVERALL"].iloc[0]

    assert overall["precision"] == 1.0
    assert overall["recall"] == 1.0
    assert overall["f1"] == 1.0
    assert overall["fpr"] == 0.0


def test_testset_does_not_compute_kaggle_columns():
    """Testset mode must not produce ap_* columns."""
    labels_df, predictions = _make_testset_fixture(seed=5)
    ev = _make_evaluator(labels_df, predictions, dataset_format="testset")
    df = ev.compute_metrics()
    kaggle_cols = [c for c in df.columns if c.startswith("ap_")]
    assert kaggle_cols == [], f"Unexpected Kaggle columns in testset mode: {kaggle_cols}"


def test_kaggle_does_not_compute_testset_columns():
    """Kaggle mode must not produce mtta / brier columns."""
    solution_df, flat_preds = _make_kaggle_fixture(seed=9)
    preds_dict = _make_preds_dict(flat_preds.keys(), flat_preds.values())
    ev = _make_evaluator(solution_df, preds_dict, dataset_format="kaggle")
    df = ev.compute_metrics()
    testset_cols = [c for c in df.columns if c in ("mtta", "brier", "early_warning_recall")]
    assert testset_cols == [], f"Unexpected testset columns in kaggle mode: {testset_cols}"
