"""
Distributed training script for 1B parameter Transformer.

Launch: torchrun --nproc_per_node=8 train.py

Stack: PyTorch DDP + BF16 autocast + 8x H100 80GB
"""

import os
import sys
import math
import time
import json
import datetime

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model.config import ModelConfig, TrainConfig
from model.transformer import Transformer
from model.data import get_tokenizer, create_dataloader


def get_wsd_lr(step, warmup_steps, total_steps, max_lr, min_lr):
    """Warmup-Stable-Decay: linear warmup -> constant -> cosine decay (last 20%)."""
    stable_end = int(total_steps * 0.8)
    if step < warmup_steps:
        return max_lr * step / max(warmup_steps, 1)
    elif step < stable_end:
        return max_lr
    else:
        progress = (step - stable_end) / max(total_steps - stable_end, 1)
        return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))


def find_latest_checkpoint(checkpoint_dir):
    """Find the latest step_*.pt checkpoint in the directory."""
    import glob
    pattern = os.path.join(checkpoint_dir, "step_*.pt")
    files = glob.glob(pattern)
    if not files:
        return None, 0
    latest = max(files, key=lambda f: int(os.path.basename(f).replace("step_", "").replace(".pt", "")))
    step = int(os.path.basename(latest).replace("step_", "").replace(".pt", ""))
    return latest, step


