"""Soft-reference metrics, temperature fitting and bounded invariance evaluation."""

import copy
import math
from collections import defaultdict

import torch

from .data import encoded_row
from .model import collate, soft_ce_sum


def probabilities(record, temperatures=None):
    temperature = (temperatures or {}).get(record["type"], 1.0)
    return (torch.tensor(record["logits"], dtype=torch.float64) / temperature).softmax(-1)


def metric_summary(records, temperatures=None):
    if not records:
        return {"count": 0}
    loss, kl, brier, agreement, ordinal = [], [], [], [], []
    hard_bins = [[] for _ in range(15)]
    soft_bins = [[] for _ in range(15)]
    for row in records:
        p = probabilities(row, temperatures)
        t = torch.tensor(row["target"], dtype=torch.float64)
        logp = p.clamp_min(1e-15).log()
        loss.append(float(-(t * logp).sum()))
        entropy = float(-(t * t.clamp_min(1e-15).log()).sum())
        kl.append(max(0.0, loss[-1] - entropy))
        brier.append(float(0.5 * (p - t).square().sum()))
        predicted, target = int(p.argmax()), int(t.argmax())
        agreement.append(float(predicted == target))
        confidence = float(p.max())
        bucket = min(14, int(confidence * 15))
        soft_bins[bucket].append((confidence, float(t[predicted])))
        if float(t.max()) >= 1 - 1e-8:
            hard_bins[bucket].append((confidence, float(predicted == target)))
        if row["type"] == "score":
            levels = torch.arange(len(p), dtype=torch.float64)
            ordinal.append(abs(float(((p - t) * levels).sum())))

    def reliability(bins):
        result = []
        for values in bins:
            if values:
                result.append(
                    {
                        "count": len(values),
                        "confidence": sum(v[0] for v in values) / len(values),
                        "reference_agreement": sum(v[1] for v in values) / len(values),
                    }
                )
        return result

    hard = reliability(hard_bins)
    total_hard = sum(b["count"] for b in hard)
    result = {
        "count": len(records),
        "soft_ce": sum(loss) / len(loss),
        "kl": sum(kl) / len(kl),
        "brier": sum(brier) / len(brier),
        "argmax_agreement": sum(agreement) / len(agreement),
        "hard_count": total_hard,
        "hard_ece": (
            sum(b["count"] * abs(b["confidence"] - b["reference_agreement"]) for b in hard)
            / total_hard
        )
        if total_hard
        else None,
        "hard_reliability": hard,
        "soft_reliability": reliability(soft_bins),
    }
    if ordinal:
        result["expected_score_mae"] = sum(ordinal) / len(ordinal)
    return result


def grouped_metrics(records, temperatures=None):
    result = {"overall": metric_summary(records, temperatures)}
    for field in ("type", "category"):
        groups = defaultdict(list)
        for row in records:
            groups[row[field]].append(row)
        result[field] = {name: metric_summary(rows, temperatures) for name, rows in groups.items()}
    return result


@torch.no_grad()
def evaluate(model, codec, rows, batch_size=8, budget=None):
    device = next(model.parameters()).device
    was_training = model.training
    model.eval()
    records = []
    completed = True
    try:
        for start in range(0, len(rows), batch_size):
            if budget and budget.stop():
                completed = False
                break
            chunk = rows[start : start + batch_size]
            encoded = [encoded_row(row, codec) for row in chunk]
            batch = collate(encoded, codec.pad, device)
            with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
                logits = model(batch).float().cpu()
            index = 0
            for row, encoding in zip(chunk, encoded, strict=True):
                for name, keys, target, typ in zip(
                    encoding["question_names"],
                    encoding["option_keys"],
                    encoding["targets"],
                    encoding["types"],
                    strict=True,
                ):
                    values = logits[index, : len(keys)].tolist()
                    if not all(math.isfinite(v) for v in values):
                        raise ValueError("Non-finite evaluation logits")
                    records.append(
                        {
                            "id": row["decision_ids"][name],
                            "question": name,
                            "type": typ,
                            "category": row["categories"][name],
                            "option_keys": keys,
                            "target": target,
                            "logits": values,
                        }
                    )
                    index += 1
            if budget:
                budget.persist()
    finally:
        model.train(was_training)
    return records, completed


