#!/usr/bin/env python3
import argparse
import gc
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from mla.backend import NAME, DTYPE
from mla.optim import AdamW, clip_grad_norm
from mla.loss import cross_entropy
from mla.schedule import lr_schedule
from mla.checkpoint import load_checkpoint, save_checkpoint

SFT = Path("data/sft")
CKPT_DIR = Path("checkpoints")
BASE = CKPT_DIR / "pretrain_final.npz"


def get_sft_batch(ids, mask, block_size, batch_size, rng):
    hi = len(ids) - block_size - 1
    ix = rng.integers(0, hi, size=batch_size)
    x = np.stack([ids[i:i + block_size] for i in ix]).astype(np.int64)
    y = np.stack([ids[i + 1:i + 1 + block_size] for i in ix]).astype(np.int64)
    ym = np.stack([mask[i + 1:i + 1 + block_size] for i in ix])
    y[ym == 0] = -1
    return x, y


def eval_masked(model, ids, mask, block_size, batch_size, n_batches, rng):
    total = 0.0
    for _ in range(n_batches):
        x, y = get_sft_batch(ids, mask, block_size, batch_size, rng)
        loss = cross_entropy(model(x), y)
        total += float(loss.data)
        del loss
        gc.collect()
    mean = total / max(1, n_batches)
    return mean, float(np.exp(mean))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=800)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--block-size", type=int, default=256)
    ap.add_argument("--peak-lr", type=float, default=1e-4)
    ap.add_argument("--min-lr", type=float, default=1e-5)
    ap.add_argument("--warmup", type=int, default=50)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--max-norm", type=float, default=1.0)
    ap.add_argument("--log-every", type=int, default=50)
    ap.add_argument("--eval-every", type=int, default=100)
    ap.add_argument("--eval-batches", type=int, default=20)
    ap.add_argument("--ckpt-every", type=int, default=400)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--base", type=str, default=str(BASE))
    args = ap.parse_args()

    if not (SFT / "train_ids.npy").exists():
        sys.exit(f"missing {SFT}/train_ids.npy — run scripts/tokenize_sft.py first")
    if not Path(args.base).exists():
        sys.exit(f"missing base checkpoint {args.base} — pull it from HF first")

    train_ids = np.load(SFT / "train_ids.npy")
    train_mask = np.load(SFT / "train_mask.npy")
    val_ids = np.load(SFT / "val_ids.npy")
    val_mask = np.load(SFT / "val_mask.npy")

    model, _, base_step = load_checkpoint(args.base)
    opt = AdamW(model.parameters(), lr=args.peak_lr, weight_decay=args.weight_decay)
    rng = np.random.default_rng(args.seed)
    val_rng = np.random.default_rng(args.seed + 1)

    CKPT_DIR.mkdir(exist_ok=True)
    print(f"backend={NAME} dtype={DTYPE} params={model.n_params():,} "
          f"base_step={base_step} train_tokens={len(train_ids):,} val_tokens={len(val_ids):,}")
    print(f"steps={args.steps} batch={args.batch_size} block={args.block_size} "
          f"peak_lr={args.peak_lr} warmup={args.warmup}")

    t0 = time.time()
    for step in range(args.steps):
        opt.lr = lr_schedule(step, args.peak_lr, args.warmup, args.steps, args.min_lr)
        x, y = get_sft_batch(train_ids, train_mask, args.block_size, args.batch_size, rng)
        opt.zero_grad()
        loss = cross_entropy(model(x), y)
        loss.backward()
        clip_grad_norm(opt.params, args.max_norm)
        opt.step()
        train_loss = float(loss.data)
        del loss
        gc.collect()

        if (step + 1) % args.log_every == 0:
            it_s = (step + 1) / (time.time() - t0)
            print(f"step {step + 1}/{args.steps} lr={opt.lr:.2e} "
                  f"loss={train_loss:.4f} {it_s:.2f} it/s")
        if (step + 1) % args.eval_every == 0:
            vl, ppl = eval_masked(model, val_ids, val_mask, args.block_size,
                                  args.batch_size, args.eval_batches, val_rng)
            print(f"  [eval] step {step + 1} val_loss={vl:.4f} ppl={ppl:.2f}")
        if (step + 1) % args.ckpt_every == 0:
            save_checkpoint(str(CKPT_DIR / "sft.npz"), model, opt, step + 1)
            print(f"  [ckpt] saved -> {CKPT_DIR / 'sft.npz'} @ step {step + 1}")

    save_checkpoint(str(CKPT_DIR / "sft_final.npz"), model, opt, args.steps)
    print(f"done. final checkpoint -> {CKPT_DIR / 'sft_final.npz'}")


if __name__ == "__main__":
    main()