import os
import sys
import argparse
import time
import math
import torch
from torch.utils.data import DataLoader

# Ensure the module can find our files
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from model import RecursiveCausalLM, ModelConfig
from dataset import MemmappedDataset

def get_lr_scheduler(step, total_steps, warmup_steps, peak_lr, min_lr):
    """Calculates cosine learning rate decay with a linear warmup phase."""
    if step < warmup_steps:
        # Linear warmup
        return peak_lr * (step + 1) / warmup_steps
    if step >= total_steps or total_steps <= warmup_steps:
        return min_lr
    # Cosine decay down to min_lr
    decay_ratio = (step - warmup_steps) / (total_steps - warmup_steps)
    coefficient = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coefficient * (peak_lr - min_lr)

@torch.no_grad()
def run_live_sample(model, tokenizer, device, prompt="Question: A train of length 150 meters passes a pole in 15 seconds. What is the speed of", max_new_tokens=50):
    """Pauses training to run a quick evaluation prompt on the active model state."""
    was_training = model.training
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    generated = list(input_ids[0].cpu().numpy())
    
    use_amp = (device.type == "cuda")
    for _ in range(max_new_tokens):
        curr_input = torch.tensor([generated[-model.config.max_seq_len:]], dtype=torch.long, device=device)
        with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
            logits, _ = model(curr_input)
        next_token_logits = logits[0, -1, :] / 0.7  # Temperature = 0.7
        
        # Simple Top-K filtering (k=50)
        v, _ = torch.topk(next_token_logits, min(50, next_token_logits.size(-1)))
        next_token_logits[next_token_logits < v[-1]] = -float('inf')
        
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1).item()
        generated.append(next_token)
        
        if next_token == tokenizer.eos_token_id:
            break
            
    decoded = tokenizer.decode(generated)
    print(f"\n[LIVE GENERATION CALLBACK]:\n   \"{decoded}\"\n")
    if was_training:
        model.train()

def save_checkpoint(model, optimizer, step, val_loss, config, base_path, max_to_keep=3):
    """Saves an atomic checkpoint and maintains a rotating list of the last N checkpoints on disk."""
    import glob
    import re
    import shutil
    
    dir_name = os.path.dirname(base_path)
    if dir_name:
        os.makedirs(dir_name, exist_ok=True)
    
    base_name = os.path.basename(base_path)
    name_no_ext, ext = os.path.splitext(base_name)
    
    # Strip any existing '_step_\d+' suffix to prevent compounding step patterns
    name_no_ext = re.sub(r"_step_\d+", "", name_no_ext)
    
    # Step-specific filename (e.g., micro_llm_200m/uct_pretrain_step_100.pt)
    step_path = os.path.join(dir_name, f"{name_no_ext}_step_{step}{ext}")
    
    # Save step checkpoint atomically
    tmp_path = step_path + ".tmp"
    torch.save({
        "step": step,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "val_loss": val_loss,
        "config": config
    }, tmp_path)
    os.replace(tmp_path, step_path)
    
    # Save standard copy to main base_path for seamless resume/inference compatibility
    try:
        shutil.copy(step_path, base_path)
    except Exception as e:
        print(f"⚠️ Warning: Could not copy step checkpoint to base path: {e}")
        
    print(f"-> Saved checkpoint to '{step_path}' and copied to '{base_path}'")
    
    # Prune old checkpoints to keep max_to_keep
    pattern = os.path.join(dir_name, f"{name_no_ext}_step_*{ext}")
    ckpt_files = glob.glob(pattern)
    
    ckpts_with_steps = []
    for f in ckpt_files:
        match = re.search(r"_step_(\d+)" + re.escape(ext) + r"$", f)
        if match:
            ckpts_with_steps.append((int(match.group(1)), f))
            
    ckpts_with_steps.sort(key=lambda x: x[0])
    
    while len(ckpts_with_steps) > max_to_keep:
        oldest_step, oldest_file = ckpts_with_steps.pop(0)
        try:
            os.remove(oldest_file)
            print(f"-> Pruned oldest checkpoint file: '{oldest_file}' (keeping last {max_to_keep})")
        except Exception as e:
            print(f"⚠️ Warning: Could not remove old checkpoint '{oldest_file}': {e}")

