#!/usr/bin/env python3
"""Distributed pretraining on FineWeb-EDU with checkpoint resume/init support.

Launch on a training machine with all GPUs, for example:

  torchrun --standalone --nproc_per_node=$(nvidia-smi -L | wc -l) train.py \
    --tokenizer_path /path/to/custom-128k-tokenizer \
    --out_dir runs/one_b_fwe \
    --total_tokens 10000000000 \
    --block_size 4096 \
    --micro_batch_size 1 \
    --grad_accum_steps 64

Resume the same run, including optimizer/scaler/RNG state:
  torchrun ... train.py --tokenizer_path ... --out_dir runs/one_b_fwe --resume runs/one_b_fwe/ckpt_last.pt

Start a new run from model weights only:
  torchrun ... train.py --tokenizer_path ... --out_dir runs/continued --init_from runs/one_b_fwe/ckpt_last.pt
"""

from __future__ import annotations

import argparse
import json
import math
import os
import random
import time
from dataclasses import asdict
from pathlib import Path
from typing import Iterable

import numpy as np
import torch
import torch.distributed as dist
from datasets import load_dataset
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, IterableDataset, get_worker_info
from transformers import AutoTokenizer

from model import GPT, GPTConfig


class FineWebEduTokenDataset(IterableDataset):
    """Streams FineWeb-EDU text, tokenizes it, and yields contiguous LM blocks.

    For deterministic DDP resume, keep num_workers=0. The dataset will skip the
    number of per-rank sequences already consumed when --resume is used.
    """

    def __init__(
        self,
        tokenizer_path: str,
        dataset_name: str,
        dataset_config: str | None,
        split: str,
        text_column: str,
        block_size: int,
        seed: int,
        rank: int,
        world_size: int,
        shuffle_buffer: int,
        skip_sequences: int = 0,
        add_eos: bool = True,
    ):
        self.tokenizer_path = tokenizer_path
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.split = split
        self.text_column = text_column
        self.block_size = block_size
        self.seed = seed
        self.rank = rank
        self.world_size = world_size
        self.shuffle_buffer = shuffle_buffer
        self.skip_sequences = skip_sequences
        self.add_eos = add_eos

    def _iter_token_blocks(self) -> Iterable[torch.Tensor]:
        worker = get_worker_info()
        if worker is None:
            shard_rank = self.rank
            num_shards = self.world_size
            worker_seed = self.seed
        else:
            # Additional sharding across DataLoader workers. Exact resume is only
            # guaranteed for num_workers=0, but this prevents duplicate examples.
            shard_rank = self.rank * worker.num_workers + worker.id
            num_shards = self.world_size * worker.num_workers
            worker_seed = self.seed + worker.id

        tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path, use_fast=True)
        if tokenizer.eos_token_id is None:
            raise ValueError("Tokenizer must define an eos_token_id for document boundaries.")

        ds = load_dataset(
            self.dataset_name,
            self.dataset_config,
            split=self.split,
            streaming=True,
        )
        if self.shuffle_buffer > 0:
            ds = ds.shuffle(seed=worker_seed, buffer_size=self.shuffle_buffer)
        ds = ds.shard(num_shards=num_shards, index=shard_rank)

        buf: list[int] = []
        yielded = 0
        for row in ds:
            text = row.get(self.text_column)
            if not text:
                continue
            ids = tokenizer.encode(text, add_special_tokens=False)
            if self.add_eos:
                ids.append(tokenizer.eos_token_id)
            buf.extend(ids)

            while len(buf) >= self.block_size + 1:
                block = buf[: self.block_size + 1]
                del buf[: self.block_size]
                if yielded < self.skip_sequences:
                    yielded += 1
                    continue
                yielded += 1
                yield torch.tensor(block, dtype=torch.long)

    def __iter__(self):
        return self._iter_token_blocks()


