from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from metrics_lie.metrics.core import (
    METRICS,
    THRESHOLD_METRICS,
    CALIBRATION_METRICS,
    RANKING_METRICS,
)
from metrics_lie.model.surface import PredictionSurface, SurfaceType
from metrics_lie.surface_compat import DEFAULT_THRESHOLD

# Score-based metrics that don't use a threshold (union of ranking + calibration).
SCORE_METRIC_IDS: set[str] = RANKING_METRICS | CALIBRATION_METRICS


@dataclass(frozen=True)
class ThresholdSweepResult:
    thresholds: np.ndarray
    metric_curves: dict[str, np.ndarray]
    optimal_thresholds: dict[str, float]
    crossover_points: list[dict]

    def to_jsonable(self) -> dict:
        return {
            "thresholds": self.thresholds.tolist(),
            "metric_curves": {k: v.tolist() for k, v in self.metric_curves.items()},
            "optimal_thresholds": self.optimal_thresholds,
            "crossover_points": self.crossover_points,
        }


def run_threshold_sweep(
    *,
    y_true: np.ndarray,
    surface: PredictionSurface,
    metrics: list[str],
    n_points: int = 101,
) -> ThresholdSweepResult:
    if surface.surface_type not in (SurfaceType.PROBABILITY, SurfaceType.SCORE):
        raise ValueError(
            f"threshold sweep requires PROBABILITY or SCORE surface, "
            f"got {surface.surface_type.value!r}"
        )

    if surface.surface_type == SurfaceType.PROBABILITY:
        low, high = 0.0, 1.0
    else:
        finite_vals = np.asarray(surface.values, dtype=float)
        finite_vals = finite_vals[np.isfinite(finite_vals)]
        if finite_vals.size == 0:
            raise ValueError(
                "threshold sweep requires score variance; surface values are constant"
            )
        vmin = float(np.min(finite_vals))
        vmax = float(np.max(finite_vals))
        if vmin == vmax:
            raise ValueError(
                "threshold sweep requires score variance; surface values are constant"
            )
        low, high = vmin, vmax
    thresholds = np.linspace(low, high, n_points)
    curves: dict[str, np.ndarray] = {}
    optimal_thresholds: dict[str, float] = {}

    for metric_id in metrics:
        fn = METRICS.get(metric_id)
        if fn is None:
            continue
        if metric_id in THRESHOLD_METRICS:
            vals = [fn(y_true, surface.values, threshold=t) for t in thresholds]
            curves[metric_id] = np.array(vals, dtype=float)
            best_idx = int(np.argmax(curves[metric_id]))
            optimal_thresholds[metric_id] = float(thresholds[best_idx])
        elif metric_id in SCORE_METRIC_IDS and surface.surface_type == SurfaceType.PROBABILITY:
            v = fn(y_true, surface.values)
            curves[metric_id] = np.full_like(thresholds, float(v), dtype=float)
            optimal_thresholds[metric_id] = float(surface.threshold or DEFAULT_THRESHOLD)

    crossover_points: list[dict] = []
    metric_ids = [m for m in metrics if m in THRESHOLD_METRICS]
    for i in range(len(metric_ids)):
        for j in range(i + 1, len(metric_ids)):
            a = metric_ids[i]
            b = metric_ids[j]
            ta = optimal_thresholds.get(a)
            tb = optimal_thresholds.get(b)
            if ta is None or tb is None:
                continue
            if abs(ta - tb) > 1e-9:
                crossover_points.append(
                    {"metric_a": a, "metric_b": b, "threshold_a": ta, "threshold_b": tb}
                )

    return ThresholdSweepResult(
        thresholds=thresholds,
        metric_curves=curves,
        optimal_thresholds=optimal_thresholds,
        crossover_points=crossover_points,
    )
