#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["numpy>=1.26", "torch>=2.4", "tokenizers>=0.20"]
# ///
"""Train the dense ~50M GPT after running `uv run dataset.py`.

Run: uv run train.py
Logs are emitted every 100 optimizer steps by default.
"""

from __future__ import annotations

import argparse
import json
import math
import random
import shutil
import time
from dataclasses import asdict
from pathlib import Path

import numpy as np
import torch
from tokenizers import Tokenizer
from torch.utils.data import DataLoader, Dataset

from config import CHECKPOINT_DIR, DATA_METADATA_PATH, MODEL, TOKENS_PATH, TRAIN, TOKENIZER_PATH
from model import GPT, parameter_count


class TokenBlocks(Dataset):
    def __init__(self, path, num_tokens: int, block_size: int, seed: int):
        self.path, self.num_tokens, self.block_size = str(path), num_tokens, block_size
        self.n_blocks = (num_tokens - 1) // block_size
        self.order = np.random.default_rng(seed).permutation(self.n_blocks)
        self.memmap = None

    def __len__(self):
        return self.n_blocks

    def __getitem__(self, index):
        if self.memmap is None:
            self.memmap = np.memmap(self.path, mode="r", dtype=np.uint16, shape=(self.num_tokens,))
        start = int(self.order[index]) * self.block_size
        return torch.from_numpy(np.asarray(self.memmap[start : start + self.block_size + 1], dtype=np.int64))


def make_optimizer(model: GPT):
    decay, no_decay = [], []
    for parameter in model.parameters():
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)
    groups = [{"params": decay, "weight_decay": TRAIN.weight_decay}, {"params": no_decay, "weight_decay": 0.0}]
    try:
        return torch.optim.AdamW(groups, lr=TRAIN.learning_rate, betas=(0.9, 0.95), eps=1e-8, fused=True)
    except (RuntimeError, TypeError):
        return torch.optim.AdamW(groups, lr=TRAIN.learning_rate, betas=(0.9, 0.95), eps=1e-8)


def lr_scheduler(optimizer):
    warmup = max(1, int(TRAIN.max_steps * TRAIN.warmup_ratio))

    def scale(step: int):
        if step < warmup:
            return (step + 1) / warmup
        progress = min(1.0, (step - warmup) / max(1, TRAIN.max_steps - warmup))
        return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, scale)


def format_eta(seconds: float) -> str:
    total_seconds = int(max(0, seconds))
    days, remainder = divmod(total_seconds, 86400)
    hours, remainder = divmod(remainder, 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{days:02d}d:{hours:02d}h:{minutes:02d}m:{secs:02d}s"


def find_checkpoint(resume_arg: str | bool) -> Path:
    if not CHECKPOINT_DIR.exists():
        raise FileNotFoundError(f"Checkpoint directory {CHECKPOINT_DIR} does not exist.")

    if isinstance(resume_arg, str):
        candidate = Path(resume_arg)
        if candidate.is_dir() and (candidate / "model.pt").exists():
            return candidate
        if (CHECKPOINT_DIR / resume_arg).is_dir() and ((CHECKPOINT_DIR / resume_arg) / "model.pt").exists():
            return CHECKPOINT_DIR / resume_arg
        try:
            step_num = int(resume_arg)
            formatted_dir = CHECKPOINT_DIR / f"step-{step_num:07d}"
            if formatted_dir.is_dir() and (formatted_dir / "model.pt").exists():
                return formatted_dir
        except ValueError:
            pass
        raise FileNotFoundError(f"Specified checkpoint '{resume_arg}' not found.")

    checkpoints = []
    for p in CHECKPOINT_DIR.glob("step-*"):
        if p.is_dir() and (p / "model.pt").exists():
            try:
                step_num = int(p.name.split("-")[1])
                checkpoints.append((step_num, p))
            except (IndexError, ValueError):
                pass

    if not checkpoints:
        raise FileNotFoundError(f"No valid checkpoints found in {CHECKPOINT_DIR}.")

    checkpoints.sort(key=lambda x: x[0])
    return checkpoints[-1][1]


def save_checkpoint(model, optimizer, scheduler, step: int, scaler=None):
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    final = CHECKPOINT_DIR / f"step-{step:07d}"
    temporary = CHECKPOINT_DIR / f".step-{step:07d}.tmp"
    shutil.rmtree(temporary, ignore_errors=True)
    temporary.mkdir()
    ckpt_dict = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "step": step,
    }
    if scaler is not None:
        ckpt_dict["scaler"] = scaler.state_dict()
    torch.save(ckpt_dict, temporary / "model.pt")
    (temporary / "config.json").write_text(json.dumps(asdict(MODEL), indent=2), encoding="utf-8")
    shutil.copy2(TOKENIZER_PATH, temporary / "tokenizer.json")
    if final.exists():
        shutil.rmtree(final)
    temporary.rename(final)
    for old in sorted(CHECKPOINT_DIR.glob("step-*"))[:-2]:
        shutil.rmtree(old, ignore_errors=True)
    print(f"[+] checkpoint saved: {final}")