def setup_distributed():
    ddp = "RANK" in os.environ and "WORLD_SIZE" in os.environ
    if ddp:
        dist.init_process_group(backend="nccl")
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        master = rank == 0
    else:
        rank = 0
        local_rank = 0
        world_size = 1
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        master = True
    return ddp, rank, local_rank, world_size, device, master


def cleanup_distributed(ddp: bool) -> None:
    if ddp:
        dist.destroy_process_group()


def set_seed(seed: int, rank: int) -> None:
    seed = seed + rank
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def get_lr(tokens_seen: int, args: argparse.Namespace) -> float:
    if tokens_seen < args.warmup_tokens:
        return args.learning_rate * max(1, tokens_seen) / max(1, args.warmup_tokens)
    if tokens_seen >= args.total_tokens:
        return args.min_lr
    progress = (tokens_seen - args.warmup_tokens) / max(1, args.total_tokens - args.warmup_tokens)
    coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
    return args.min_lr + coeff * (args.learning_rate - args.min_lr)


def set_lr(optimizer: torch.optim.Optimizer, lr: float) -> None:
    for group in optimizer.param_groups:
        group["lr"] = lr


def unwrap_model(model: torch.nn.Module) -> GPT:
    return model.module if isinstance(model, DDP) else model


def save_checkpoint(
    path: Path,
    raw_model: GPT,
    optimizer: torch.optim.Optimizer,
    scaler: torch.cuda.amp.GradScaler,
    args: argparse.Namespace,
    tokens_seen: int,
    step: int,
) -> None:
    ckpt = {
        "model": raw_model.state_dict(),
        "model_config": asdict(raw_model.config),
        "optimizer": optimizer.state_dict(),
        "scaler": scaler.state_dict(),
        "tokens_seen": tokens_seen,
        "step": step,
        "args": vars(args),
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state_all": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
        "numpy_rng_state": np.random.get_state(),
        "python_rng_state": random.getstate(),
    }
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(ckpt, tmp)
    os.replace(tmp, path)


def load_checkpoint(path: str, map_location: str | torch.device = "cpu"):
    try:
        return torch.load(path, map_location=map_location, weights_only=False)
    except TypeError:  # older PyTorch
        return torch.load(path, map_location=map_location)


def restore_rng(ckpt) -> None:
    if "torch_rng_state" in ckpt:
        torch.set_rng_state(ckpt["torch_rng_state"])
    if torch.cuda.is_available() and ckpt.get("cuda_rng_state_all") is not None:
        torch.cuda.set_rng_state_all(ckpt["cuda_rng_state_all"])
    if "numpy_rng_state" in ckpt:
        np.random.set_state(ckpt["numpy_rng_state"])
    if "python_rng_state" in ckpt:
        random.setstate(ckpt["python_rng_state"])


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    p.add_argument("--tokenizer_path", type=str, required=True, help="Local path or HF id for your 128k tokenizer")
    p.add_argument("--out_dir", type=str, default="runs/one_b_fineweb_edu")
    p.add_argument("--resume", type=str, default=None, help="Resume full training state from checkpoint")
    p.add_argument("--init_from", type=str, default=None, help="Initialize model weights from checkpoint, new optimizer/run")

    p.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu")
    p.add_argument("--dataset_config", type=str, default="sample-10BT")
    p.add_argument("--split", type=str, default="train")
    p.add_argument("--text_column", type=str, default="text")
    p.add_argument("--shuffle_buffer", type=int, default=10_000)

    p.add_argument("--total_tokens", type=int, default=10_000_000_000)
    p.add_argument("--block_size", type=int, default=4096)
    p.add_argument("--micro_batch_size", type=int, default=1)
    p.add_argument("--grad_accum_steps", type=int, default=64)

    p.add_argument("--n_layer", type=int, default=16)
    p.add_argument("--n_head", type=int, default=16)
    p.add_argument("--n_kv_head", type=int, default=8)
    p.add_argument("--n_embd", type=int, default=2048)
    p.add_argument("--intermediate_size", type=int, default=5632)
    p.add_argument("--dropout", type=float, default=0.0)
    p.add_argument("--rope_theta", type=float, default=10_000.0)
    p.add_argument("--untie_embeddings", action="store_true")

    p.add_argument("--learning_rate", type=float, default=3e-4)
    p.add_argument("--min_lr", type=float, default=3e-5)
    p.add_argument("--warmup_tokens", type=int, default=200_000_000)
    p.add_argument("--weight_decay", type=float, default=0.1)
    p.add_argument("--beta1", type=float, default=0.9)
    p.add_argument("--beta2", type=float, default=0.95)
    p.add_argument("--grad_clip", type=float, default=1.0)

    p.add_argument("--precision", choices=["bf16", "fp16", "fp32"], default="bf16")
    p.add_argument("--compile", action="store_true")
    p.add_argument("--seed", type=int, default=1337)
    p.add_argument("--num_workers", type=int, default=0)
    p.add_argument("--log_interval", type=int, default=10)
    p.add_argument("--save_every_tokens", type=int, default=250_000_000)
    return p


