"""
Adaptation Evaluation — measures per-camera LoRA quality vs static routing.

Metrics:
  - Per-camera decode loss (adapted vs base)
  - Adaptation improvement percentage
  - Confidence calibration quality
  - Adapter diversity (how different are adapters for different cameras?)

Used by scripts/eval_adaptation.py — this module provides the core evaluation
functions without the CLI/checkpoint-loading boilerplate.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn.functional as F


def compute_adaptation_improvement(
    base_loss: float,
    adapted_loss: float,
) -> float:
    """Compute improvement percentage: positive means adapted is better."""
    if base_loss == 0:
        return 0.0
    return (base_loss - adapted_loss) / base_loss * 100


def compute_adapter_diversity(
    adapter_params: list[torch.Tensor],
) -> dict[str, float]:
    """
    Measure how diverse the generated adapters are across cameras.

    Low diversity suggests the hypernetwork is ignoring conditioning.

    Args:
        adapter_params: List of flat LoRA param tensors (one per camera)
    Returns:
        Dict with mean_cosine_sim (lower = more diverse), std, min, max
    """
    if len(adapter_params) < 2:
        return {"mean_cosine_sim": 0.0, "std": 0.0, "min": 0.0, "max": 0.0}

    # Stack and compute pairwise cosine similarity
    stacked = torch.stack(adapter_params)  # [N, params]
    stacked_norm = F.normalize(stacked, dim=-1)
    sim_matrix = stacked_norm @ stacked_norm.T  # [N, N]

    # Get upper triangle (exclude diagonal)
    mask = torch.triu(torch.ones_like(sim_matrix, dtype=torch.bool), diagonal=1)
    pairwise_sims = sim_matrix[mask]

    return {
        "mean_cosine_sim": pairwise_sims.mean().item(),
        "std": pairwise_sims.std().item(),
        "min": pairwise_sims.min().item(),
        "max": pairwise_sims.max().item(),
    }


def compute_calibration_quality(
    confidences: list[float],
    errors: list[float],
    n_bins: int = 10,
) -> dict[str, float]:
    """
    Measure calibration quality: does confidence predict actual accuracy?

    Uses Expected Calibration Error (ECE) — lower is better.

    Args:
        confidences: List of confidence scores [0, 1]
        errors: List of actual error values (lower = better prediction)
    Returns:
        Dict with ece, mean_confidence, mean_error
    """
    if not confidences:
        return {"ece": 0.0, "mean_confidence": 0.0, "mean_error": 0.0}

    conf_t = torch.tensor(confidences)
    err_t = torch.tensor(errors)

    # Convert errors to accuracy (1 - normalized_error)
    max_err = err_t.max()
    if max_err > 0:
        accuracy = 1.0 - err_t / max_err
    else:
        accuracy = torch.ones_like(err_t)

    # Bin by confidence
    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        mask = (conf_t >= bin_boundaries[i]) & (conf_t < bin_boundaries[i + 1])
        if mask.sum() > 0:
            bin_conf = conf_t[mask].mean()
            bin_acc = accuracy[mask].mean()
            ece += mask.float().mean() * abs(bin_conf - bin_acc)

    return {
        "ece": ece.item(),
        "mean_confidence": conf_t.mean().item(),
        "mean_error": err_t.mean().item(),
        "mean_accuracy": accuracy.mean().item(),
    }