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

Prediction rows (written by ``predict.py``) hold raw logits per field in taxonomy option order:
``{"id", "split", "n_tokens", "truncated", "logits": {field: [float, ...]}}``. Calibrated probabilities are
softmax(logits / T_field). Temperatures are fitted by NLL grid search on dev rows with a non-null target;
``crossfit`` fits them on each of two author-group halves of dev and scores the other half, which gives an
out-of-sample calibration check without touching the benchmark.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path

from .taxonomy import Taxonomy, 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, 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(row: dict, tax: Taxonomy, temps: dict | None = None) -> dict[str, dict[str, float]]:
    temps = temps or {}
    return {f: dict(zip(tax.fields[f], softmax(row['logits'][f], temps.get(f, 1.0)))) for f in tax.names}


def gated(row: dict, tax: Taxonomy, temps: dict | None = None) -> dict[str, dict[str, float]]:
    return gate(probs(row, tax, temps))


def nll(pairs, tax: Taxonomy, field: str, t: float) -> float | None:
    """Mean negative log-likelihood of the gold option; pairs are (prediction row, targets dict)."""
    total, n = 0.0, 0
    options = tax.fields[field]
    for row, targets in pairs:
        gold = targets.get(field)
        if gold is None:
            continue
        p = softmax(row['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]:
    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_id: str, seed: str = 'classifier-crossfit') -> int:
    return int(hashlib.sha256(json.dumps([seed, author_id]).encode()).hexdigest()[:8], 16) % 2


def crossfit(pairs_with_author, tax: Taxonomy) -> dict:
    """Temperatures fitted on one author half and applied to the other: (row, targets, author) triples.

    Returns the two fitted temperature sets and the out-of-half calibrated rows (id -> calibrated probs)."""
    halves = {0: [], 1: []}
    for row, targets, author in pairs_with_author:
        halves[author_half(author)].append((row, targets))
    temps = {h: fit(halves[h], tax) for h in (0, 1)}
    out = {}
    for h in (0, 1):
        for row, _ in halves[h]:
            out[row['id']] = probs(row, tax, temps[1 - h])
    return {'temperatures_by_half': {str(h): temps[h] for h in (0, 1)}, 'calibrated': out,
            'sizes': {str(h): len(halves[h]) for h in (0, 1)}}


def load_jsonl(path) -> list[dict]:
    with open(path, encoding='utf-8') as fh:
        return [json.loads(line) for line in fh if line.strip()]


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = parser.add_subparsers(dest='command', required=True)
    f = sub.add_parser('fit', help='Fit temperatures on dev predictions')
    f.add_argument('predictions')
    f.add_argument('--data', required=True, help='Export directory (taxonomy.json, dev.jsonl)')
    f.add_argument('--out', required=True, help='temperatures.json')
    args = parser.parse_args(argv)
    if Path(args.out).exists():
        raise SystemExit('Refusing to overwrite ' + args.out)
    tax = Taxonomy.load(args.data)
    gold = {r['id']: r['targets'] for r in load_jsonl(Path(args.data) / 'dev.jsonl')}
    preds = load_jsonl(args.predictions)
    ids = [p['id'] for p in preds]
    if len(set(ids)) != len(ids) or set(ids) != set(gold):
        raise SystemExit('Predictions must cover each of the %d dev rows exactly once (got %d rows, %d distinct)'
                         % (len(gold), len(ids), len(set(ids))))
    temps = fit([(p, gold[p['id']]) for p in preds], tax)
    edge = sorted(f for f, 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())