def main() -> None:
    args = build_arg_parser().parse_args()
    ddp, rank, local_rank, world_size, device, master = setup_distributed()
    set_seed(args.seed, rank)

    out_dir = Path(args.out_dir)
    if master:
        out_dir.mkdir(parents=True, exist_ok=True)
        with open(out_dir / "args.json", "w", encoding="utf-8") as f:
            json.dump(vars(args), f, indent=2)

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, use_fast=True)
    vocab_size = len(tokenizer)
    if master and vocab_size != 128_000:
        print(f"[info] tokenizer vocab size is {vocab_size:,}; model will use that size (expected custom 128k).")

    ckpt = None
    tokens_seen = 0
    step = 0
    if args.resume and args.init_from:
        raise ValueError("Use only one of --resume or --init_from.")
    if args.resume:
        ckpt = load_checkpoint(args.resume, map_location="cpu")
        config = GPTConfig(**ckpt["model_config"])
        tokens_seen = int(ckpt.get("tokens_seen", 0))
        step = int(ckpt.get("step", 0))
    elif args.init_from:
        ckpt = load_checkpoint(args.init_from, map_location="cpu")
        config = GPTConfig(**ckpt["model_config"])
        config.block_size = args.block_size
    else:
        config = GPTConfig(
            vocab_size=vocab_size,
            block_size=args.block_size,
            n_layer=args.n_layer,
            n_head=args.n_head,
            n_kv_head=args.n_kv_head,
            n_embd=args.n_embd,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
            rope_theta=args.rope_theta,
            tie_embeddings=not args.untie_embeddings,
        )

    if config.vocab_size != vocab_size:
        raise ValueError(
            f"Checkpoint/model vocab_size={config.vocab_size} but tokenizer has {vocab_size}. "
            "Use the same tokenizer that was used for the checkpoint."
        )

    raw_model = GPT(config)
    if ckpt is not None:
        raw_model.load_state_dict(ckpt["model"], strict=True)
        if master:
            mode = "resumed full state from" if args.resume else "initialized weights from"
            print(f"[info] {mode} {args.resume or args.init_from}")
    raw_model.to(device)

    optimizer = raw_model.configure_optimizers(
        weight_decay=args.weight_decay,
        learning_rate=args.learning_rate,
        betas=(args.beta1, args.beta2),
    )
    model: torch.nn.Module = raw_model
    if args.compile:
        model = torch.compile(model)
    use_fp16_scaler = args.precision == "fp16" and device.type == "cuda"
    scaler = torch.cuda.amp.GradScaler(enabled=use_fp16_scaler)

    if args.resume and ckpt is not None:
        optimizer.load_state_dict(ckpt["optimizer"])
        scaler.load_state_dict(ckpt["scaler"])
        restore_rng(ckpt)

    if ddp:
        model = DDP(model, device_ids=[local_rank], output_device=local_rank, gradient_as_bucket_view=True)

    dtype = {
        "bf16": torch.bfloat16,
        "fp16": torch.float16,
        "fp32": torch.float32,
    }[args.precision]
    use_amp = args.precision != "fp32" and device.type == "cuda"

    tokens_per_iter = args.micro_batch_size * config.block_size * args.grad_accum_steps * world_size
    per_rank_sequences_seen = tokens_seen // max(1, world_size) // config.block_size
    if args.num_workers > 0 and per_rank_sequences_seen > 0 and master:
        print("[warn] exact dataset-position resume is only guaranteed with --num_workers 0")

    dataset = FineWebEduTokenDataset(
        tokenizer_path=args.tokenizer_path,
        dataset_name=args.dataset_name,
        dataset_config=args.dataset_config,
        split=args.split,
        text_column=args.text_column,
        block_size=config.block_size,
        seed=args.seed,
        rank=rank,
        world_size=world_size,
        shuffle_buffer=args.shuffle_buffer,
        skip_sequences=per_rank_sequences_seen,
    )
    loader = DataLoader(
        dataset,
        batch_size=args.micro_batch_size,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
    )
    data_iter = iter(loader)

    if master:
        print(f"[info] device={device}, world_size={world_size}, tokens/iter={tokens_per_iter:,}")
        print(f"[info] parameters={raw_model.num_parameters():,} ({raw_model.num_parameters()/1e9:.3f}B)")
        print(f"[info] starting at step={step:,}, tokens_seen={tokens_seen:,}")
        config.to_json(out_dir / "model_config.json")

    model.train()
    last_log = time.time()
    last_tokens = tokens_seen
    next_save_at = ((tokens_seen // args.save_every_tokens) + 1) * args.save_every_tokens

    while tokens_seen < args.total_tokens:
        lr = get_lr(tokens_seen, args)
        set_lr(optimizer, lr)
        optimizer.zero_grad(set_to_none=True)
        total_loss = 0.0

        for micro_step in range(args.grad_accum_steps):
            batch = next(data_iter).to(device, non_blocking=True)
            x = batch[:, :-1].contiguous()
            y = batch[:, 1:].contiguous()

            sync_context = (
                model.no_sync()
                if ddp and micro_step < args.grad_accum_steps - 1
                else torch.enable_grad()
            )
            with sync_context:
                with torch.autocast(device_type=device.type, dtype=dtype, enabled=use_amp):
                    _, loss = model(x, y)
                    loss = loss / args.grad_accum_steps
                total_loss += loss.detach().float().item()
                scaler.scale(loss).backward()

        if args.grad_clip > 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(raw_model.parameters(), args.grad_clip)
        scaler.step(optimizer)
        scaler.update()

        step += 1
        tokens_seen += tokens_per_iter

        if master and step % args.log_interval == 0:
            now = time.time()
            dt = max(now - last_log, 1e-9)
            toks_per_s = (tokens_seen - last_tokens) / dt
            print(
                f"step {step:7d} | tokens {tokens_seen:13,d}/{args.total_tokens:,} | "
                f"loss {total_loss:.4f} | lr {lr:.2e} | {toks_per_s:,.0f} tok/s"
            )
            last_log = now
            last_tokens = tokens_seen

        if master and (tokens_seen >= next_save_at or tokens_seen >= args.total_tokens):
            save_checkpoint(out_dir / "ckpt_last.pt", raw_model, optimizer, scaler, args, tokens_seen, step)
            save_checkpoint(out_dir / f"ckpt_{tokens_seen:013d}.pt", raw_model, optimizer, scaler, args, tokens_seen, step)
            print(f"[info] saved checkpoint at {tokens_seen:,} tokens")
            next_save_at += args.save_every_tokens

    cleanup_distributed(ddp)


if __name__ == "__main__":
    main()