"""Per-field temperature scaling on stored logits (pure Python). Run as ``python -m training.calibrate``.

Prediction rows (``training.predict``, the pod, and the node12 server) hold raw logits per field in taxonomy option
order: ``{"id", "split", "n_tokens", "truncated", "logits": {field: [float, ...]}}``. Calibrated probabilities are
softmax(logits / T_field); the served/stored probabilities are those, gated (``taxonomy.gate``). One temperature
per field is fitted on val by NLL (grid search) over the rows whose target for that field is not masked, and frozen
before test (classifier-plan §5.3). ``crossfit`` fits on each of two author halves of val and scores the other half.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
from typing import Iterable, Mapping

from .taxonomy import Taxonomy, argmax, gate

T_GRID = tuple(round(0.1 * (100 ** (i / 160)), 4) for i in range(161))  # 0.1 .. 10, geometric
FLOOR = 1e-12


def softmax(logits: Iterable[float], t: float = 1.0) -> list[float]:
    scaled = [x / t for x in logits]
    m = max(scaled)
    exp = [math.exp(x - m) for x in scaled]
    total = sum(exp)
    return [e / total for e in exp]


def probs(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None
          ) -> dict[str, dict[str, float]]:
    """Calibrated (ungated) distributions from per-field logits."""
    temps = temps or {}
    return {f: dict(zip(tax.fields[f], softmax(logits[f], temps.get(f, 1.0)))) for f in tax.names}


def gated(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None
          ) -> dict[str, dict[str, float]]:
    """Calibrated and gated distributions: what the server returns and ``post_predictions.probs`` stores."""
    return gate(probs(logits, tax, temps))


def outputs(logits: Mapping[str, list[float]], tax: Taxonomy, temps: Mapping[str, float] | None = None) -> dict:
    """Everything a prediction carries besides the logits: gated probabilities and gated argmax labels."""
    p = gated(logits, tax, temps)
    return {"probs": p, "labels": argmax(p)}


def max_prob_diff(a: Mapping[str, Mapping[str, float]], b: Mapping[str, Mapping[str, float]]) -> float:
    """Largest absolute difference between two {field: {option: p}} sets (the same fields and options)."""
    if set(a) != set(b):
        raise ValueError("field sets differ")
    out = 0.0
    for f in a:
        if set(a[f]) != set(b[f]):
            raise ValueError(f"option sets of {f} differ")
        out = max(out, max(abs(a[f][o] - b[f][o]) for o in a[f]))
    return out


def _gold(targets: Mapping[str, str | None], mask: Iterable[str], field: str) -> str | None:
    return None if field in set(mask) else targets.get(field)


def nll(pairs, tax: Taxonomy, field: str, t: float) -> float | None:
    """Mean NLL of the target option; pairs are (prediction row, export row with targets and mask)."""
    total, n = 0.0, 0
    options = tax.fields[field]
    for pred, row in pairs:
        gold = _gold(row["targets"], row.get("mask", ()), field)
        if gold is None:
            continue
        p = softmax(pred["logits"][field], t)[options.index(gold)]
        total -= math.log(max(p, FLOOR))
        n += 1
    return total / n if n else None


def fit(pairs, tax: Taxonomy) -> dict[str, float]:
    pairs = list(pairs)
    out = {}
    for f in tax.names:
        if nll(pairs, tax, f, 1.0) is None:
            out[f] = 1.0
            continue
        out[f] = min(T_GRID, key=lambda t: nll(pairs, tax, f, t))
    return out


def author_half(author_key: str, seed: str = "nsb-crossfit") -> int:
    return int(hashlib.sha256(json.dumps([seed, author_key]).encode()).hexdigest()[:8], 16) % 2


def crossfit(pairs, tax: Taxonomy) -> dict:
    """Temperatures fitted on one author half of val and applied to the other.

    Returns both temperature sets and, per row id, the out-of-half calibrated (ungated) distributions."""
    halves: dict[int, list] = {0: [], 1: []}
    for pred, row in pairs:
        halves[author_half(row["author_key"])].append((pred, row))
    temps = {h: fit(halves[h], tax) for h in (0, 1)}
    calibrated = {}
    for h in (0, 1):
        for pred, _ in halves[h]:
            calibrated[pred["id"]] = probs(pred["logits"], tax, temps[1 - h])
    return {"temperatures_by_half": {str(h): temps[h] for h in (0, 1)}, "calibrated": calibrated,
            "sizes": {str(h): len(halves[h]) for h in (0, 1)}}


def pair_up(preds: list[dict], rows: list[dict]) -> list[tuple[dict, dict]]:
    """(prediction, export row) pairs; every row needs exactly one prediction."""
    by_id: dict[str, dict] = {}
    for p in preds:
        if p["id"] in by_id:
            raise ValueError(f"duplicate prediction {p['id']}")
        by_id[p["id"]] = p
    missing = [r["id"] for r in rows if r["id"] not in by_id]
    if missing:
        raise ValueError(f"{len(missing)} rows have no prediction (first: {missing[0]})")
    return [(by_id[r["id"]], r) for r in rows]


def main(argv=None) -> int:
    from .data import load_jsonl, load_split

    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = parser.add_subparsers(dest="command", required=True)
    f = sub.add_parser("fit", help="Fit temperatures on val predictions")
    f.add_argument("predictions")
    f.add_argument("--data", required=True, help="Bundle directory (taxonomy.json, val.jsonl)")
    f.add_argument("--out", required=True, help="temperatures.json")
    f.add_argument("--split", default="val")
    args = parser.parse_args(argv)
    if Path(args.out).exists():
        raise SystemExit("refusing to overwrite " + args.out)
    tax = Taxonomy.load(args.data)
    rows = load_split(args.data, args.split)
    pairs = pair_up(load_jsonl(args.predictions), rows)
    temps = fit(pairs, tax)
    edge = sorted(k for k, t in temps.items() if t in (T_GRID[0], T_GRID[-1]))
    Path(args.out).write_text(json.dumps(temps, indent=1) + "\n")
    print(json.dumps({"temperatures": temps, "at_grid_edge": edge}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())