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

Input: an export split (``--split dev|benchmark``) or any JSONL of export-format rows (``--input``, e.g. the
pool or the robustness set). Output rows: ``{"id", "split", "n_tokens", "truncated", "logits"}`` with raw logits
per field in taxonomy option order; ``classifier.calibrate`` turns them into calibrated, gated probabilities.
The summary line reports posts/s and tokens/s (the serving throughput) and, with ``--repeat N``, the largest
probability difference between repeats (the determinism check).
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import time

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


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument('--run', required=True, help='Run directory written by classifier.train')
    p.add_argument('--which', default='best', choices=['best', 'last'])
    source = p.add_mutually_exclusive_group(required=True)
    source.add_argument('--split', choices=['train', 'dev', 'benchmark'])
    source.add_argument('--input', help='JSONL of export-format rows')
    p.add_argument('--data', help='Export directory (default: the run config\'s)')
    p.add_argument('--model', help='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')
    p.add_argument('--kernels', default='fast', choices=['fast', 'torch'])
    args = p.parse_args(argv)
    import torch
    from .model import build, load_tokenizer, predict, tokenizer_fn

    run = Path(args.run)
    cfg = json.loads((run / 'config.json').read_text())
    if args.model:
        cfg['model'] = args.model
    data = args.data or cfg['data']
    tax = Taxonomy.load(run / 'taxonomy.json')
    device = torch.device(args.device)
    ckpt = run / args.which
    model = build(cfg, tax.sizes(), device, train=False, adapter_dir=ckpt / 'adapter', heads_path=ckpt / 'heads.pt')
    from .check_env import fast_path, torch_path
    if args.kernels == 'torch':
        torch_path().__enter__()
    kernels = fast_path(model)
    if args.kernels == 'fast' and not kernels['ok']:
        raise SystemExit('Fast kernels are not active: %s' % json.dumps(kernels))
    tok = load_tokenizer(cfg['model'])
    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, tokenizer_fn(tok), cfg['max_len']) 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) / seconds, 1), tokens_per_second=round(tokens / seconds),
                   tokens=tokens, truncated_rows=sum(1 for e in encoded if e[1]['cut']), which=args.which,
                   checkpoint=str(ckpt), kernels=args.kernels)
    if len(runs) > 1:
        diff, same = 0.0, True
        for other, _ in runs[1:]:
            for i in logits:
                for f in tax.names:
                    a, b = softmax(logits[i][f]), softmax(other[i][f])
                    diff = max(diff, max(abs(x - y) for x, y in zip(a, b)))
                    same &= a.index(max(a)) == b.index(max(b))
        summary.update(max_prob_diff=diff, argmax_identical=same)
    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())