from __future__ import annotations

from typing import Any

import numpy as np
import torch

from tiny_router.constants import HEAD_LABELS


def scale_logits(logits: np.ndarray, temperature: float) -> np.ndarray:
    safe_temperature = max(float(temperature), 1e-3)
    return logits / safe_temperature


def apply_temperature_scaling(
    logits_by_head: dict[str, np.ndarray],
    temperatures: dict[str, float] | None,
) -> dict[str, np.ndarray]:
    if not temperatures:
        return logits_by_head
    return {
        head: scale_logits(logits, temperatures.get(head, 1.0))
        for head, logits in logits_by_head.items()
    }


def fit_temperature(logits: np.ndarray, labels: np.ndarray) -> float:
    logits_tensor = torch.tensor(logits, dtype=torch.float32)
    labels_tensor = torch.tensor(labels, dtype=torch.long)
    loss_fn = torch.nn.CrossEntropyLoss()
    # Use a bounded search instead of LBFGS to avoid numerical blowups on sharp logits.
    candidates = torch.logspace(-1, 1, steps=81, dtype=torch.float32)
    losses = []
    with torch.no_grad():
        for temperature in candidates:
            loss = loss_fn(logits_tensor / temperature, labels_tensor)
            losses.append(float(loss.item()))
    best_index = min(range(len(losses)), key=losses.__getitem__)
    return round(float(candidates[best_index].item()), 6)


def fit_temperature_by_head(
    logits_by_head: dict[str, np.ndarray],
    labels_by_head: dict[str, np.ndarray],
) -> dict[str, Any]:
    per_head = {}
    for head in HEAD_LABELS:
        per_head[head] = fit_temperature(logits_by_head[head], labels_by_head[head])
    return {
        "method": "per_head_temperature_scaling",
        "source_split": "validation",
        "per_head": per_head,
    }