"""
This training script can be run both on a single gpu in debug mode,
and also in a larger training run with distributed data parallel (ddp).

REQUIRED:
  1. You must specify a config file from the config/ directory
  2. All configuration must be in the config file. No CLI overrides allowed
  3. Each config must specify model_config to select a model from models/

Usage:
  python train.py <config_file>

Examples:
  # Train with Shakespeare char-level model
  python train.py config/train_shakespeare_char.py

  # Train with GPT-2 124M on OpenWebText
  python train.py config/train_gpt2.py

  # Train with Reflow model
  python train.py config/train_reflow.py

  # DDP on 4 gpus:
  torchrun --standalone --nproc_per_node=4 train.py config/train_gpt2.py

Available configs in config/:
  - train_gpt2.py           GPT-2 124M on OpenWebText
  - train_shakespeare_char.py  Character-level Shakespeare
  - train_reflow.py         Reflow model (modernized GPT)
  - train_sft.py            SFT with GPT-2
  - train_sft_lima.py       SFT with LIMA dataset
  - finetune_shakespeare.py    Fine-tune GPT-2-XL on Shakespeare
"""

# -----------------------------------------------------------------------------
# Configuration loading (BEFORE imports to validate config first)
# Usage:
#   python train.py config/train_gpt2.py
# Note: All configuration must be specified in the config file. No CLI overrides.
# -----------------------------------------------------------------------------
import sys
import os

# Parse command line - only accept config file, no --key=value allowed
if len(sys.argv) != 2:
    print("ERROR: Invalid arguments!")
    print("Usage: python train.py <config_file>")
    print("Note: All configuration must be in the config file. No CLI overrides.")
    print("Available configs in config/:")
    print("  - train_gpt2.py              (GPT-2 124M)")
    print("  - train_shakespeare_char.py  (Shakespeare char-level)")
    print("  - train_reflow.py           (Reflow model)")
    print("  - finetune_shakespeare.py   (Fine-tune GPT-2-XL)")
    sys.exit(1)

config_file = sys.argv[1]

# Disallow --key=value arguments
for arg in sys.argv[1:]:
    if arg.startswith('--'):
        print(f"ERROR: CLI overrides are not supported. All config must be in file: {config_file}")
        sys.exit(1)

# Load the specified config file
print(f"Loading config from: {config_file}")
exec(open(config_file).read())

# Validate required config keys
required_keys = [
    'out_dir', 'dataset', 'batch_size', 'block_size',
    'n_layer', 'n_head', 'n_embd', 'learning_rate', 'max_iters',
    'model_config'  # Must specify which model file to use
]

# Optional log file path (e.g., '/path/to/train.log')
log_file = globals().get('log_file', None)
log_f = None
if log_file:
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    log_f = open(log_file, 'a')  # append mode
    print(f"Logging to: {log_file}")
missing_keys = [k for k in required_keys if k not in globals()]
if missing_keys:
    print(f"ERROR: Missing required config keys: {missing_keys}")
    sys.exit(1)

# Load model configuration - import GPTConfig and GPT from specified model file
model_config = globals()['model_config']
model_file = f"models/{model_config}.py"
try:
    exec(open(model_file).read())
except FileNotFoundError:
    print(f"ERROR: Model file not found: {model_file}")
    print(f"Available models in models/:")
    import os
    for f in os.listdir('models'):
        if f.endswith('.py') and not f.startswith('_'):
            print(f"  - {f[:-3]}")
    sys.exit(1)

# Get model-specific required config keys from GPTConfig
model_required_keys = []
if 'GPTConfig' in globals():
    config_class = globals()['GPTConfig']
    # Get all fields from the dataclass
    import dataclasses
    for field in dataclasses.fields(config_class):
        model_required_keys.append(field.name)

# Validate model-specific config keys
missing_model_keys = [k for k in model_required_keys if k not in globals()]
if missing_model_keys:
    print(f"ERROR: Missing required model config keys for {model_config}: {missing_model_keys}")
    print(f"Required keys: {model_required_keys}")
    sys.exit(1)

# Validate SFT-specific config
if init_from == 'finetune':
    if 'base_model_dir' not in globals():
        print("ERROR: 'base_model_dir' is required when init_from='finetune'")
        sys.exit(1)

# Freeze config - collect all config keys and make them immutable
config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
config = {k: globals()[k] for k in config_keys}

# Print all configuration
print("\n" + "=" * 60)
print("CURRENT CONFIGURATION")
print("=" * 60)
for key in sorted(config.keys()):
    print(f"  {key:30s} = {config[key]}")
print("=" * 60 + "\n")