def main():
    parser = argparse.ArgumentParser(description="Train the dense ~50M GPT")
    parser.add_argument(
        "--resume",
        nargs="?",
        const=True,
        default=False,
        help="Resume training from the latest checkpoint (or specify checkpoint path/step)",
    )
    args = parser.parse_args()

    if not TOKENS_PATH.exists() or not TOKENIZER_PATH.exists() or not DATA_METADATA_PATH.exists():
        raise FileNotFoundError("Dataset artifacts are missing. Run `uv run dataset.py` first.")
    if not torch.cuda.is_available():
        raise RuntimeError("No CUDA GPU detected. This configuration targets your RTX 3050 6GB.")

    torch.manual_seed(TRAIN.seed)
    np.random.seed(TRAIN.seed)
    random.seed(TRAIN.seed)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.set_float32_matmul_precision("high")
    device = torch.device("cuda")
    print(f"[*] GPU: {torch.cuda.get_device_name(0)} | VRAM: {torch.cuda.get_device_properties(0).total_memory / 2**30:.1f} GiB")

    metadata = json.loads(DATA_METADATA_PATH.read_text(encoding="utf-8"))
    tokenizer = Tokenizer.from_file(str(TOKENIZER_PATH))
    if tokenizer.get_vocab_size() != MODEL.vocab_size:
        raise ValueError("Tokenizer/model vocabulary mismatch. Delete .data and run `uv run dataset.py` again.")

    model = GPT(MODEL).to(device=device)
    model.gradient_checkpointing = True
    parameters = parameter_count(model)
    print(f"[*] Dense model parameters: {parameters:,} ({parameters / 1e6:.2f}M)")
    print(f"[*] Sequence: {MODEL.block_size} | micro-batch: {TRAIN.micro_batch_size} | accumulation: {TRAIN.gradient_accumulation} | effective tokens/update: {MODEL.block_size * TRAIN.micro_batch_size * TRAIN.gradient_accumulation:,}")

    optimizer = make_optimizer(model)
    scheduler = lr_scheduler(optimizer)
    scaler = torch.amp.GradScaler("cuda")

    start_step = 0
    if args.resume:
        ckpt_path = find_checkpoint(args.resume)
        print(f"[*] Resuming training from checkpoint: {ckpt_path}")
        checkpoint = torch.load(ckpt_path / "model.pt", map_location=device)
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        scheduler.load_state_dict(checkpoint["scheduler"])
        if "scaler" in checkpoint:
            scaler.load_state_dict(checkpoint["scaler"])
        start_step = checkpoint.get("step", 0)
        print(f"[*] Resumed at step {start_step:,}/{TRAIN.max_steps:,}")

    if start_step >= TRAIN.max_steps:
        print(f"[*] Training already completed ({start_step}/{TRAIN.max_steps} steps). Nothing to do.")
        return

    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    log_file = CHECKPOINT_DIR / "train.log"

    dataset = TokenBlocks(TOKENS_PATH, int(metadata["tokens"]), MODEL.block_size, TRAIN.seed)
    total_samples = (len(dataset) // TRAIN.micro_batch_size) * TRAIN.micro_batch_size
    consumed_samples = start_step * TRAIN.gradient_accumulation * TRAIN.micro_batch_size
    start_idx = consumed_samples % total_samples

    full_loader = DataLoader(dataset, batch_size=TRAIN.micro_batch_size, drop_last=True, num_workers=0, pin_memory=True)

    if start_idx > 0:
        sampler = range(start_idx, total_samples)
        first_loader = DataLoader(dataset, batch_size=TRAIN.micro_batch_size, sampler=sampler, pin_memory=True)
        iterator = iter(first_loader)
        print(f"[*] Fast-forwarded dataset to sample index {start_idx:,}/{total_samples:,} (batch {consumed_samples // TRAIN.micro_batch_size:,})")
    else:
        iterator = iter(full_loader)

    optimizer.zero_grad(set_to_none=True)
    model.train()
    print(f"[*] Training for steps {start_step + 1:,} -> {TRAIN.max_steps:,}; logging every {TRAIN.log_every} steps to {log_file}")

    terminal_log_every = 10
    file_log_every = TRAIN.log_every  # 100 steps

    start_train_time = time.time()
    started_term = time.time()
    started_file = time.time()
    loss_sum_term = 0.0
    loss_sum_file = 0.0
    term_steps_count = 0
    file_steps_count = 0

    for step in range(start_step + 1, TRAIN.max_steps + 1):
        step_loss = 0.0
        for _ in range(TRAIN.gradient_accumulation):
            try:
                block = next(iterator)
            except StopIteration:
                iterator = iter(full_loader)
                block = next(iterator)
            block = block.to(device, non_blocking=True)
            with torch.autocast("cuda", dtype=torch.float16):
                _, loss = model(block[:, :-1], block[:, 1:])
                loss = loss / TRAIN.gradient_accumulation
            scaler.scale(loss).backward()
            step_loss += loss.detach().float().item()

        loss_sum_term += step_loss
        loss_sum_file += step_loss
        term_steps_count += 1
        file_steps_count += 1

        scaler.unscale_(optimizer)
        gradient_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)
        scheduler.step()

        now = time.time()
        overall_elapsed = max(now - start_train_time, 1e-6)
        steps_done_this_run = step - start_step
        sec_per_step = overall_elapsed / steps_done_this_run
        eta_seconds = (TRAIN.max_steps - step) * sec_per_step
        eta_str = format_eta(eta_seconds)

        # Print to terminal every 10 steps
        if step % terminal_log_every == 0 or step == TRAIN.max_steps:
            elapsed_term = max(now - started_term, 1e-6)
            avg_loss_term = loss_sum_term / term_steps_count
            throughput_term = term_steps_count * TRAIN.gradient_accumulation * TRAIN.micro_batch_size * MODEL.block_size / elapsed_term

            term_msg = (
                f"[step {step:,}/{TRAIN.max_steps:,}] "
                f"loss={avg_loss_term:.4f} "
                f"ppl={math.exp(min(avg_loss_term, 20)):.2f} "
                f"lr={scheduler.get_last_lr()[0]:.3e} "
                f"grad={float(gradient_norm):.3f} "
                f"speed={throughput_term:,.0f} tok/s "
                f"eta={eta_str}"
            )
            print(term_msg)
            started_term, loss_sum_term, term_steps_count = time.time(), 0.0, 0

        # Log to file every 100 steps
        if step % file_log_every == 0 or step == TRAIN.max_steps:
            elapsed_file = max(now - started_file, 1e-6)
            avg_loss_file = loss_sum_file / file_steps_count
            throughput_file = file_steps_count * TRAIN.gradient_accumulation * TRAIN.micro_batch_size * MODEL.block_size / elapsed_file

            file_msg = (
                f"[step {step:,}/{TRAIN.max_steps:,}] "
                f"loss={avg_loss_file:.4f} "
                f"ppl={math.exp(min(avg_loss_file, 20)):.2f} "
                f"lr={scheduler.get_last_lr()[0]:.3e} "
                f"grad={float(gradient_norm):.3f} "
                f"speed={throughput_file:,.0f} tok/s "
                f"eta={eta_str}"
            )
            with log_file.open("a", encoding="utf-8") as f:
                f.write(file_msg + "\n")
            started_file, loss_sum_file, file_steps_count = time.time(), 0.0, 0

        if step % TRAIN.save_every == 0 or step == TRAIN.max_steps:
            save_checkpoint(model, optimizer, scheduler, step, scaler)


if __name__ == "__main__":
    main()