def main():
    dist.init_process_group("nccl", timeout=datetime.timedelta(minutes=30))
    rank = int(os.environ.get("RANK", 0))
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    model_config = ModelConfig()
    train_config = TrainConfig()

    eff_batch = train_config.batch_size_per_gpu * world_size * train_config.gradient_accumulation_steps
    tokens_per_step = eff_batch * model_config.max_seq_len
    total_steps = train_config.total_tokens // tokens_per_step

    if rank == 0:
        os.makedirs(train_config.log_dir, exist_ok=True)
        os.makedirs(train_config.checkpoint_dir, exist_ok=True)
        print("=" * 70)
        print(f"  TRAINING 1B TRANSFORMER FROM SCRATCH")
        print(f"  Arch: {model_config.num_layers}L / {model_config.hidden_dim}D / "
              f"{model_config.num_attention_heads}H / GQA-{model_config.num_kv_heads}KV / "
              f"SwiGLU-{model_config.intermediate_dim}")
        print(f"  Seq: {model_config.max_seq_len} | Vocab: {model_config.vocab_size}")
        print(f"  GPUs: {world_size}x H100 80GB | Backend: DDP + BF16 autocast")
        print(f"  Batch: {eff_batch} seqs = {tokens_per_step:,} tok/step")
        print(f"  Steps: {total_steps:,} | Target: {train_config.total_tokens:,} tokens")
        print("=" * 70)

    # Tokenizer
    tokenizer = get_tokenizer()

    # Model
    torch.manual_seed(train_config.seed)
    model = Transformer(model_config).to(device)

    if rank == 0:
        n = sum(p.numel() for p in model.parameters())
        print(f"[Init] Params: {n:,} ({n/1e9:.3f}B)")

    model = DDP(model, device_ids=[local_rank])

    # Optimizer
    decay_params = [p for n, p in model.named_parameters() if p.dim() >= 2 and p.requires_grad]
    nodecay_params = [p for n, p in model.named_parameters() if p.dim() < 2 and p.requires_grad]
    optimizer = torch.optim.AdamW([
        {"params": decay_params, "weight_decay": train_config.weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ], lr=train_config.learning_rate, betas=(train_config.beta1, train_config.beta2), fused=True)

    if rank == 0:
        dp = sum(p.numel() for p in decay_params)
        ndp = sum(p.numel() for p in nodecay_params)
        print(f"[Init] Optimizer: {dp:,} decay + {ndp:,} no-decay params")

    # Resume from checkpoint
    resume_step = 0
    ckpt_path, ckpt_step = find_latest_checkpoint(train_config.checkpoint_dir)
    if ckpt_path is not None:
        if rank == 0:
            print(f"[Resume] Loading checkpoint: {ckpt_path} (step {ckpt_step})")
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        model.module.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        resume_step = ckpt["step"]
        if rank == 0:
            print(f"[Resume] Restored model + optimizer at step {resume_step}, "
                  f"loss was {ckpt.get('loss', 'N/A')}")
        del ckpt
        torch.cuda.empty_cache()
    else:
        if rank == 0:
            print("[Init] No checkpoint found, starting from scratch")

    # Data — use (seed + resume_step) so resumed runs see different shuffled data
    effective_seed = train_config.seed + resume_step
    dataloader = create_dataloader(tokenizer, train_config, rank=rank, world_size=world_size,
                                   seed_override=effective_seed)
    data_iter = iter(dataloader)

    if rank == 0:
        print(f"[Init] Dataloader ready (streaming FineWeb-Edu 10BT)")
        print(f"[Schedule] WSD: warmup {train_config.warmup_steps} -> "
              f"stable {int(total_steps*0.8)} -> decay {total_steps}")
        if resume_step > 0:
            remaining = total_steps - resume_step
            print(f"[Resume] Continuing from step {resume_step}, {remaining:,} steps remaining")
        print("-" * 70)
        sys.stdout.flush()

    # ===== TRAINING LOOP =====
    model.train()
    global_step = resume_step
    running_loss = 0.0
    best_loss = float("inf")
    tokens_done = resume_step * tokens_per_step
    t0 = time.time()
    step_t0 = time.time()

    log_file = open(os.path.join(train_config.log_dir, "train_log.jsonl"), "a") if rank == 0 else None

    while global_step < total_steps:
        optimizer.zero_grad(set_to_none=True)
        micro_loss = 0.0

        for micro in range(train_config.gradient_accumulation_steps):
            try:
                input_ids, labels = next(data_iter)
            except StopIteration:
                data_iter = iter(dataloader)
                input_ids, labels = next(data_iter)

            input_ids = input_ids.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            # BF16 autocast — no scaler needed (BF16 has enough dynamic range)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                _, loss = model(input_ids, labels)
                loss = loss / train_config.gradient_accumulation_steps

            loss.backward()
            micro_loss += loss.item()

        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.grad_clip)

        # LR schedule
        lr = get_wsd_lr(global_step, train_config.warmup_steps, total_steps,
                        train_config.learning_rate, train_config.min_lr)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        optimizer.step()
        global_step += 1
        running_loss += micro_loss
        tokens_done += tokens_per_step

        # Log
        if global_step % train_config.log_interval == 0:
            dt = time.time() - step_t0
            tps = (train_config.log_interval * tokens_per_step) / max(dt, 1e-9)
            avg = running_loss / train_config.log_interval
            elapsed = time.time() - t0
            pct = 100.0 * global_step / total_steps
            eta = (elapsed / max(global_step, 1)) * (total_steps - global_step)

            if rank == 0:
                gpu_mem = torch.cuda.max_memory_allocated(device) / 1e9
                print(
                    f"[Step {global_step:>6d}/{total_steps}] "
                    f"loss={avg:.4f} | lr={lr:.2e} | "
                    f"tok/s={tps:,.0f} | GPU={gpu_mem:.1f}GB | "
                    f"{pct:.1f}% | ETA={eta/3600:.1f}h",
                    flush=True,
                )
                if log_file:
                    log_file.write(json.dumps({
                        "step": global_step, "loss": round(avg, 4), "lr": lr,
                        "tps": round(tps), "tokens": tokens_done,
                        "gpu_gb": round(gpu_mem, 1), "elapsed_s": round(elapsed, 1),
                    }) + "\n")
                    log_file.flush()

            if avg < best_loss:
                best_loss = avg
            running_loss = 0.0
            step_t0 = time.time()

        # Checkpoint
        if global_step % train_config.save_interval == 0:
            dist.barrier()
            if rank == 0:
                ckpt_path = os.path.join(train_config.checkpoint_dir, f"step_{global_step}.pt")
                torch.save({
                    "step": global_step,
                    "model": model.module.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "loss": avg if global_step % train_config.log_interval == 0 else micro_loss,
                    "config": {"model": model_config.__dict__, "train": train_config.__dict__},
                }, ckpt_path)
                print(f"  >> Checkpoint: {ckpt_path}", flush=True)
            dist.barrier()

    # Final
    dist.barrier()
    if rank == 0:
        final_path = os.path.join(train_config.checkpoint_dir, "final.pt")
        torch.save({
            "step": global_step,
            "model": model.module.state_dict(),
            "config": {"model": model_config.__dict__, "train": train_config.__dict__},
        }, final_path)
        total_time = time.time() - t0
        print("=" * 70)
        print(f"  TRAINING COMPLETE")
        print(f"  Steps: {global_step:,} | Tokens: {tokens_done:,}")
        print(f"  Time: {total_time/3600:.2f}h | Throughput: {tokens_done/total_time:,.0f} tok/s")
        print(f"  Best loss: {best_loss:.4f}")
        print(f"  Final model: {final_path}")
        print("=" * 70)
        if log_file:
            log_file.close()

    dist.destroy_process_group()


if __name__ == "__main__":
    main()