# Helper function for logging to both stdout and file
def log_print(*args, **kwargs):
    print(*args, **kwargs)
    if log_f:
        print(*args, **kwargs, file=log_f)
        log_f.flush()

# Store frozen config
_frozen_config = config.copy()
# -----------------------------------------------------------------------------

# Now import dependencies
import os
import time
import math
import pickle
from contextlib import nullcontext

import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group

# Import GPTConfig and GPT from the model file specified in model_config
GPTConfig = globals()['GPTConfig']
GPT = globals()['GPT']

# Auto-detect dtype
if dtype == 'bfloat16' and not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()):
    dtype = 'float16'
    # Update the frozen config as well
    _frozen_config['dtype'] = dtype

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
if ddp:
    init_process_group(backend=backend)
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    ddp_world_size = int(os.environ['WORLD_SIZE'])
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
    seed_offset = ddp_rank # each process gets a different seed
    # world_size number of processes will be training simultaneously, so we can scale
    # down the desired gradient accumulation iterations per process proportionally
    assert gradient_accumulation_steps % ddp_world_size == 0
    gradient_accumulation_steps //= ddp_world_size
else:
    # if not ddp, we are running on a single gpu, and one process
    master_process = True
    seed_offset = 0
    ddp_world_size = 1
tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size
log_print(f"tokens per iteration will be: {tokens_per_iter:,}")

if master_process:
    os.makedirs(out_dir, exist_ok=True)
torch.manual_seed(1337 + seed_offset)
torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
# note: float16 data type will automatically use a GradScaler
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

# poor man's data loader
data_dir = os.path.join('data', dataset)

# Check if SFT masking is enabled
sft_masking = globals().get('sft_masking', False)
sft_mask_available = False

# Check if SFT mask files exist (dual-bin format: train.bin + train_mask.bin)
if sft_masking:
    train_mask_path = os.path.join(data_dir, 'train_mask.bin')
    val_mask_path = os.path.join(data_dir, 'val_mask.bin')
    sft_mask_available = os.path.exists(train_mask_path) and os.path.exists(val_mask_path)
    if not sft_mask_available:
        log_print(f"WARNING: sft_masking=True but mask files not found at {data_dir}/train_mask.bin")
        sft_masking = False

def get_batch(split):
    # We recreate np.memmap every batch to avoid a memory leak, as per
    # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122
    if split == 'train':
        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
    else:
        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])

    # SFT masking: only compute loss on response part
    # If sft_mask_available, load mask data and apply masking
    if sft_mask_available:
        if split == 'train':
            mask_data = np.memmap(os.path.join(data_dir, 'train_mask.bin'), dtype=np.uint16, mode='r')
        else:
            mask_data = np.memmap(os.path.join(data_dir, 'val_mask.bin'), dtype=np.uint16, mode='r')
        # mask_data[i+1] corresponds to y[i], so we shift by 1
        m = torch.stack([torch.from_numpy((mask_data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
        # Apply mask: set y to -100 where mask is 0 (ignore these tokens in loss)
        y = torch.where(m == 0, torch.full_like(y, -100), y)

    if device_type == 'cuda':
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

# init these up here, can override if init_from='resume' (i.e. from a checkpoint)
iter_num = 0
best_val_loss = 1e9

# attempt to derive vocab_size from the dataset
meta_path = os.path.join(data_dir, 'meta.pkl')
meta_vocab_size = None
if os.path.exists(meta_path):
    with open(meta_path, 'rb') as f:
        meta = pickle.load(f)
    meta_vocab_size = meta['vocab_size']
    log_print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")

# model init - collect all model-specific config from GPTConfig
model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
                  bias=bias, vocab_size=None, dropout=dropout)

# Add any additional model-specific parameters from GPTConfig
if 'GPTConfig' in globals():
    config_class = globals()['GPTConfig']
    import dataclasses
    for field in dataclasses.fields(config_class):
        if field.name not in model_args and field.name in globals():
            model_args[field.name] = globals()[field.name]
if init_from == 'scratch':
    # init a new model from scratch
    print("Initializing a new model from scratch")
    # determine the vocab size we'll use for from-scratch training
    if meta_vocab_size is None:
        print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)")
    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
elif init_from == 'resume':
    log_print(f"Resuming training from {out_dir}")
    # resume training from a checkpoint.
    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
    checkpoint = torch.load(ckpt_path, map_location=device)
    checkpoint_model_args = checkpoint['model_args']
    # force these config attributes to be equal otherwise we can't even resume training
    # the rest of the attributes (e.g. dropout) can stay as desired from command line
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = checkpoint_model_args[k]
    # create the model
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
    state_dict = checkpoint['model']
    # fix the keys of the state dictionary :(
    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    model.load_state_dict(state_dict)
    iter_num = checkpoint['iter_num']
    best_val_loss = checkpoint['best_val_loss']
elif init_from == 'finetune':
    # Fine-tuning from a pretrained base model (load weights but not optimizer)
    log_print(f"Fine-tuning from base model in {base_model_dir}")
    ckpt_path = os.path.join(base_model_dir, 'ckpt.pt')
    checkpoint = torch.load(ckpt_path, map_location=device)
    checkpoint_model_args = checkpoint['model_args']
    # force these config attributes to be equal
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = checkpoint_model_args[k]
    # create the model
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    model.load_state_dict(state_dict)
    # Do NOT load optimizer or iter_num - fresh SFT training
    checkpoint = None
elif init_from.startswith('gpt2'):
    log_print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
    # initialize from OpenAI GPT-2 weights
    override_args = dict(dropout=dropout)
    model = GPT.from_pretrained(init_from, override_args)
    # read off the created config params, so we can store them into checkpoint correctly
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = getattr(model.config, k)
# crop down the model block size if desired, using model surgery
if block_size < model.config.block_size:
    model.crop_block_size(block_size)
    model_args['block_size'] = block_size # so that the checkpoint will have the right value
model.to(device)

# initialize a GradScaler. If enabled=False scaler is a no-op
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))

# optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
if init_from == 'resume':
    optimizer.load_state_dict(checkpoint['optimizer'])
checkpoint = None # free up memory

# compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    unoptimized_model = model
    model = torch.compile(model) # requires PyTorch 2.0

# wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])

# helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            with ctx:
                logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * (it + 1) / (warmup_iters + 1)
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
    return min_lr + coeff * (learning_rate - min_lr)

# logging
if wandb_log and master_process:
    import wandb
    wandb.init(project=wandb_project, name=wandb_run_name, config=config)

# training loop
X, Y = get_batch('train') # fetch the very first batch
t0 = time.time()
local_iter_num = 0 # number of iterations in the lifetime of this process
raw_model = model.module if ddp else model # unwrap DDP container if needed
running_mfu = -1.0
while True:

    # determine and set the learning rate for this iteration
    lr = get_lr(iter_num) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # evaluate the loss on train/val sets and write checkpoints
    if iter_num % eval_interval == 0 and master_process:
        losses = estimate_loss()
        log_print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
        if wandb_log:
            wandb.log({
                "iter": iter_num,
                "train/loss": losses['train'],
                "val/loss": losses['val'],
                "lr": lr,
                "mfu": running_mfu*100, # convert to percentage
            })
        if losses['val'] < best_val_loss or always_save_checkpoint:
            best_val_loss = losses['val']
            if iter_num > 0:
                checkpoint = {
                    'model': raw_model.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'model_args': model_args,
                    'iter_num': iter_num,
                    'best_val_loss': best_val_loss,
                    'config': config,
                }
                log_print(f"saving checkpoint to {out_dir}")
                torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
    if iter_num == 0 and eval_only:
        break

    # forward backward update, with optional gradient accumulation to simulate larger batch size
    # and using the GradScaler if data type is float16
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            # in DDP training we only need to sync gradients at the last micro step.
            # the official way to do this is with model.no_sync() context manager, but
            # I really dislike that this bloats the code and forces us to repeat code
            # looking at the source of that context manager, it just toggles this variable
            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
        with ctx:
            logits, loss = model(X, Y)
            loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation
        # immediately async prefetch next batch while model is doing the forward pass on the GPU
        X, Y = get_batch('train')
        # backward pass, with gradient scaling if training in fp16
        scaler.scale(loss).backward()
    # clip the gradient
    if grad_clip != 0.0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    # step the optimizer and scaler if training in fp16
    scaler.step(optimizer)
    scaler.update()
    # flush the gradients as soon as we can, no need for this memory anymore
    optimizer.zero_grad(set_to_none=True)

    # timing and logging
    t1 = time.time()
    dt = t1 - t0
    t0 = t1
    if iter_num % log_interval == 0 and master_process:
        # get loss as float. note: this is a CPU-GPU sync point
        # scale up to undo the division above, approximating the true total loss (exact would have been a sum)
        lossf = loss.item() * gradient_accumulation_steps
        if local_iter_num >= 5: # let the training loop settle a bit
            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
            running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
        log_print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%")
    iter_num += 1
    local_iter_num += 1

    # termination conditions
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()

# Close log file
if log_f:
    log_f.close()