"""Recompute YellowCab's public held-out metrics from sanitized predictions."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    brier_score_loss,
    f1_score,
    log_loss,
)


CLASSES = ("continue", "slow", "stop", "turn_left", "turn_right")
SEED = 20260726
BOOTSTRAP_ITERATIONS = 500


def expected_calibration_error(
    probabilities: np.ndarray,
    targets: np.ndarray,
    *,
    bins: int = 15,
) -> float:
    confidence = probabilities.max(axis=1)
    predictions = probabilities.argmax(axis=1)
    correct = predictions == targets
    edges = np.linspace(0.0, 1.0, bins + 1)
    value = 0.0
    for index in range(bins):
        if index == bins - 1:
            mask = (confidence >= edges[index]) & (confidence <= edges[index + 1])
        else:
            mask = (confidence >= edges[index]) & (confidence < edges[index + 1])
        if mask.any():
            value += float(mask.mean()) * abs(
                float(correct[mask].mean()) - float(confidence[mask].mean())
            )
    return value


def route_bootstrap_macro_f1(
    frame: pd.DataFrame,
    targets: np.ndarray,
    predictions: np.ndarray,
) -> tuple[float, float]:
    route_values = frame["route_group"].astype(str).to_numpy()
    routes = np.unique(route_values)
    rng = np.random.default_rng(SEED)
    values: list[float] = []
    for _ in range(BOOTSTRAP_ITERATIONS):
        sampled = rng.choice(routes, size=len(routes), replace=True)
        indices = np.concatenate(
            [np.flatnonzero(route_values == route) for route in sampled]
        )
        values.append(
            float(
                f1_score(
                    targets[indices],
                    predictions[indices],
                    labels=np.arange(len(CLASSES)),
                    average="macro",
                    zero_division=0,
                )
            )
        )
    return tuple(float(value) for value in np.quantile(values, [0.025, 0.975]))


def main() -> None:
    root = Path(__file__).resolve().parent
    frame = pd.read_csv(root / "eval_predictions.csv")
    class_to_index = {name: index for index, name in enumerate(CLASSES)}
    targets = frame["label"].map(class_to_index).to_numpy(dtype=np.int64)
    predictions = frame["prediction"].map(class_to_index).to_numpy(dtype=np.int64)
    probabilities = frame[
        [f"prob_{class_name}" for class_name in CLASSES]
    ].to_numpy(dtype=np.float64)
    if not np.allclose(probabilities.sum(axis=1), 1.0, atol=1e-6):
        raise ValueError("Probability rows do not sum to one")
    one_hot = np.eye(len(CLASSES), dtype=np.float64)[targets]
    ci_low, ci_high = route_bootstrap_macro_f1(frame, targets, predictions)
    result = {
        "samples": int(len(frame)),
        "routes": int(frame["route_group"].nunique()),
        "accuracy": float(accuracy_score(targets, predictions)),
        "balanced_accuracy": float(balanced_accuracy_score(targets, predictions)),
        "macro_f1": float(
            f1_score(
                targets,
                predictions,
                labels=np.arange(len(CLASSES)),
                average="macro",
                zero_division=0,
            )
        ),
        "weighted_f1": float(
            f1_score(
                targets,
                predictions,
                labels=np.arange(len(CLASSES)),
                average="weighted",
                zero_division=0,
            )
        ),
        "log_loss": float(log_loss(targets, probabilities, labels=range(len(CLASSES)))),
        "brier_score": float(np.mean(np.sum((probabilities - one_hot) ** 2, axis=1))),
        "ece_15_bin": expected_calibration_error(probabilities, targets),
        "macro_f1_route_ci_95": [ci_low, ci_high],
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()