import argparse
import math
import os
import time
import numpy as np
import torch
import torch.nn.functional as F
from modeling import FelaTab, FelaTabConfig, count_params
from prior import PriorConfig, gen_pool


def collate(tasks, task_type, tmax, rng, maxF, s_frac=(0.5, 0.9)):
    B = len(tasks)
    used, sup = ([], [])
    for t in tasks:
        N = t["X"].shape[0]
        n = min(N, tmax)
        frac = rng.uniform(*s_frac)
        s = int(np.clip(int(frac * n), 16, n - 8))
        used.append(n)
        sup.append(s)
    T = max(used)
    X = np.zeros((B, T, maxF), np.float32)
    y = np.zeros((B, T), np.int64 if task_type == "cls" else np.float32)
    lm = np.zeros((B, T), np.float32)
    ncl = np.zeros(B, np.int64)
    ns = np.zeros(B, np.int64)
    for b, t in enumerate(tasks):
        N = t["X"].shape[0]
        n = used[b]
        s = sup[b]
        Ft = min(t["X"].shape[1], maxF)
        perm = rng.permutation(N)[:n]
        X[b, :n, :Ft] = t["X"][perm][:, :Ft]
        y[b, :n] = t["y"][perm]
        lm[b, s:n] = 1.0
        ns[b] = s
        ncl[b] = t["n_classes"] if task_type == "cls" else 0
    return (
        torch.from_numpy(X),
        torch.from_numpy(y),
        torch.from_numpy(ns),
        torch.from_numpy(lm),
        torch.from_numpy(ncl),
    )


def cls_loss_acc(logits, y, lm, ncl):
    B, T, C = logits.shape
    cmask = torch.arange(C, device=logits.device).view(1, C) < ncl.view(B, 1)
    logits = logits.masked_fill(~cmask.unsqueeze(1), -1000000000.0)
    lf = logits.reshape(B * T, C)
    yf = y.reshape(B * T).clamp(min=0)
    mf = lm.reshape(B * T).bool()
    if mf.sum() == 0:
        return (logits.sum() * 0.0, 0.0)
    loss = F.cross_entropy(lf[mf], yf[mf])
    acc = (lf[mf].argmax(-1) == yf[mf]).float().mean().item()
    return (loss, acc)


def reg_loss_r2(out, y, lm):
    mean = out[..., 0]
    logvar = out[..., 1].clamp(-8, 8)
    mf = lm.bool()
    if mf.sum() == 0:
        return (out.sum() * 0.0, 0.0)
    m, lv, t = (mean[mf], logvar[mf], y[mf])
    nll = 0.5 * (lv + (t - m) ** 2 / torch.exp(lv) + math.log(2 * math.pi))
    ss_res = ((t - m) ** 2).sum().item()
    ss_tot = ((t - t.mean()) ** 2).sum().item() + 1e-08
    return (nll.mean(), 1 - ss_res / ss_tot)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=20000)
    ap.add_argument("--bs", type=int, default=16)
    ap.add_argument("--tmax", type=int, default=1024)
    ap.add_argument("--lr", type=float, default=0.0003)
    ap.add_argument("--dim", type=int, default=512)
    ap.add_argument("--layers", type=int, default=14)
    ap.add_argument("--heads", type=int, default=8)
    ap.add_argument("--head_dim", type=int, default=64)
    ap.add_argument("--pool", type=int, default=1000)
    ap.add_argument("--refresh", type=int, default=400)
    ap.add_argument("--ckpt", type=str, default="ckpt/felatab.pt")
    ap.add_argument(
        "--smoke", action="store_true", help="tiny fast loop to prove it trains"
    )
    args = ap.parse_args()
    if args.smoke:
        args.steps, args.dim, args.layers, args.heads = (200, 128, 4, 4)
        args.pool, args.bs, args.tmax = (128, 8, 256)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    prior_cfg = PriorConfig(max_classes=10, max_rows=6000)
    cfg = FelaTabConfig(
        dim=args.dim,
        n_layers=args.layers,
        n_heads=args.heads,
        head_dim=args.head_dim,
        max_features=prior_cfg.max_features,
        max_classes=prior_cfg.max_classes,
    )
    model = FelaTab(cfg).to(dev)
    print(
        f"params {count_params(model) / 1000000.0:.2f}M  dim={cfg.dim} L={cfg.n_layers} H={cfg.n_heads}  dev={dev}"
    )
    opt = torch.optim.AdamW(
        model.parameters(), lr=args.lr, weight_decay=0.01, betas=(0.9, 0.95)
    )
    sched = torch.optim.lr_scheduler.OneCycleLR(
        opt,
        max_lr=args.lr,
        total_steps=args.steps,
        pct_start=0.03,
        anneal_strategy="cos",
    )
    rng = np.random.default_rng(1234)

    def refill(seed):
        pool = gen_pool(prior_cfg, args.pool, seed=seed, workers=None)
        return (
            [t for t in pool if t["task_type"] == "cls"],
            [t for t in pool if t["task_type"] == "reg"],
        )

    print("generating initial synthetic pool...")
    t0 = time.time()
    cls_pool, reg_pool = refill(int(rng.integers(1000000000.0)))
    print(
        f"pool ready {time.time() - t0:.1f}s (cls {len(cls_pool)} reg {len(reg_pool)})"
    )
    model.train()
    run_loss = run_metric = 0.0
    for step in range(1, args.steps + 1):
        typ = "cls" if step % 2 == 0 else "reg"
        src = cls_pool if typ == "cls" else reg_pool
        if not src:
            continue
        T_target = int(rng.integers(128, args.tmax + 1))
        idx = rng.integers(0, len(src), size=args.bs)
        batch = [src[i] for i in idx]
        X, y, ns, lm, ncl = collate(batch, typ, T_target, rng, cfg.max_features)
        X, y, ns, lm, ncl = [z.to(dev) for z in (X, y, ns, lm, ncl)]
        out = model(X, y, ns, torch.tensor(0 if typ == "cls" else 1, device=dev))
        if typ == "cls":
            loss, metric = cls_loss_acc(out.float(), y, lm, ncl)
        else:
            loss, metric = reg_loss_r2(out.float(), y, lm)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        sched.step()
        run_loss += loss.item()
        run_metric += metric
        if step % 50 == 0:
            print(
                f"step {step:6d} {typ} loss {run_loss / 50:.4f} metric {run_metric / 50:.4f} lr {sched.get_last_lr()[0]:.2e}",
                flush=True,
            )
            run_loss = run_metric = 0.0
        if step % args.refresh == 0:
            cls_pool, reg_pool = refill(int(rng.integers(1000000000.0)))
    os.makedirs(os.path.dirname(args.ckpt) or ".", exist_ok=True)
    torch.save(model.state_dict(), args.ckpt)
    print(f"done. saved {args.ckpt}")


if __name__ == "__main__":
    main()