"""Train the heads classifier on the export. Run on the pod as ``python -m classifier.train``.

Adapted from the read-only reference ~/yt_live_analysis/training/train.py (arm H only), with the spec's
fixes: steps of a fixed number of posts split into token-budget micro-batches with per-step loss
normalisation; nulls masked; a cumulative ``--max-hours`` kept in state.json (a resume cannot extend it);
``--stop-after`` to test resume without changing the schedule; and a forced step on the longest rows to
measure peak memory.

Writes OUT/config.json, OUT/train_log.jsonl, OUT/last/ (resume point), OUT/best/ (lowest dev mean field
log-loss), OUT/predictions/dev.jsonl and OUT/metrics_dev.json for the best checkpoint. Exit code 3: stopped by
--max-hours or --stop-after at a saved resume point; rerun with --resume.
"""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
import platform
import random
import time

from .calibrate import nll, probs as calibrated_probs
from .data import inference_batches, load_split, micro_batches, step_counts, steps, subset, write_jsonl
from .formats import FORMAT, encode
from .taxonomy import Taxonomy, argmax

CORE = ('data', 'model', 'subset', 'epochs', 'max_steps', 'posts_per_step', 'token_budget', 'max_len', 'seed',
        'lr_lora', 'lr_heads', 'warmup', 'weight_decay', 'lora_r', 'lora_alpha', 'lora_dropout', 'head_dropout',
        'dtype', 'kernels')


