"""Train Modular-TS from scratch with DHF + pinball loss.

python -m modular_ts.train --size 2.5M --ctx 4096 --tokens_per_batch 65536 --steps 20000 \
    --lr 2e-3 --warmup 1000 --out runs/smoke --max_minutes 35 [--compile] [--data synthetic]

Logs JSONL to <out>/log.jsonl (step, loss, val, lr, enc tokens/s, grad norm),
checkpoints to <out>/ckpt.pt and an inference export (config.json + model.safetensors,
loadable by timestral.TimestralForecaster) to <out>/export/.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import time

import torch

from modular_ts.model import ModularTS, count_params, make_config
from modular_ts.synth import SynthStream


def export(model: ModularTS, out_dir: str) -> None:
    from safetensors.torch import save_file
    os.makedirs(out_dir, exist_ok=True)
    sd = {k: v.detach().float().cpu().contiguous() for k, v in model.state_dict().items()}
    save_file(sd, os.path.join(out_dir, "model.safetensors"))
    model.config.save(os.path.join(out_dir, "config.json"))


def lr_at(step, total, warmup, peak, final_frac=0.05):
    if step < warmup:
        return peak * (step + 1) / warmup
    prog = min(1.0, (step - warmup) / max(1, total - warmup))
    return peak * (final_frac + (1 - final_frac) * 0.5 * (1 + math.cos(math.pi * prog)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--size", default="2.5M")
    ap.add_argument("--ctx", type=int, default=4096)
    ap.add_argument("--tokens_per_batch", type=int, default=65536, help="encoder patch tokens per step")
    ap.add_argument("--steps", type=int, default=20000, help="schedule length (cosine horizon)")
    ap.add_argument("--max_minutes", type=float, default=0, help="stop early after this wall time (0 = off)")
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--warmup", type=int, default=1000)
    ap.add_argument("--wd", type=float, default=0.05)
    ap.add_argument("--clip", type=float, default=1.0)
    ap.add_argument("--data", default="synthetic", choices=["synthetic", "synthetic+gifteval"])
    ap.add_argument("--real_frac", type=float, default=0.5)
    ap.add_argument("--gifteval_subset", default="", help="comma list of GiftEvalPretrain sub-datasets (default all kept)")
    ap.add_argument("--gep_selection", default="", help="json {dataset: [shard files]} -> fast GPU-resident GiftEvalPretrain stream")
    ap.add_argument("--synth", default="v1", choices=["v1", "v2"], help="v2 = modular_ts.synth2 prior (v3)")
    ap.add_argument("--p_new", type=float, default=0.55, help="synth v2: share of new generator families")
    ap.add_argument("--real_v2", action="store_true", help="v3 real stream: resolution sampling + TSMixup + ds_temp 0.3")
    ap.add_argument("--min_lr_frac", type=float, default=0.05)
    ap.add_argument("--time_schedule", action="store_true", help="cosine progress = elapsed / max_minutes (smoke runs)")
    ap.add_argument("--compile", action="store_true")
    ap.add_argument("--bf16", action="store_true", default=True)
    ap.add_argument("--no_bf16", dest="bf16", action="store_false")
    ap.add_argument("--log_every", type=int, default=25)
    ap.add_argument("--val_every", type=int, default=250)
    ap.add_argument("--ckpt_every", type=int, default=1000)
    ap.add_argument("--out", default="runs/dev")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--resume_from", default="", help="ckpt to resume model+opt from (default <out>/ckpt.pt)")
    ap.add_argument("--reset_schedule", action="store_true", help="on resume, restart warmup+cosine at the resumed step")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    torch.manual_seed(args.seed)
    dev = torch.device(args.device)
    if dev.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    cfg = make_config(args.size, context_length=args.ctx)
    model = ModularTS(cfg).to(dev)
    pc = count_params(model)
    print("params", pc, flush=True)

    if args.synth == "v2":
        from modular_ts.synth2 import SynthStream2
        synth = SynthStream2(args.ctx, args.tokens_per_batch, dev, seed=args.seed, p_new=args.p_new)
    else:
        synth = SynthStream(args.ctx, args.tokens_per_batch, dev, seed=args.seed)
    stream = synth
    manifest = {"data": args.data, "synth_version": args.synth, "real_v2": args.real_v2, "synthetic": "modular_ts.synth (KernelSynth-style GP + TSI + AR + intermittent/bursty + multivariatizers + covariate tasks)"}
    if args.data == "synthetic+gifteval":
        from modular_ts.realdata import GiftEvalPretrainStream, MixedStream
        sub = [s for s in args.gifteval_subset.split(",") if s] or None
        if args.gep_selection:
            from modular_ts.realdata import FastGiftEvalPretrain
            if args.real_v2:
                from modular_ts.realdata2 import FastGiftEvalPretrain2 as FastGiftEvalPretrain
            real = FastGiftEvalPretrain(args.ctx, args.tokens_per_batch, dev, json.load(open(args.gep_selection)), seed=args.seed)
        else:
            real = GiftEvalPretrainStream(args.ctx, args.tokens_per_batch, dev, seed=args.seed, datasets_subset=sub)
        stream = MixedStream(synth, real, args.real_frac, seed=args.seed)
        manifest.update(real.manifest())
    json.dump(manifest, open(os.path.join(args.out, "data_manifest.json"), "w"), indent=1)

    # fixed held-out validation batches from a different seed (one per V bucket)
    vgen = SynthStream(args.ctx, args.tokens_per_batch // 2, dev, seed=10_000 + args.seed, chunk=1)
    val_batches = [vgen.next_batch(V) for V in (1, 4, 16)]

    decay, no_decay = [], []
    for n, p in model.named_parameters():
        (decay if p.dim() >= 2 else no_decay).append(p)
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": args.wd}, {"params": no_decay, "weight_decay": 0.0}],
                            lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=dev.type == "cuda")
    step = 0
    ck = os.path.join(args.out, "ckpt.pt")
    rck = args.resume_from or ck
    step0 = 0
    prev_elapsed = 0.0
    if args.resume and os.path.exists(rck):
        st = torch.load(rck, map_location=dev)
        model.load_state_dict(st["model"])
        opt.load_state_dict(st["opt"])
        step = st["step"]
        if args.reset_schedule:
            step0 = step
        elif args.time_schedule and "elapsed_min" in st:
            prev_elapsed = st["elapsed_min"]  # continue the same time-based schedule after a crash
        print("resumed at", step, "from", rck, flush=True)

    fwd = model.forward_train
    if args.compile:
        try:
            fwd = torch.compile(model.forward_train, dynamic=False)
        except Exception as e:  # pragma: no cover
            print("compile failed, eager:", e, flush=True)

    amp = torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=args.bf16 and dev.type == "cuda")

    def evaluate():
        model.eval()
        tot = 0.0
        with torch.no_grad(), amp:
            for b in val_batches:
                l, _ = model.forward_train(b["values"], b["observed"], b["roles"], b["s"])
                tot += l.item()
        model.train()
        return tot / len(val_batches)

    log = open(os.path.join(args.out, "log.jsonl"), "a")
    t0 = time.time() - prev_elapsed * 60
    t_last, tok_acc, gen_t = time.time(), 0, 0.0
    ema = None
    model.train()
    while step < args.steps:
        if args.max_minutes and (time.time() - t0) / 60 > args.max_minutes:
            print("time limit reached", flush=True)
            break
        prof = step % args.log_every == 0  # measure data-gen share only on log steps (sync is costly)
        if prof and dev.type == "cuda":
            torch.cuda.synchronize()
        tg = time.time()
        b = stream.next_batch()
        if prof and dev.type == "cuda":
            torch.cuda.synchronize()
        if prof:
            gen_t += (time.time() - tg) * args.log_every
        ss = step - step0
        if args.time_schedule and args.max_minutes:
            frac = (time.time() - t0) / 60 / args.max_minutes
            lr = lr_at(ss, max(ss + 1, int(ss / max(frac, 1e-6))) if ss > args.warmup else args.steps, args.warmup, args.lr, args.min_lr_frac)
        else:
            lr = lr_at(ss, args.steps, args.warmup, args.lr, args.min_lr_frac)
        for g in opt.param_groups:
            g["lr"] = lr
        try:
            with amp:
                loss, st = fwd(b["values"], b["observed"], b["roles"], b["s"])
        except Exception as e:
            if fwd is model.forward_train:
                raise
            print("compiled forward failed, falling back to eager:", repr(e)[:500], flush=True)
            fwd = model.forward_train
            with amp:
                loss, st = fwd(b["values"], b["observed"], b["roles"], b["s"])
        if not torch.isfinite(loss):
            print(f"non-finite loss at step {step}; skipping batch", flush=True)
            opt.zero_grad(set_to_none=True)
            step += 1
            continue
        loss.backward()
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
        opt.step()
        opt.zero_grad(set_to_none=True)
        tok_acc += b["values"].shape[0] * b["values"].shape[1] * (args.ctx // cfg.patch_size)
        step += 1
        lv = loss.item() if step % args.log_every == 0 or step == 1 else None
        if lv is not None:
            ema = lv if ema is None else 0.9 * ema + 0.1 * lv
            if dev.type == "cuda":
                torch.cuda.synchronize()
            dt = time.time() - t_last
            rec = {"step": step, "loss": lv, "loss_ema": ema, "lr": lr, "gnorm": float(gn),
                   "enc_tok_s": tok_acc / dt, "gen_frac": gen_t / dt, "elapsed_min": (time.time() - t0) / 60}
            if step % args.val_every == 0:
                rec["val"] = evaluate()
            if dev.type == "cuda":
                rec["mem_gb"] = torch.cuda.max_memory_allocated() / 1e9
            print(json.dumps(rec), flush=True)
            log.write(json.dumps(rec) + "\n")
            log.flush()
            t_last, tok_acc, gen_t = time.time(), 0, 0.0
        if step % args.ckpt_every == 0:
            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "args": vars(args), "elapsed_min": (time.time() - t0) / 60}, ck)
            export(model, os.path.join(args.out, "export"))
    torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "args": vars(args), "elapsed_min": (time.time() - t0) / 60}, ck)
    export(model, os.path.join(args.out, "export"))
    fin = {"final_step": step, "val": evaluate(), "params": pc, "minutes": (time.time() - t0) / 60}
    print("FINAL", json.dumps(fin), flush=True)
    json.dump(fin, open(os.path.join(args.out, "final.json"), "w"))


if __name__ == "__main__":
    main()