import math
import os
import time
import argparse

import torch
import torch.nn.functional as F

from ar.model.config import ARConfig, small_config, tiny_config
from ar.model.model import ARModel
from ar.data.dataset import BinDataset
from ar.train.train_config import TrainConfig


def lr_at(step: int, cfg: TrainConfig) -> float:
    """Warmup then cosine decay: ramp up linearly, then ease down to min_lr."""
    if step < cfg.warmup_steps:
        return cfg.peak_lr * (step + 1) / cfg.warmup_steps
    # cosine decay from peak_lr to min_lr over the remaining steps
    progress = (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps)
    progress = min(1.0, progress)
    coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
    return cfg.min_lr + coeff * (cfg.peak_lr - cfg.min_lr)


def train(bin_paths, weights, tcfg: TrainConfig, mcfg: ARConfig, device="cuda"):
    torch.manual_seed(0)
    model = ARModel(mcfg).to(device)
    model.train()

    data = BinDataset(bin_paths, tcfg.seq_len, weights)

    # AdamW with the LLM-pretraining betas; fused=True is a faster CUDA path
    opt = torch.optim.AdamW(
        model.parameters(), lr=tcfg.peak_lr,
        betas=tcfg.betas, weight_decay=tcfg.weight_decay, fused=True,
    )

    os.makedirs(tcfg.ckpt_dir, exist_ok=True)
    t0 = time.time()

    for step in range(tcfg.max_steps):
        # set this step's learning rate on the optimizer
        lr = lr_at(step, tcfg)
        for g in opt.param_groups:
            g["lr"] = lr

        x, y = data.get_batch(tcfg.batch_size, device)

        # forward: logits are (B, T, vocab); cross-entropy wants (B*T, vocab)
        logits = model(x)
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)), y.view(-1)
        )

        # backward: clear old grads, compute new ones, clip, then step
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), tcfg.grad_clip)
        opt.step()

        if step % tcfg.log_every == 0:
            dt = time.time() - t0
            toks = (step + 1) * tcfg.batch_size * tcfg.seq_len
            print(f"step {step:6d} | loss {loss.item():.4f} | lr {lr:.2e} "
                  f"| {toks/dt/1e3:.1f}k tok/s")

        if step > 0 and step % tcfg.ckpt_every == 0:
            path = os.path.join(tcfg.ckpt_dir, f"ar_step{step}.pt")
            torch.save({"model": model.state_dict(), "step": step}, path)
            print(f"  saved {path}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--bin", nargs="+", required=True, help="one or more .bin paths")
    p.add_argument("--weights", nargs="+", type=float, default=None)
    p.add_argument("--max_steps", type=int, default=100)
    p.add_argument("--batch_size", type=int, default=4)
    p.add_argument("--seq_len", type=int, default=4096)
    p.add_argument("--warmup_steps", type=int, default=2000)
    p.add_argument("--model_size", choices=["small", "full", "tiny"], default="small")
    args = p.parse_args()

    tcfg = TrainConfig(max_steps=args.max_steps, batch_size=args.batch_size,
                       seq_len=args.seq_len, warmup_steps=args.warmup_steps)
    mcfg = {"small": small_config, "full": ARConfig, "tiny": tiny_config}[args.model_size]()
    train(args.bin, args.weights, tcfg, mcfg)