def parse(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument('--data', required=True, help='Export directory (train/dev.jsonl, taxonomy.json, manifest.json)')
    p.add_argument('--model', required=True, help='Local checkpoint directory')
    p.add_argument('--out', required=True)
    p.add_argument('--subset', type=int, default=0, help='Train on this many rows in seeded order (0 = all)')
    p.add_argument('--epochs', type=float, default=2.0)
    p.add_argument('--max-steps', type=int, default=0, help='Cap the schedule at this many optimiser steps')
    p.add_argument('--posts-per-step', type=int, default=32)
    p.add_argument('--token-budget', type=int, default=16384, help='Padded tokens per training micro-batch')
    p.add_argument('--eval-token-budget', type=int, default=65536)
    p.add_argument('--max-len', type=int, default=2048)
    p.add_argument('--lr-lora', type=float, default=2e-4)
    p.add_argument('--lr-heads', type=float, default=1e-3)
    p.add_argument('--warmup', type=float, default=0.03)
    p.add_argument('--weight-decay', type=float, default=0.01)
    p.add_argument('--lora-r', type=int, default=32)
    p.add_argument('--lora-alpha', type=int, default=32)
    p.add_argument('--lora-dropout', type=float, default=0.05)
    p.add_argument('--head-dropout', type=float, default=0.1)
    p.add_argument('--eval-every', type=int, default=200, help='Optimiser steps between dev evaluations and saves')
    p.add_argument('--dev-limit', type=int, default=0, help='Evaluate on the first N dev rows in seeded order')
    p.add_argument('--log-every', type=int, default=10)
    p.add_argument('--max-hours', type=float, default=0, help='Cumulative training hours across resumes')
    p.add_argument('--stop-after', type=int, default=0, help='Stop this process after N more steps (resume test)')
    p.add_argument('--force-longest', type=int, default=0, help='First run one forward/backward on the N longest rows')
    p.add_argument('--seed', type=int, default=20261017)
    p.add_argument('--resume', action='store_true',
                   help='Continue from last/ if it exists (settings must match); otherwise start fresh')
    p.add_argument('--kernels', default='fast', choices=['fast', 'torch'],
                   help='fast: require fla + causal-conv1d (asserted); torch: the PyTorch reference path')
    p.add_argument('--device', default='cuda')
    p.add_argument('--dtype', default='bfloat16', choices=['bfloat16', 'float32'])
    return p.parse_args(argv)


def versions() -> dict:
    import peft
    import torch
    import transformers
    out = dict(python=platform.python_version(), torch=torch.__version__, transformers=transformers.__version__,
               peft=peft.__version__, cuda=torch.version.cuda)
    for pkg in ('triton', 'fla', 'causal_conv1d', 'tilelang'):
        try:
            out[pkg] = getattr(__import__(pkg), '__version__', 'unknown')
        except Exception:
            out[pkg] = None
    commit = Path(__file__).with_name('.source_commit')
    out['source_commit'] = commit.read_text().strip() if commit.exists() else None
    return out


class Log:
    def __init__(self, path: Path):
        self.fh = open(path, 'a', encoding='utf-8')

    def __call__(self, **row):
        row = dict(time=round(time.time(), 1), **row)
        line = json.dumps(row, ensure_ascii=False)
        self.fh.write(line + '\n')
        self.fh.flush()
        print(line, flush=True)


def prediction_rows(rows, logits_by_index, infos, split=None) -> list[dict]:
    return [{'id': r['id'], 'split': split or r.get('split'), 'n_tokens': infos[i]['n_tokens'], 'truncated': infos[i]['cut'],
             'logits': logits_by_index[i]} for i, r in enumerate(rows)]


def dev_report(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict:
    pairs = [(p, r['targets']) for p, r in zip(preds, rows)]
    per_field = {f: nll(pairs, tax, f, 1.0) for f in tax.names}
    scored = [v for v in per_field.values() if v is not None]
    acc = {}
    for f in tax.names:
        gold = [(p, t[f]) for p, t in pairs if t.get(f) is not None]
        acc[f] = sum(1 for p, g in gold if argmax(calibrated_probs(p, tax))[f] == g) / len(gold) if gold else None
    return {'selection_loss': sum(scored) / len(scored), 'log_loss': per_field, 'accuracy': acc, 'rows': len(rows)}


def main(argv=None):
    args = parse(argv)
    import torch
    from .model import build, pad, predict, save_checkpoint, step_loss, load_tokenizer, tokenizer_fn

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    cfg_path = out / 'config.json'
    cfg = {k: v for k, v in vars(args).items() if k not in ('resume', 'device', 'stop_after', 'force_longest')}
    if cfg_path.exists():
        if not args.resume:
            raise SystemExit('Refusing to start a new run over %s; use --resume or a new --out' % out)
        old = json.loads(cfg_path.read_text())
        changed = sorted(k for k in CORE if old.get(k) != cfg.get(k))
        if changed:
            raise SystemExit('--resume with different settings: %s' % changed)
        if old.get('max_hours') != cfg.get('max_hours'):
            raise SystemExit('--resume with a different --max-hours; the cap is cumulative')
    manifest = json.loads((Path(args.data) / 'manifest.json').read_text())
    cfg.update(versions=versions(), format=FORMAT, data_version=manifest.get('version'),
               data_files_sha256=manifest.get('files_sha256') or manifest.get('sha256'))
    cfg_path.write_text(json.dumps(cfg, indent=1))
    log = Log(out / 'train_log.jsonl')

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = torch.device(args.device)
    tax = Taxonomy.load(args.data)
    (out / 'taxonomy.json').write_text(json.dumps(tax.to_json(), ensure_ascii=False, indent=1))
    tok = load_tokenizer(args.model)
    tokenize = tokenizer_fn(tok)

    train_rows = subset(load_split(args.data, 'train'), args.subset, args.seed)
    dev_rows = subset(load_split(args.data, 'dev'), args.dev_limit, args.seed)
    t0 = time.time()
    enc_train = [encode(r, tokenize, args.max_len) for r in train_rows]
    enc_dev = [encode(r, tokenize, args.max_len) for r in dev_rows]
    ids = [e[0] for e in enc_train]
    lengths = [len(x) for x in ids]
    labels = [tax.encode(r['targets']) for r in train_rows]
    dev_ids = [e[0] for e in enc_dev]
    dev_infos = [e[1] for e in enc_dev]
    dev_batches = inference_batches([len(x) for x in dev_ids], args.eval_token_budget)
    cut = sum(1 for e in enc_train if e[1]['cut'])
    log(event='data', train=len(train_rows), dev=len(dev_rows), tokens=sum(lengths), max_len=max(lengths),
        mean_len=round(sum(lengths) / len(lengths), 1), truncated_rows=cut,
        dev_tokens=sum(len(x) for x in dev_ids), seconds=round(time.time() - t0, 1))

    last_dir, best_dir = out / 'last', out / 'best'
    resume = args.resume and (last_dir / 'state.json').exists()
    model = build(cfg, tax.sizes(), device, train=True)
    from .check_env import fast_path, torch_path
    if args.kernels == 'torch':
        torch_path().__enter__()  # for the whole process: the pre-declared PyTorch-path fallback
    cfg['fast_path'] = fast_path(model)
    cfg_path.write_text(json.dumps(cfg, indent=1))
    if args.kernels == 'fast' and not cfg['fast_path']['ok']:
        raise SystemExit('Fast kernels are not active: %s' % json.dumps(cfg['fast_path']))
    if resume:
        from peft import set_peft_model_state_dict
        from safetensors.torch import load_file
        set_peft_model_state_dict(model.backbone, load_file(str(last_dir / 'adapter' / 'adapter_model.safetensors')))
        model.heads.load_state_dict(torch.load(last_dir / 'heads.pt', map_location='cpu'))
    lora = [p for n, p in model.named_parameters() if p.requires_grad and 'lora_' in n]
    heads = [p for n, p in model.named_parameters() if p.requires_grad and 'lora_' not in n]
    opt = torch.optim.AdamW([dict(params=lora, lr=args.lr_lora), dict(params=heads, lr=args.lr_heads)],
                            weight_decay=args.weight_decay)
    steps_per_epoch = math.ceil(len(train_rows) / args.posts_per_step)
    total = max(1, int(steps_per_epoch * args.epochs))
    if args.max_steps:
        total = min(total, args.max_steps)
    from transformers import get_cosine_schedule_with_warmup
    sched = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total)
    state = dict(step=0, epoch=0, step_in_epoch=0, best=None, best_step=None, hours_used=0.0)
    if resume:
        state = json.loads((last_dir / 'state.json').read_text())
        trainer = torch.load(last_dir / 'trainer.pt', map_location='cpu', weights_only=False)
        opt.load_state_dict(trainer['opt'])
        sched.load_state_dict(trainer['sched'])
        log(event='resume', **state)
    log(event='start', trainable=sum(p.numel() for p in lora + heads), steps=total, steps_per_epoch=steps_per_epoch)

    amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == 'cuda')
    process_start = time.time()
    hours_before = state['hours_used']

    def hours_now():
        return hours_before + (time.time() - process_start) / 3600

    def save_last():
        state['hours_used'] = round(hours_now(), 4)
        save_checkpoint(model, last_dir, dict(state), dict(opt=opt.state_dict(), sched=sched.state_dict()))

    def evaluate():
        t = time.time()
        logits = predict(model, dev_ids, tok.pad_token_id, dev_batches, device, tax.names)
        preds = prediction_rows(dev_rows, logits, dev_infos, 'dev')
        rep = dev_report(preds, dev_rows, tax)
        seconds = time.time() - t
        log(event='eval', step=state['step'], selection_loss=round(rep['selection_loss'], 5),
            log_loss={f: (round(v, 4) if v is not None else None) for f, v in rep['log_loss'].items()},
            accuracy={f: (round(v, 4) if v is not None else None) for f, v in rep['accuracy'].items()},
            seconds=round(seconds, 1), posts_per_second=round(len(dev_rows) / seconds, 1))
        if state['best'] is None or rep['selection_loss'] < state['best']:
            state['best'], state['best_step'] = rep['selection_loss'], state['step']
            save_checkpoint(model, best_dir, dict(state))
            write_jsonl(out / 'predictions' / 'dev.jsonl', preds)
            (out / 'metrics_dev.json').write_text(json.dumps(dict(step=state['step'], **rep), indent=1))

    if args.force_longest and not resume:
        # The worst-case micro-batch for peak memory: the N longest rows of the whole train split (not the
        # subset), which fill the token budget at max_len. One forward/backward, no optimiser step.
        pool = sorted(load_split(args.data, 'train'), key=lambda r: -sum(len(x['text']) for x in r['segments']))
        cand = [(encode(r, tokenize, args.max_len)[0], tax.encode(r['targets'])) for r in pool[:4 * args.force_longest]]
        cand = sorted(cand, key=lambda c: -len(c[0]))[:args.force_longest]
        f_ids, f_labels = [c[0] for c in cand], [c[1] for c in cand]
        f_len = [len(x) for x in f_ids]
        parts = micro_batches(list(range(len(cand))), f_len, args.token_budget)
        for mb in parts:
            x, m, last = pad([f_ids[i] for i in mb], tok.pad_token_id, device)
            y = torch.tensor([f_labels[i] for i in mb], device=device)
            with torch.autocast(**amp):
                loss, _ = step_loss(model(x, m, last), y, step_counts(f_labels))
            loss.backward()
        opt.zero_grad(set_to_none=True)
        mem = round(torch.cuda.max_memory_allocated() / 2 ** 30, 2) if device.type == 'cuda' else None
        log(event='forced_longest', rows=len(cand), lengths=f_len, micro_batches=len(parts),
            padded_tokens=[max(f_len[i] for i in mb) * len(mb) for mb in parts], max_memory_gb=mem)
        if device.type == 'cuda':
            torch.cuda.reset_peak_memory_stats()

    model.train()
    try:
        stop = run_steps(args, state, total, model, opt, sched, lora, heads, ids, lengths, labels, tax, device, amp,
                         tok, log, hours_now, evaluate, save_last)
    except KeyboardInterrupt:  # timeout sends SIGINT: keep a resume point at the last completed step
        log(event='interrupted', step=state['step'])
        save_last()
        return 3
    if stop == 4:
        return 4
    save_last()
    if state['best'] is None:
        evaluate()
    log(event='done' if stop is None else 'stopped', **state)
    return stop or 0