def main():
    parser = argparse.ArgumentParser(description="Micro-Pretraining a Recursive Universal Causal Transformer (UCT)")
    parser.add_argument("--config", type=str, default="mini", choices=["mini", "target"],
                        help="Choose model config: 'mini' (13.5M unique params) or 'target' (124.13M unique params, unrolled to 176.29M virtual)")
    parser.add_argument("--epochs", type=float, default=1.0, help="Number of pre-training epochs (default: 1.0)")
    parser.add_argument("--steps", type=int, default=None, help="Explicit number of training steps (overrides --epochs calculation)")
    parser.add_argument("--batch_size", type=int, default=1, help="Physical batch size loaded per GPU forward pass")
    parser.add_argument("--accumulate", type=int, default=8, help="Gradient accumulation steps (simulates larger batch size)")
    parser.add_argument("--eval_interval", type=int, default=None, help="Evaluate validation loss every N steps (default: auto-detected, ~10% of total steps)")
    parser.add_argument("--warmup_steps", type=int, default=None, help="Linear LR warmup steps (default: auto-detected, ~5% of total steps)")
    parser.add_argument("--lr", type=float, default=None, help="Custom learning rate (defaults to config preset)")
    parser.add_argument("--dry_run", action="store_true", help="Load dataset, print step-budgeting and scheduling statistics, and exit immediately without instantiating model or allocating VRAM.")
    parser.add_argument("--checkpoint_path", type=str, default=None, help="Custom checkpoint output file path (e.g. micro_llm_200m/uct_pretrain_new.pt)")
    parser.add_argument("--keep_checkpoints", type=int, default=3, help="Max number of rotating checkpoints to keep (default: 3)")
    parser.add_argument("--profile", action="store_true", help="Run PyTorch Profiler on the first 5 steps and output trace file.")
    args = parser.parse_args()
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("====================================================")
    print(f"--- LAUNCHING RECURSIVE TRANSFORMER PRE-TRAINING ---")
    print("====================================================")
    print(f"Target Device: {device}")
    if torch.cuda.is_available():
        print(f"GPU Name: {torch.cuda.get_device_name(0)}")
    
    # Define Model Configurations
    if args.config == "mini":
        config = ModelConfig(
            d_model=256,
            n_iterations=4,
            n_heads=4,
            n_kv_heads=2,
            d_ff=512,
            max_seq_len=256
        )
        peak_lr = args.lr if args.lr else 1e-3
        min_lr = peak_lr * 0.1
    else:  # target (124.13M physical, 176.29M virtual)
        config = ModelConfig(
            d_model=768,
            n_iterations=16,
            n_heads=12,
            n_kv_heads=4,
            d_ff=2048,
            max_seq_len=512
        )
        peak_lr = args.lr if args.lr else 4e-4  # Lower learning rate for stable recursive gradient scaling
        min_lr = peak_lr * 0.1
        
    print(f"\nModel Configuration Type: {args.config.upper()}")
    print(f"-> Vocabulary Size: {config.vocab_size:,}")
    print(f"-> Hidden State Dim: {config.d_model}")
    print(f"-> Attention Heads: {config.n_heads} (Query), {config.n_kv_heads} (KV for GQA)")
    print(f"-> Recurrence Loops (Virtual Depth): {config.n_iterations}")
    print(f"-> Context Size window: {config.max_seq_len} tokens")
    
    # Load Datasets first to detect actual sizes and show steps before initializing model or VRAM
    print("\nInitializing memory-mapped data loaders...")
    train_bin = os.path.join("micro_llm_200m", "data_train_trinity.bin")
    val_bin = os.path.join("micro_llm_200m", "data_val_trinity.bin")
    
    train_ds = MemmappedDataset(train_bin, seq_len=config.max_seq_len)
    val_ds = MemmappedDataset(val_bin, seq_len=config.max_seq_len)
    
    # Pass physical batch size per device dynamically from command-line arguments
    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True)
    val_loader = DataLoader(val_ds, batch_size=min(args.batch_size, 4), shuffle=False)
    
    print(f"-> Train samples: {len(train_ds):,} slices")
    print(f"-> Val samples: {len(val_ds):,} slices")
    
    # Determine the checkpoint path first to pre-load start_step if resuming
    checkpoint_path = args.checkpoint_path if args.checkpoint_path is not None else f"micro_llm_200m/uct_{args.config}_new.pt"
    start_step = 0
    if os.path.exists(checkpoint_path):
        try:
            ckpt_meta = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
            if "step" in ckpt_meta:
                start_step = ckpt_meta["step"]
        except Exception as e:
            pass

    # Clamp accumulate to prevent division-by-zero or empty micro-step loops
    if args.accumulate <= 0:
        args.accumulate = 1

    # Apply 1-Epoch Step-Budgeting Law to calculate total steps dynamically
    if args.steps is not None:
        total_steps = args.steps
        print(f"-> Explicit steps provided: {total_steps:,} (overriding epoch step budget calculation)")
    else:
        steps_per_epoch = len(train_ds) // (args.batch_size * args.accumulate)
        total_steps = int(args.epochs * steps_per_epoch)
        print(f"-> 1-Epoch Step-Budgeting Law Applied:")
        print(f"   * Total training samples (slices): {len(train_ds):,}")
        print(f"   * Context Length: {config.max_seq_len} tokens")
        print(f"   * Effective batch size: {args.batch_size * args.accumulate} slices ({args.batch_size * args.accumulate * config.max_seq_len:,} tokens per step)")
        print(f"   * Steps per Epoch: {steps_per_epoch:,}")
        print(f"   * Training Budget: {args.epochs} epoch(s)")
        print(f"   * Calculated Total steps: {total_steps:,}")
        
    # Set dynamic warmup and eval intervals
    if args.warmup_steps is not None:
        warmup_steps = args.warmup_steps
        print(f"-> Warmup steps: {warmup_steps:,} (explicitly specified)")
    else:
        # 5% of total steps, min 10
        warmup_steps = max(10, int(0.05 * total_steps))
        print(f"-> Warmup steps: {warmup_steps:,} (automatically calculated as 5% of total steps)")
        
    if args.eval_interval is not None:
        eval_interval = args.eval_interval
        print(f"-> Eval interval steps: {eval_interval:,} (explicitly specified)")
    else:
        # 10% of total steps, min 20, max 300
        eval_interval = min(300, max(20, int(0.10 * total_steps)))
        print(f"-> Eval interval steps: {eval_interval:,} (automatically calculated as 10% of total steps)")
        
    if args.dry_run:
        print("\n[DRY RUN SUCCESSFUL] Dataset detected and steps dynamically computed. Exiting gracefully without instantiating model.")
        sys.exit(0)
        
    # Initialize Model
    print("\nInstantiating network weights...")
    model = RecursiveCausalLM(config).to(device)
    unique_params = model.get_num_params(unique_only=True)
    effective_params = model.get_num_params(unique_only=False)
    print(f"-> Unique Parameters Saved (VRAM footprint): {unique_params / 1e6:.2f}M")
    print(f"-> Unrolled Virtual Parameters (Capacity): {effective_params / 1e6:.2f}M")
        
    # Load Optimizer (with Windows compatibility fallbacks)
    print("\nSetting up optimizer...")
    try:
        import bitsandbytes as bnb
        print("-> Community bitsandbytes detected! Using 8-bit AdamW to save VRAM.")
        # Group parameters to apply weight decay strictly to weights and exclude biases/norms
        decay_params = []
        nodecay_params = []
        gate_params = []
        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if "depth_gate" in name:
                gate_params.append(param)
            elif param.ndim >= 2:
                decay_params.append(param)
            else:
                nodecay_params.append(param)
                
        optim_groups = [
            {"params": decay_params, "weight_decay": 0.05},
            {"params": gate_params, "weight_decay": 0.10}, # Strict 0.1 decay prevents gate saturation
            {"params": nodecay_params, "weight_decay": 0.0}
        ]
        optimizer = bnb.optim.Adam8bit(optim_groups, lr=peak_lr)
    except Exception as e:
        print(f"-> bitsandbytes not fully linked: {e}")
        print("-> Falling back to native PyTorch high-performance Fused AdamW optimizer.")
        decay_params = []
        nodecay_params = []
        gate_params = []
        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if "depth_gate" in name:
                gate_params.append(param)
            elif param.ndim >= 2:
                decay_params.append(param)
            else:
                nodecay_params.append(param)
                
        optim_groups = [
            {"params": decay_params, "weight_decay": 0.05},
            {"params": gate_params, "weight_decay": 0.10}, # Strict 0.1 decay prevents gate saturation
            {"params": nodecay_params, "weight_decay": 0.0}
        ]
        # fused=True provides an immediate ~2x speedup in native PyTorch on compatible devices
        use_fused = (device.type == "cuda")
        optimizer = torch.optim.AdamW(optim_groups, lr=peak_lr, fused=use_fused)
        
    # Auto-detect AMP (use float16 for stable Windows laptop standard autocasting)
    use_amp = (device.type == "cuda")
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    
    # Load Offline Tokenizer for live generation callback
    from transformers import AutoTokenizer
    tokenizer_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tokenizer")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    
    best_val_loss = float("inf")
    if args.checkpoint_path is not None:
        checkpoint_path = args.checkpoint_path
    else:
        checkpoint_path = f"micro_llm_200m/uct_{args.config}_new.pt"
    
    # Check if a prior checkpoint exists to resume
    start_step = 0
    if os.path.exists(checkpoint_path):
        print(f"Prior checkpoint detected at '{checkpoint_path}'. Loading weights...")
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        
        # Check if the checkpoint has the old vocabulary size (50257)
        checkpoint_vocab_size = checkpoint["model_state_dict"].get("embeddings.weight", None)
        if checkpoint_vocab_size is not None:
            old_size = checkpoint_vocab_size.shape[0]
            new_size = model.config.vocab_size
            
            if old_size < new_size:
                print(f"-> Upgrading checkpoint vocab size from {old_size} to {new_size} for 128-bit CUDA alignment.")
                # We manually expand the checkpoint tensor to prevent strict=False uninitialized garbage
                old_weight = checkpoint["model_state_dict"]["embeddings.weight"]
                mean_weight = old_weight.mean(dim=0, keepdim=True)
                
                # Expand weights with mean + slight random noise (cast to target GPU device!)
                padding_weight = mean_weight.repeat(new_size - old_size, 1) + torch.randn(new_size - old_size, model.config.d_model, device=device) * 0.02
                checkpoint["model_state_dict"]["embeddings.weight"] = torch.cat([old_weight, padding_weight], dim=0)
                
                # Check if there is an lm_head_bias parameter in the checkpoint and model
                if "lm_head_bias" in checkpoint["model_state_dict"] and model.lm_head_bias is not None:
                    old_bias = checkpoint["model_state_dict"]["lm_head_bias"]
                    padding_bias = torch.zeros(new_size - old_size, device=device, dtype=old_bias.dtype)
                    checkpoint["model_state_dict"]["lm_head_bias"] = torch.cat([old_bias, padding_bias], dim=0)

                # Upgraded Optimizer Momentum Buffers to completely prevent shape mismatches!
                # Identify the saved index corresponding to self.embeddings.weight and self.lm_head_bias
                all_params = []
                for group in optimizer.param_groups:
                    for p in group["params"]:
                        if any(p is x for x in all_params):
                            pass
                        else:
                            all_tokens = all_params.append(p)
                
                # Map model parameters to their index in the optimizer
                param_to_idx = {p: i for i, p in enumerate(all_params)}
                
                # Handle embeddings.weight optimizer state
                embed_param = model.embeddings.weight
                if embed_param in param_to_idx:
                    embed_idx = param_to_idx[embed_param]
                    if "optimizer_state_dict" in checkpoint and "state" in checkpoint["optimizer_state_dict"]:
                        for key in [embed_idx, str(embed_idx)]:
                            if key in checkpoint["optimizer_state_dict"]["state"]:
                                param_state = checkpoint["optimizer_state_dict"]["state"][key]
                                for state_key in ["exp_avg", "exp_avg_sq"]:
                                    if state_key in param_state:
                                        old_buf = param_state[state_key]
                                        padding_shape = (new_size - old_size,) + old_buf.shape[1:]
                                        padding_buf = torch.zeros(padding_shape, dtype=old_buf.dtype, device=old_buf.device)
                                        param_state[state_key] = torch.cat([old_buf, padding_buf], dim=0)
                                print(f"-> Successfully upgraded optimizer momentum buffers for embeddings.weight.")
                                
                # Handle lm_head_bias optimizer state if it exists
                bias_param = model.lm_head_bias
                if bias_param is not None and bias_param in param_to_idx:
                    bias_idx = param_to_idx[bias_param]
                    if "optimizer_state_dict" in checkpoint and "state" in checkpoint["optimizer_state_dict"]:
                        for key in [bias_idx, str(bias_idx)]:
                            if key in checkpoint["optimizer_state_dict"]["state"]:
                                param_state = checkpoint["optimizer_state_dict"]["state"][key]
                                for state_key in ["exp_avg", "exp_avg_sq"]:
                                    if state_key in param_state:
                                        old_buf = param_state[state_key]
                                        padding_buf = torch.zeros(new_size - old_size, dtype=old_buf.dtype, device=old_buf.device)
                                        param_state[state_key] = torch.cat([old_buf, padding_buf], dim=0)
                                print(f"-> Successfully upgraded optimizer momentum buffers for lm_head_bias.")
                
        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if "step" in checkpoint:
            start_step = checkpoint["step"]
            print(f"Resuming training smoothly from step {start_step:,}!")
            
    print(f"\nTraining execution window initialized for {total_steps:,} steps.")
    print("----------------------------------------------------------------")
    
    train_iter = iter(train_loader)
    model.train()
    
    t0 = time.time()
    accumulated_loss = 0.0
    
    # Initialize PyTorch Profiler only if explicitly requested to prevent severe CPU/GPU tracing overhead
    prof = None
    if args.profile:
        print("-> Profiling enabled. Starting PyTorch Profiler for the first 5 steps...")
        prof = torch.profiler.profile(
            schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
            on_trace_ready=torch.profiler.tensorboard_trace_handler('micro_llm_200m/profiler_logs'),
            record_shapes=True,
            profile_memory=True,
            with_stack=True,
            acc_events=True
        )
        prof.start()
    
    for step in range(start_step, total_steps):
        # Calculate dynamic learning rate
        lr = get_lr_scheduler(step, total_steps, warmup_steps, peak_lr, min_lr)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr
            
        optimizer.zero_grad(set_to_none=True)
        
        step_loss = 0.0
        for micro_step in range(args.accumulate):
            try:
                x, y = next(train_iter)
            except StopIteration:
                train_iter = iter(train_loader)
                x, y = next(train_iter)
                
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            
            with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
                logits, loss = model(x, y)
                # Scale loss by accumulation steps
                loss = loss / args.accumulate
                
            step_loss += loss.item()
            scaler.scale(loss).backward()
            
        # Unscale gradients for clipping
        scaler.unscale_(optimizer)
        # Gradient clipping max norm=0.5 protects recursive models from divergence
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
        
        scaler.step(optimizer)
        scaler.update()
        
        accumulated_loss += step_loss
        
        # Training loop stdout logging every 20 steps or on the first step of resuming
        if (step + 1) % 20 == 0 or step == start_step:
            torch.cuda.synchronize() if device.type == "cuda" else None
            t1 = time.time()
            is_interval = (step > start_step and (step + 1) % 20 == 0)
            step_time = (t1 - t0) / (20 if is_interval else 1)
            # Prevent division by zero under ultra-fast or empty steps
            step_time_safe = max(step_time, 1e-5)
            # Calculate running tokens processed per second
            tokens_per_sec = (config.max_seq_len * args.batch_size * args.accumulate) / step_time_safe
            
            vram_gb = 0.0
            if device.type == "cuda":
                vram_gb = torch.cuda.memory_allocated(device) / 1024 / 1024 / 1024
                
            print(f"Step {step+1:5d}/{total_steps:5d} | Loss: {accumulated_loss / (20 if is_interval else 1):.4f} | LR: {lr:.2e} | Speed: {tokens_per_sec:.0f} tok/s ({step_time*1000:.0f} ms) | VRAM: {vram_gb:.2f} GB")
            accumulated_loss = 0.0
            t0 = time.time()
            
        # Evaluation loop
        if (step + 1) % eval_interval == 0:
            # Delete training tensors to free memory prior to evaluation
            if 'logits' in locals(): del logits
            if 'loss' in locals(): del loss
            
            print("\nEvaluating validation set performance...")
            
            # Clear CUDA cache before starting evaluation to maximize free memory
            import gc
            gc.collect()
            if device.type == "cuda":
                torch.cuda.empty_cache()
                
            model.eval()
            val_loss = 0.0
            val_count = min(len(val_loader), 50)  # Limit to 50 batches for fast validation check
            
            val_iter = iter(val_loader)
            with torch.no_grad():
                for v_step in range(val_count):
                    x_v, y_v = next(val_iter)
                    x_v, y_v = x_v.to(device), y_v.to(device)
                    with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
                        logits_v, loss = model(x_v, y_v)
                    val_loss += loss.item()
                    
                    # Delete step variables immediately to free evaluation cache
                    del x_v, y_v, loss, logits_v
                    if device.type == "cuda" and (v_step + 1) % 10 == 0:
                        torch.cuda.empty_cache()
                    
            val_loss /= val_count
            
            # Calculate validation Perplexity (PPL)
            perplexity = math.exp(val_loss) if val_loss < 50 else float("inf")
            print(f"-> Validation Loss: {val_loss:.4f} | Perplexity (PPL): {perplexity:.2f} (Best: {best_val_loss:.4f})")
            
            # Append metrics to local CSV log for plotting curves
            log_path = os.path.join("micro_llm_200m", "metrics_log.csv")
            is_new = not os.path.exists(log_path)
            try:
                with open(log_path, "a", encoding="utf-8") as f:
                    if is_new:
                        f.write("step,val_loss,perplexity\n")
                    f.write(f"{step+1},{val_loss:.6f},{perplexity:.6f}\n")
            except Exception as e:
                print(f"⚠️ Warning: Could not write metrics to CSV: {e}")
                
            # Trigger Live Generation Callback!
            print("Triggering Live Generation Callback...")
            run_live_sample(model, tokenizer, device)
            
            # Save best checkpoint
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                print(f"-> New best loss!")
                save_checkpoint(
                    model=model,
                    optimizer=optimizer,
                    step=step + 1,
                    val_loss=val_loss,
                    config=config,
                    base_path=checkpoint_path,
                    max_to_keep=args.keep_checkpoints
                )
            
            # Explicitly release all validation tensors from Python's local scope
            if 'x_v' in locals(): del x_v
            if 'y_v' in locals(): del y_v
            if 'logits_v' in locals(): del logits_v
            if 'val_iter' in locals(): del val_iter
            if 'loss' in locals(): del loss
            import gc
            gc.collect()
            if device.type == "cuda":
                torch.cuda.empty_cache()
                
            print("----------------------------------------------------------------")
            model.train()
            t0 = time.time()  # reset timer after eval lag
            
        if prof is not None:
            prof.step()

    if prof is not None:
        prof.stop()
    print("\nPre-training run completed successfully! Best model weights saved.")
    print(f"Weights ready for interactive autoregressive generation at '{checkpoint_path}'.")

if __name__ == "__main__":
    main()