import torch
import torch.nn as nn
from torch.nn import functional as F
from model.transformer import GPT
from model.config import ModelConfig
from model.tokenizer import AdvancedTokenizer
from train.dataset import TextDataset
from torch.utils.data import DataLoader
import os
import time
import math
import copy
import requests
from tqdm import tqdm

# Allowlisting ModelConfig for PyTorch 2.6+ security
try:
    torch.serialization.add_safe_globals([ModelConfig])
except:
    pass


# ─────────────────────────────────────────────────────────────────────────────
# Exponential Moving Average (EMA) of model weights
# ─────────────────────────────────────────────────────────────────────────────
class EMA:
    """Maintains an exponential moving average of model parameters for better generalization."""
    def __init__(self, model, decay=0.999):
        self.decay = decay
        self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}

    def update(self, model):
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in self.shadow:
                    self.shadow[name].mul_(self.decay).add_(param, alpha=1 - self.decay)

    def apply(self, model):
        """Apply EMA weights to model (for evaluation/saving)."""
        backup = {}
        for name, param in model.named_parameters():
            if name in self.shadow:
                backup[name] = param.clone()
                param.data.copy_(self.shadow[name])
        return backup

    def restore(self, model, backup):
        """Restore original weights after EMA evaluation."""
        for name, param in model.named_parameters():
            if name in backup:
                param.data.copy_(backup[name])


# ─────────────────────────────────────────────────────────────────────────────
# Early Stopping
# ─────────────────────────────────────────────────────────────────────────────
class EarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4):
        self.patience = patience
        self.min_delta = min_delta
        self.best_loss = float('inf')
        self.counter = 0

    def check(self, val_loss) -> bool:
        """Returns True if training should stop."""
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            return False
        self.counter += 1
        return self.counter >= self.patience


# ─────────────────────────────────────────────────────────────────────────────
# Cosine LR with Warm Restarts
# ─────────────────────────────────────────────────────────────────────────────
def get_lr(step, total_steps, max_lr, min_lr_ratio=0.1, warmup_steps=100):
    """Cosine annealing with linear warmup."""
    if step < warmup_steps:
        return max_lr * (step + 1) / max(1, warmup_steps)
    decay_ratio = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return max_lr * min_lr_ratio + max_lr * (1 - min_lr_ratio) * coeff