def fit_temperatures(records):
    result = {}
    for typ in ("choice", "noul", "score"):
        selected = [r for r in records if r["type"] == typ]
        if not selected:
            result[typ] = 1.0
            continue
        width = max(len(r["target"]) for r in selected)
        logits = torch.full((len(selected), width), -torch.inf, dtype=torch.float64)
        targets = torch.zeros_like(logits)
        for i, row in enumerate(selected):
            logits[i, : len(row["logits"])] = torch.tensor(row["logits"])
            targets[i, : len(row["target"])] = torch.tensor(row["target"])
        # Finite masked values avoid inf * d(1/T)/dT in autograd.
        logits = logits.masked_fill(~torch.isfinite(logits), -1e9)
        log_t = torch.zeros((), dtype=torch.float64, requires_grad=True)
        optimizer = torch.optim.LBFGS([log_t], max_iter=60, line_search_fn="strong_wolfe")

        def closure():
            optimizer.zero_grad()
            loss = soft_ce_sum(logits / log_t.clamp(-3, 3).exp(), targets) / len(selected)
            loss.backward()
            return loss

        initial = float(closure().detach())
        optimizer.step(closure)
        temperature = float(log_t.detach().clamp(-3, 3).exp())
        fitted = float(soft_ce_sum(logits / temperature, targets) / len(selected))
        result[typ] = temperature if math.isfinite(fitted) and fitted <= initial + 1e-9 else 1.0
    return result


def permutation_rows(rows, count, seed):
    # Round-robin categories for a bounded, deterministic stratified selection.
    from .common import digest

    groups = defaultdict(list)
    for row in rows:
        for name, question in row["input"]["questions"].items():
            if question["type"] == "choice":
                groups[row["categories"][name]].append(
                    (digest([seed, row["decision_ids"][name]]), row, name)
                )
    queues = [sorted(group, key=lambda value: value[0]) for _, group in sorted(groups.items())]
    selected = []
    while len(selected) < count and any(queues):
        for group in queues:
            if group and len(selected) < count:
                _, row, name = group.pop(0)
                from .data import subset

                selected.append(subset(row, [name]))
    return selected


def permutation_sensitivity(model, codec, rows, count=256, seed=3407, budget=None):
    selected = permutation_rows(rows, count, seed)
    base, changed = [], []
    for row in selected:
        row["encoded"] = codec.encode(
            row["input"]["state"], row["input"]["questions"], row["targets"], row["option_keys"]
        )
        variant = copy.deepcopy(row)
        name = next(iter(variant["input"]["questions"]))
        criteria = variant["input"]["questions"][name]["criteria"]
        items = list(criteria.items())
        variant["input"]["questions"][name]["criteria"] = dict(items[1:] + items[:1])
        variant["encoded"] = codec.encode(
            variant["input"]["state"],
            variant["input"]["questions"],
            variant["targets"],
            variant["option_keys"],
        )
        base.append(row)
        changed.append(variant)
    a, complete_a = evaluate(model, codec, base, budget=budget)
    b, complete_b = evaluate(model, codec, changed, budget=budget)
    total_variation, flips = [], []
    for original, variant in zip(a, b):
        p = dict(zip(original["option_keys"], probabilities(original).tolist(), strict=True))
        q = dict(zip(variant["option_keys"], probabilities(variant).tolist(), strict=True))
        total_variation.append(0.5 * sum(abs(p[k] - q[k]) for k in p))
        flips.append(max(p, key=p.get) != max(q, key=q.get))
    return {
        "count": len(flips),
        "completed": complete_a and complete_b,
        "mean_total_variation": sum(total_variation) / len(flips) if flips else None,
        "answer_flip_rate": sum(flips) / len(flips) if flips else None,
    }