"""Predict with a trained run. Run as ``python -m training.predict`` (pod, or a rented GPU for backfills).

Input: a bundle split (``--split val|test``) or any JSONL of rows with ``id`` and ``segments`` (``--input``: the
robustness set, or ``naver-psy predict --export-inputs`` for a backfill, classifier-plan §6.4). Output rows:
``{"id", "split", "n_tokens", "truncated", "logits"}`` with raw logits per field in taxonomy option order (ids and
numbers only); ``training.calibrate`` and the predict worker turn them into calibrated, gated probabilities.
The summary line reports posts/s and tokens/s and, with ``--repeat N``, the largest calibrated probability
difference between repeats (the determinism check, §5.5 item 6).
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from .calibrate import max_prob_diff, outputs
from .data import inference_batches, load_jsonl, load_split, write_jsonl
from .formats import encode
from .taxonomy import Taxonomy


def load_run(run: Path, which: str = "best", model_dir: str | None = None, device="cuda"):
    """(model, tokenizer, taxonomy, config) for a run directory."""
    from .model import build, load_tokenizer

    cfg = json.loads((run / "config.json").read_text())
    if model_dir:
        cfg["model"] = model_dir
    tax = Taxonomy.load(run / "taxonomy.json")
    ckpt = run / which
    model = build(cfg, tax.sizes(), device, train=False, adapter_dir=ckpt / "adapter",
                  heads_path=ckpt / "heads.safetensors")
    return model, load_tokenizer(cfg["model"]), tax, cfg


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--run", required=True, help="Run directory written by training.train")
    p.add_argument("--which", default="best", choices=["best", "last"])
    source = p.add_mutually_exclusive_group(required=True)
    source.add_argument("--split", choices=["train", "val", "test"])
    source.add_argument("--input", help="JSONL of rows with id and segments")
    p.add_argument("--data", help="Bundle directory (default: the run config's)")
    p.add_argument("--model", help="Base checkpoint directory (default: the run config's)")
    p.add_argument("--out", required=True)
    p.add_argument("--token-budget", type=int, default=65536, help="Padded tokens per batch")
    p.add_argument("--limit", type=int, default=0)
    p.add_argument("--repeat", type=int, default=1)
    p.add_argument("--device", default="cuda")
    args = p.parse_args(argv)
    import torch

    from .model import bos_id, predict, tokenizer_fn

    run = Path(args.run)
    device = torch.device(args.device)
    model, tok, tax, cfg = load_run(run, args.which, args.model, device)
    data = args.data or cfg["data"]
    rows = load_split(data, args.split) if args.split else load_jsonl(args.input)
    if args.limit:
        rows = rows[:args.limit]
    t = time.time()
    encoded = [encode(r["segments"], tokenizer_fn(tok), cfg["max_len"], bos_id(tok)) for r in rows]
    ids = [e[0] for e in encoded]
    encode_seconds = time.time() - t
    batches = inference_batches([len(x) for x in ids], args.token_budget)
    runs = []
    for _ in range(max(1, args.repeat)):
        t = time.time()
        logits = predict(model, ids, tok.pad_token_id, batches, device, tax.names)
        if device.type == "cuda":
            torch.cuda.synchronize()
        runs.append((logits, time.time() - t))
    logits, seconds = runs[0]
    out_rows = [{"id": r["id"], "split": r.get("split"), "n_tokens": encoded[i][1]["n_tokens"],
                 "truncated": encoded[i][1]["cut"], "logits": logits[i]} for i, r in enumerate(rows)]
    write_jsonl(args.out, out_rows)
    tokens = sum(len(x) for x in ids)
    summary = dict(rows=len(rows), seconds=round(seconds, 2), encode_seconds=round(encode_seconds, 2),
                   posts_per_second=round(len(rows) / max(seconds, 1e-9), 1),
                   tokens_per_second=round(tokens / max(seconds, 1e-9)), tokens=tokens,
                   truncated_rows=sum(1 for e in encoded if e[1]["cut"]), which=args.which)
    if len(runs) > 1:
        temps_path = run / "temperatures.json"
        temps = json.loads(temps_path.read_text()) if temps_path.exists() else None
        diff, same = 0.0, True
        for other, _ in runs[1:]:
            for i in logits:
                a, b = outputs(logits[i], tax, temps), outputs(other[i], tax, temps)
                diff = max(diff, max_prob_diff(a["probs"], b["probs"]))
                same &= a["labels"] == b["labels"]
        summary.update(max_prob_diff=diff, argmax_identical=same, deterministic=diff <= 1e-3)
    if device.type == "cuda":
        summary["max_memory_gb"] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1)
    print(json.dumps(summary))
    return 0


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