# ─────────────────────────────────────────────────────────────────────────────
# Main Training Function
# ─────────────────────────────────────────────────────────────────────────────
def train(dataset_path=None, job=None, text_content=None, category="text",
          continue_learning=True, noise_level=0.0, pin_memory=True,
          expert_offloading=True, max_epochs=100):
    """
    Optimized training for RTX 4060 (8GB VRAM).
    
    Key optimizations:
      • Mixed precision (bf16/fp16)
      • Gradient accumulation (effective batch 128)
      • Gradient checkpointing (2x VRAM savings)
      • torch.compile with CUDA graphs
      • EMA weights averaging
      • Cosine LR with warmup
      • Label smoothing
      • Early stopping
      • Expert offloading to CPU
    """
    from model.category_manager import get_category_manager

    # 1. Load Data
    text = None
    if text_content:
        text = text_content
    elif dataset_path and os.path.exists(dataset_path):
        pass

    if not text:
        data_path = 'input.txt'
        if not os.path.exists(data_path):
            print("input.txt not found. Using dummy data for test.")
            try:
                url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
                r = requests.get(url)
                with open(data_path, 'w') as f:
                    f.write(r.text)
                print("Downloaded TinyShakespeare.")
            except:
                text = "Hello world. This is a test of the Sail AI agent system. " * 100
                with open(data_path, 'w') as f:
                    f.write(text)

        if not text:
            with open(data_path, 'r', encoding='utf-8') as f:
                text = f.read()

    # 2. Check for existing model (Continuous Learning)
    existing_checkpoint = None
    if continue_learning and os.path.exists("sail.pt"):
        print("Loading existing model for continuous learning...")
        try:
            existing_checkpoint = torch.load("sail.pt", map_location='cpu', weights_only=True)
            print(f"Loaded existing model. Continuing training on category: {category}")
        except Exception as e:
            print(f"Could not load existing model: {e}. Starting fresh.")
            existing_checkpoint = None

    # 3. Train Tokenizer
    print("Initializing Advanced Tokenizer...")
    if job: job.message = "Training Tokenizer..."

    tokenizer = AdvancedTokenizer(vocab_size=5000)

    if existing_checkpoint and 'vocab' in existing_checkpoint:
        print("Merging vocabularies...")
        tokenizer.word_to_id = existing_checkpoint['vocab'].copy()
        
        # Ensure all special instruction tokens are retained
        for token in tokenizer.specials:
            if token not in tokenizer.word_to_id:
                tokenizer.word_to_id[token] = len(tokenizer.word_to_id)
                
        tokenizer.id_to_word = {v: k for k, v in tokenizer.word_to_id.items()}
        tokenizer.is_trained = True
        tokenizer.train(text)
    else:
        tokenizer.train(text)

    # 4. Prepare Config
    config = ModelConfig()
    config.vocab_size = len(tokenizer.word_to_id)
    print(f"Vocab Size: {config.vocab_size}")

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    config.device = device
    config.pin_memory = pin_memory
    config.expert_offloading = expert_offloading

    if device == 'cuda':
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
        print(f"GPU: {gpu_name} ({gpu_mem:.1f} GB)")
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    if job: job.message = f"Device: {device}. Vocab: {config.vocab_size}. Category: {category}"

    # 5. Initialize Model
    model = GPT(config)

    if existing_checkpoint and 'model_state_dict' in existing_checkpoint:
        try:
            print("Loading weights into model (CPU)...")
            state_dict = existing_checkpoint['model_state_dict']
            
            # Handle dynamic vocabulary expansion (e.g. adding new instruction tokens)
            if 'token_emb.weight' in state_dict and state_dict['token_emb.weight'].shape[0] != config.vocab_size:
                old_v = state_dict['token_emb.weight'].shape[0]
                new_v = config.vocab_size
                print(f"Resizing embeddings from {old_v} to {new_v} due to vocabulary expansion...")
                
                # Clone randomly initialized new embeddings and overwrite the overlapping part
                new_emb = model.token_emb.weight.data.clone()
                min_v = min(old_v, new_v)
                new_emb[:min_v] = state_dict['token_emb.weight'][:min_v]
                state_dict['token_emb.weight'] = new_emb
                
                if 'lm_head.weight' in state_dict:
                    new_lm = model.lm_head.weight.data.clone()
                    new_lm[:min_v] = state_dict['lm_head.weight'][:min_v]
                    state_dict['lm_head.weight'] = new_lm

            model.load_state_dict(state_dict, strict=False)
            print("Loaded existing model weights.")
        except Exception as e:
            print(f"Could not load weights: {e}. Training from scratch.")

    del existing_checkpoint
    import gc
    gc.collect()

    # Move to GPU
    print(f"Moving model to {device}...")
    model.to(device)
    if device == 'cuda':
        torch.cuda.empty_cache()

    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {n_params/1e6:.2f}M")

    # Offload idle experts if enabled
    if config.expert_offloading and device == 'cuda':
        print("Expert offloading enabled — idle experts will use CPU RAM")

    # 6. Optimizer (AdamW with fused kernels)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        weight_decay=config.weight_decay,
        fused=(device == 'cuda')
    )

    # Mixed Precision Scaler
    use_amp = config.use_amp and device == 'cuda'
    scaler = torch.amp.GradScaler('cuda', enabled=use_amp)
    if use_amp:
        print("Mixed Precision Training: ENABLED (bf16/fp16)")

    # EMA
    ema = None
    if config.use_ema:
        ema = EMA(model, decay=config.ema_decay)
        print(f"EMA: ENABLED (decay={config.ema_decay})")

    # Early Stopping
    early_stopping = EarlyStopping(patience=config.patience, min_delta=config.min_delta)

    # 7. Dataset & Dataloader
    print("Encoding dataset...")
    if job: job.message = "Encoding Dataset..."

    if dataset_path and dataset_path.endswith('.json'):
        print("Detected Instruction Dataset (JSON).")
        from train.instruction_loader import InstructionDataset
        train_ds = InstructionDataset(dataset_path, tokenizer, config.block_size)
    else:
        train_ds = TextDataset(text, tokenizer, config, pin_memory=config.pin_memory)
        if len(train_ds) == 0:
            print("Dataset too small for block_size. Repeating text...")
            text = text * (config.block_size // len(text) + 2)
            train_ds = TextDataset(text, tokenizer, config)

    num_workers = min(os.cpu_count() or 2, 4)
    train_dl = DataLoader(
        train_ds,
        batch_size=config.batch_size,
        shuffle=True,
        pin_memory=(device == 'cuda'),
        num_workers=num_workers,
        prefetch_factor=2 if (device == 'cuda') else None,
        persistent_workers=True if (device == 'cuda' and num_workers > 0) else False,
        drop_last=True,
    )

    # 8. torch.compile (CUDA Graphs)
    if config.use_compile and device == 'cuda':
        print("Compiling model with torch.compile (mode='reduce-overhead')...")
        try:
            model = torch.compile(model, mode="reduce-overhead")
            print("Model compiled with CUDA Graphs!")
        except Exception as e:
            print(f"Warning: torch.compile failed ({e}). Proceeding without.")

    # 9. Training Loop
    accumulation_steps = config.gradient_accumulation_steps
    total_steps = (len(train_dl) // accumulation_steps + 1) * max_epochs

    print(f"\n{'='*60}")
    print(f"  TRAINING CONFIG")
    print(f"  Epochs          : {max_epochs}")
    print(f"  Batch Size      : {config.batch_size} (eff. {config.batch_size * accumulation_steps})")
    print(f"  Grad Accum      : {accumulation_steps}")
    print(f"  Learning Rate   : {config.learning_rate}")
    print(f"  Label Smoothing : {config.label_smoothing}")
    print(f"  Noise Level     : {noise_level}")
    print(f"  Total Steps     : ~{total_steps}")
    print(f"  Category        : {category}")
    print(f"{'='*60}\n")

    model.train()
    step = 0
    best_loss = float('inf')
    optimizer.zero_grad(set_to_none=True)

    for epoch in range(max_epochs):
        epoch_loss = 0.0
        epoch_steps = 0

        pbar = tqdm(train_dl, desc=f"Epoch {epoch+1}/{max_epochs}")
        for i, (x, y) in enumerate(pbar):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)

            # Mixed Precision Forward Pass
            with torch.amp.autocast('cuda', enabled=use_amp):
                logits, loss = model(x, y, noise_level=noise_level)
                loss = loss / accumulation_steps

            # Scaled Backward Pass
            scaler.scale(loss).backward()

            if (i + 1) % accumulation_steps == 0 or (i + 1) == len(train_dl):
                # Gradient Clipping
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.max_grad_norm)

                # LR Schedule
                lr = get_lr(step, total_steps, config.learning_rate,
                           warmup_steps=config.warmup_steps)
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr

                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)

                # EMA update
                if ema is not None:
                    ema.update(model)

                actual_loss = loss.item() * accumulation_steps
                epoch_loss += actual_loss
                epoch_steps += 1

                pbar.set_description(
                    f"Loss: {actual_loss:.4f} LR: {lr:.2e}"
                )

                if job:
                    job.message = f"Epoch {epoch+1} - Loss: {actual_loss:.4f} - LR: {lr:.2e}"

                step += 1

        # End of epoch
        avg_loss = epoch_loss / max(epoch_steps, 1)
        print(f"  Epoch {epoch+1} avg loss: {avg_loss:.4f}")

        # Save best checkpoint
        if avg_loss < best_loss:
            best_loss = avg_loss
            _save_checkpoint(model, config, tokenizer, category, job, ema, tag="best")

        # Early stopping
        if early_stopping.check(avg_loss):
            print(f"  Early stopping triggered at epoch {epoch+1} (patience={config.patience})")
            break

        # VRAM monitoring
        if device == 'cuda' and (epoch + 1) % 5 == 0:
            alloc = torch.cuda.memory_allocated() / 1024**3
            total = torch.cuda.get_device_properties(0).total_memory / 1024**3
            print(f"  VRAM: {alloc:.1f}GB / {total:.1f}GB")

    print("Training Complete.")
    _save_checkpoint(model, config, tokenizer, category, job, ema, tag="final")

    # Record category training
    cat_manager = get_category_manager()
    details = dataset_path or "text_content"
    cat_manager.add_training(category, details, config.vocab_size)
    print(f"Category '{category}' training recorded.")


def _save_checkpoint(model, config, tokenizer, category, job, ema, tag=""):
    """Save checkpoint, optionally with EMA weights."""
    if job: job.message = "Saving Model..."

    # Apply EMA weights for saving if available
    backup = None
    if ema is not None:
        backup = ema.apply(model)

    checkpoint = {
        'model_state_dict': model.state_dict(),
        'config': config,
        'vocab': tokenizer.word_to_id,
        'vocab_size': config.vocab_size,
        'category': category,
        'timestamp': time.time(),
        'tag': tag,
    }

    # Always save to sail.pt
    torch.save(checkpoint, "sail.pt")
    print(f"Saved 'sail.pt' (Category: {category}, Tag: {tag})")

    # Also save job-specific backup
    if job:
        backup_path = f"sail_{job.id}.pt"
        torch.save(checkpoint, backup_path)
        print(f"Backup saved to '{backup_path}'")

    # Restore original weights after EMA save
    if ema is not None and backup is not None:
        ema.restore(model, backup)


if __name__ == '__main__':
    train()