"""Train one pointer-chasing model (fixed-depth or recurrent).

Example:
  python train.py --arch rec --depth 1 --iters 4 --seed 0 \
      --run-name rec1x4-seed0 --hub-repo eac123/ptr-recurrent-depth \
      --space-id <trackio space id>

Trains on fresh random mappings every batch (no fixed train set), k ~ U{1..8},
and logs to trackio. Saves a checkpoint + training history JSON locally and
uploads both to the Hub repo.
"""

import argparse
import json
import os
import time

import numpy as np
import torch
import torch.nn.functional as F

from ptr.common import sample_train_batch, sample_eval_batch
from ptr.flops import count_params, flops_budget_summary, step_flops
from ptr.models import PtrTransformer

PROBE_KS = (16, 32, 64)


def make_parser():
    p = argparse.ArgumentParser()
    p.add_argument("--arch", choices=["fixed", "rec"], required=True)
    p.add_argument("--depth", type=int, default=4)
    p.add_argument("--iters", default="1",
                   help="int, or 'dynamic' for R ~ U{2..8} per batch")
    p.add_argument("--d-model", type=int, default=128)
    p.add_argument("--S", type=int, default=64)
    p.add_argument("--k-max", type=int, default=8)
    p.add_argument("--steps", type=int, default=25000)
    p.add_argument("--batch", type=int, default=256)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--warmup", type=int, default=1000)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--run-name", required=True)
    p.add_argument("--out-dir", default="out")
    p.add_argument("--hub-repo", default=None)
    p.add_argument("--space-id", default=None,
                   help="trackio space id for the metrics dashboard")
    p.add_argument("--project", default="ptr-recurrent-depth")
    p.add_argument("--eval-every", type=int, default=2500)
    p.add_argument("--eval-n", type=int, default=400)
    return p


def resolve_iters(args):
    if args.iters == "dynamic":
        return ("dynamic", 2, 8)
    if args.iters == "aligned":
        # one unroll per hop: R = k per example (all examples in a batch
        # share the same k so a single unroll count serves the batch)
        return ("aligned", 1, args.k_max)
    return int(args.iters)


def sample_iters(spec, rng):
    if isinstance(spec, tuple):
        return int(rng.integers(spec[1], spec[2] + 1))
    return spec


def lr_at(step, args):
    if step < args.warmup:
        return args.lr * (step + 1) / args.warmup
    t = (step - args.warmup) / max(1, args.steps - args.warmup)
    return args.lr * 0.5 * (1 + np.cos(np.pi * t))


@torch.no_grad()
def probe(model, rng, device, S, ks, n, n_iters, aligned=False):
    accs = {}
    for k in ks:
        b = sample_eval_batch(rng, n, S, k=k)
        f = torch.as_tensor(b["f"], device=device)
        x = torch.as_tensor(b["x"], device=device)
        kt = torch.as_tensor(b["k"], device=device)
        y = torch.as_tensor(b["answer"], device=device)
        r = k if aligned else n_iters
        logits = model(f, x, kt, n_iters=r)
        accs[k] = float((logits.argmax(-1) == y).float().mean())
    return accs