def run_steps(args, state, total, model, opt, sched, lora, heads, ids, lengths, labels, tax, device, amp, tok, log,
              hours_now, evaluate, save_last):
    """The optimiser loop. Returns None (finished), 3 (cap or planned stop, resume point saved by the caller)
    or 4 (non-finite loss or gradient: nothing is saved, so last/ keeps the previous good resume point)."""
    import torch
    from .model import pad, step_loss
    window = dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names),
                  field_counts=[0] * len(tax.names), n=0)
    stop, done_here = None, 0
    while state['step'] < total and stop is None:
        epoch_steps = steps(lengths, args.posts_per_step, args.seed * 1000 + state['epoch'])
        while state['step_in_epoch'] < len(epoch_steps) and state['step'] < total:
            if args.max_hours and hours_now() >= args.max_hours:
                stop = 3
                break
            step = epoch_steps[state['step_in_epoch']]
            step_labels = [labels[i] for i in step]
            counts = step_counts(step_labels)
            step_total = 0.0
            for mb in micro_batches(step, lengths, args.token_budget):
                x, m, last = pad([ids[i] for i in mb], tok.pad_token_id, device)
                y = torch.tensor([labels[i] for i in mb], device=device)
                with torch.autocast(**amp):
                    loss, per_field = step_loss(model(x, m, last), y, counts)
                loss.backward()
                step_total += float(loss.detach())
                window['field_sums'] = [a + b for a, b in zip(window['field_sums'], per_field)]
                window['tokens'] += int(m.sum())
                window['padded'] += int(m.numel())
            window['field_counts'] = [a + sum(1 for r in step_labels if r[f] >= 0)
                                      for f, a in enumerate(window['field_counts'])]
            norm = float(torch.nn.utils.clip_grad_norm_(lora + heads, 1.0))
            if not (math.isfinite(step_total) and math.isfinite(norm)):
                log(event='non_finite', step=state['step'] + 1, loss=step_total, grad_norm=norm)
                opt.zero_grad(set_to_none=True)
                return 4
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            state['step'] += 1
            state['step_in_epoch'] += 1
            done_here += 1
            window['loss'] += step_total
            window['n'] += 1
            if state['step'] % args.log_every == 0 or state['step'] == total:
                dt = time.time() - window['t']
                row = dict(event='train', step=state['step'], epoch=state['epoch'], loss=round(window['loss'] / window['n'], 4),
                           fields={f: round(s / c, 4) if c else None for f, s, c in
                                   zip(tax.names, window['field_sums'], window['field_counts'])},
                           lr=sched.get_last_lr()[0], tokens_per_second=round(window['tokens'] / dt),
                           padded_tokens_per_second=round(window['padded'] / dt), hours=round(hours_now(), 3))
                if device.type == 'cuda':
                    row['max_memory_gb'] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1)
                log(**row)
                window = dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names),
                              field_counts=[0] * len(tax.names), n=0)
            if state['step'] % args.eval_every == 0 or state['step'] == total:
                evaluate()
                save_last()
            if args.stop_after and done_here >= args.stop_after:
                stop = 3
                break
        if state['step_in_epoch'] >= len(epoch_steps):
            state['epoch'] += 1
            state['step_in_epoch'] = 0
    return stop


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