def main():
    args = make_parser().parse_args()
    os.makedirs(args.out_dir, exist_ok=True)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    rng = np.random.default_rng(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    iters_spec = resolve_iters(args)

    model = PtrTransformer(args.S, args.d_model, 4, 4 * args.d_model,
                           depth=args.depth, zero_init=True).to(device)
    nparams = count_params(model)
    budget = flops_budget_summary(model, args.S, args.batch, iters_spec,
                                  args.steps)
    print(f"[budget] {json.dumps(budget)}")
    # cheap assertions up front
    assert budget["params_total"] > 0 and budget["train_flops_total"] > 0

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)

    trackio = None
    if args.space_id:
        import trackio
        trackio.init(project=args.project, space_id=args.space_id,
                     config={**vars(args), **budget})

    history = {"loss": [], "id_acc": [], "ood": [], "iters": [],
               "flops_cum": 0.0}
    flops_cum = 0.0
    t0 = time.time()
    bytes_per_batch = 0
    for step in range(args.steps):
        lr = lr_at(step, args)
        for g in opt.param_groups:
            g["lr"] = lr
        if isinstance(iters_spec, tuple) and iters_spec[0] == "aligned":
            kb = int(rng.integers(iters_spec[1], iters_spec[2] + 1))
            n_iters = kb
        else:
            n_iters = sample_iters(iters_spec, rng)
            kb = None
        b = sample_train_batch(rng, args.batch, args.S, k_max=args.k_max)
        if kb is not None:
            b["k"][:] = kb
        f = torch.as_tensor(b["f"], device=device)
        x = torch.as_tensor(b["x"], device=device)
        k = torch.as_tensor(b["k"], device=device)
        y = torch.as_tensor(b["answer"], device=device)
        logits = model(f, x, k, n_iters=n_iters)
        loss = F.cross_entropy(logits, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        flops_cum += step_flops(model, args.S, args.batch, n_iters, True)

        if step % 100 == 0 or step == args.steps - 1:
            ips = 100 / max(1e-9, time.time() - t0) if step > 0 else 0.0
            t0 = time.time()
            rec = {"step": step, "loss": round(float(loss), 4), "lr": lr,
                   "iters": n_iters, "steps_per_s": round(ips, 2),
                   "flops_cum": flops_cum}
            history["loss"].append(rec)
            if device == "cuda" and step == 100:
                bytes_per_batch = torch.cuda.max_memory_allocated()
                print(f"[mem] peak GPU memory {bytes_per_batch/1e6:.1f} MB")
            if trackio:
                trackio.log({"train/loss": float(loss), "train/lr": lr,
                             "train/steps_per_s": ips,
                             "train/flops_cum": flops_cum,
                             "train/iters": n_iters}, step=step)
        if (step + 1) % args.eval_every == 0 or step == args.steps - 1:
            aligned = isinstance(iters_spec, tuple) \
                and iters_spec[0] == "aligned"
            n_eval = sample_iters(iters_spec, rng) if args.arch == "rec" \
                and not aligned else 1
            id_accs = probe(model, rng, device, args.S,
                            list(range(1, args.k_max + 1)), args.eval_n,
                            n_eval, aligned=aligned)
            ood = probe(model, rng, device, args.S, PROBE_KS, args.eval_n,
                        n_eval, aligned=aligned)
            id_mean = float(np.mean(list(id_accs.values())))
            print(f"[eval] step {step+1} id_acc(k=1..8) {id_mean:.4f} "
                  f"ood {ood}")
            history["id_acc"].append({"step": step + 1, "mean": id_mean,
                                      "per_k": id_accs})
            history["ood"].append({"step": step + 1, "acc": ood})
            if trackio:
                trackio.log({"eval/id_acc_mean": id_mean,
                             "eval/ood_k16": ood[16], "eval/ood_k32": ood[32],
                             "eval/ood_k64": ood[64]}, step=step + 1)
            history["flops_cum"] = flops_cum

    ckpt_path = os.path.join(args.out_dir, f"{args.run_name}.pt")
    torch.save({"state_dict": model.state_dict(),
                "config": vars(args), "budget": budget,
                "n_iters_spec": iters_spec}, ckpt_path)
    hist_path = os.path.join(args.out_dir, f"{args.run_name}_history.json")
    with open(hist_path, "w") as fh:
        json.dump(history, fh, indent=1)

    if args.hub_repo:
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(path_or_fileobj=ckpt_path,
                        path_in_repo=f"ckpts/{args.run_name}.pt",
                        repo_id=args.hub_repo, repo_type="model")
        api.upload_file(path_or_fileobj=hist_path,
                        path_in_repo=f"results/{args.run_name}_history.json",
                        repo_id=args.hub_repo, repo_type="model")
        print(f"[hub] uploaded ckpt + history to {args.hub_repo}")
    if trackio:
        trackio.finish()
    print(f"[done] {args.run_name}: final loss "
          f"{history['loss'][-1]['loss']}, "
          f"id_acc {history['id_acc'][-1]['mean']:.4f}")


if __name__ == "__main__":
    main()