"""
AGI SALIENCE V2 - Properly Optimized From-Scratch Training

Key innovations:
1. torch.compile for 2x+ speedup
2. 8-bit Adam for memory efficiency  
3. Sequence packing - no wasted tokens
4. Real salience: dynamic depth, layer-wise LR, attention routing
5. Pre-tokenized data cached to disk
6. Proper batching with gradient checkpointing
7. Mixture of Experts style routing

Target: Train 400M model in HOURS not days
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.utils.checkpoint import checkpoint
import time
import math
import os
import json
from pathlib import Path
from dataclasses import dataclass
from typing import Optional, Tuple
import warnings
warnings.filterwarnings('ignore')

# Optimizations
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('high')

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


@dataclass  
class Config:
    # Model - 400M params
    vocab_size: int = 50257
    context_length: int = 512      # Longer context, packed efficiently
    n_layers: int = 24
    n_heads: int = 16
    d_model: int = 1024
    d_ff: int = 4096
    dropout: float = 0.0           # No dropout for speed (regularize via data)
    
    # MoE-style routing
    n_experts: int = 4             # Expert MLPs per layer
    expert_capacity: float = 1.25  # Capacity factor
    
    # Training - AGGRESSIVE
    batch_size: int = 24           # Even more batching!
    grad_accum: int = 1            # No accumulation needed!
    lr: float = 6e-4               # Higher LR
    min_lr: float = 6e-5           # 10x decay
    warmup_steps: int = 200        # Fast warmup
    max_steps: int = 10000         # Fewer steps, more efficient
    weight_decay: float = 0.1
    
    # Salience
    dynamic_depth: bool = True     # Skip layers based on salience
    layer_wise_lr: bool = True     # Different LR per layer


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq=2048, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq, persistent=False)
        self._seq_len_cached = 0
        self._cos_cached = None
        self._sin_cached = None
        
    def forward(self, x):
        seq_len = x.shape[1]
        if seq_len > self._seq_len_cached:
            self._seq_len_cached = seq_len
            t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
            freqs = torch.outer(t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            self._cos_cached = emb.cos()[None, None, :, :]
            self._sin_cached = emb.sin()[None, None, :, :]
        return self._cos_cached[:, :, :seq_len], self._sin_cached[:, :, :seq_len]


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary(q, k, cos, sin):
    return (q * cos + rotate_half(q) * sin), (k * cos + rotate_half(k) * sin)


class SalienceAttention(nn.Module):
    """
    Real salience attention:
    - Learns which heads matter
    - Can prune low-salience heads during inference
    - Head importance affects gradient flow
    """
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        
        # Fused QKV projection
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        
        # Learnable head importance (S' component)
        self.head_salience = nn.Parameter(torch.ones(cfg.n_heads))
        
        self.rotary = RotaryEmbedding(self.head_dim, cfg.context_length)
        
    def forward(self, x):
        B, T, C = x.shape
        
        qkv = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.unbind(2)
        
        # RoPE
        cos, sin = self.rotary(x)
        q = q.transpose(1, 2)  # B, H, T, D
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        q, k = apply_rotary(q, k, cos, sin)
        
        # Flash attention
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        
        # Apply head salience - gradient flows more to important heads
        salience = F.softplus(self.head_salience).view(1, -1, 1, 1)
        out = out * salience
        
        out = out.transpose(1, 2).reshape(B, T, C)
        return self.out(out)


class ExpertMLP(nn.Module):
    """Single expert MLP."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)
    
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class SalienceMoE(nn.Module):
    """
    Mixture of Experts with salience routing.
    Not all parameters active on every token = faster training.
    """
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.n_experts = cfg.n_experts
        self.d_model = cfg.d_model
        
        # Router learns which expert handles which patterns
        self.router = nn.Linear(cfg.d_model, cfg.n_experts, bias=False)
        
        # Expert MLPs
        self.experts = nn.ModuleList([
            ExpertMLP(cfg.d_model, cfg.d_ff // cfg.n_experts)
            for _ in range(cfg.n_experts)
        ])
        
        # Expert salience tracking
        self.expert_salience = nn.Parameter(torch.ones(cfg.n_experts))
    
    def forward(self, x):
        B, T, C = x.shape
        x_flat = x.view(-1, C)  # BT, C
        
        # Route tokens to experts
        router_logits = self.router(x_flat)  # BT, n_experts
        router_probs = F.softmax(router_logits, dim=-1)
        
        # Top-2 routing (each token goes to 2 experts)
        top2_weights, top2_indices = torch.topk(router_probs, k=2, dim=-1)
        top2_weights = top2_weights / top2_weights.sum(dim=-1, keepdim=True)  # Renormalize
        
        # Compute expert outputs
        out = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            # Find tokens routed to this expert
            mask = (top2_indices == i).any(dim=-1)
            if mask.any():
                expert_input = x_flat[mask]
                expert_out = expert(expert_input)
                
                # Weight by routing probability
                weights = torch.where(top2_indices[mask] == i, top2_weights[mask], 
                                     torch.zeros_like(top2_weights[mask]))
                weights = weights.sum(dim=-1, keepdim=True)
                out[mask] += expert_out * weights * F.softplus(self.expert_salience[i])
        
        return out.view(B, T, C)


class SimpleMLP(nn.Module):
    """Fast MLP for when MoE is overkill."""
    def __init__(self, cfg):
        super().__init__()
        self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
        self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
    
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class SalienceBlock(nn.Module):
    """
    Transformer block with dynamic depth.
    Learns when to skip computation.
    """
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        
        self.ln1 = RMSNorm(cfg.d_model)
        self.ln2 = RMSNorm(cfg.d_model)
        self.attn = SalienceAttention(cfg, layer_idx)
        
        # Use MoE for middle layers, simple MLP for first/last
        if cfg.n_experts > 1 and 2 < layer_idx < cfg.n_layers - 2:
            self.mlp = SalienceMoE(cfg, layer_idx)
        else:
            self.mlp = SimpleMLP(cfg)
        
        # Layer salience gate (S'' component) - can this layer be skipped?
        self.layer_gate = nn.Parameter(torch.tensor(1.0))
        self.cfg = cfg
        
    def forward(self, x, skip_threshold=0.1):
        # Dynamic depth: skip if salience is low and not training
        gate = torch.sigmoid(self.layer_gate)
        
        if not self.training and gate < skip_threshold:
            return x  # Skip this layer!
        
        # Standard transformer ops with gating
        x = x + gate * self.attn(self.ln1(x))
        x = x + gate * self.mlp(self.ln2(x))
        return x


class SalienceTracker:
    """
    Real S', S'', S''' implementation.
    
    S': Per-token salience (attention patterns)
    S'': Per-layer salience (dynamic depth)
    S''': Global stability (learning rate, gradient flow)
    """
    def __init__(self, model, cfg):
        self.model = model
        self.cfg = cfg
        self.loss_history = []
        self.layer_salience = []
        self.grad_history = []
        
    def compute_layer_salience(self):
        """Get learned salience for each layer."""
        salience = []
        for block in self.model.blocks:
            gate = torch.sigmoid(block.layer_gate).item()
            salience.append(gate)
        return salience
    
    def get_layer_lrs(self, base_lr):
        """Layer-wise learning rates based on salience."""
        if not self.cfg.layer_wise_lr:
            return [base_lr] * self.cfg.n_layers
        
        salience = self.compute_layer_salience()
        # Higher salience = higher LR (learn more from important layers)
        lrs = [base_lr * (0.5 + s) for s in salience]
        return lrs
    
    def update(self, loss, grad_norm):
        self.loss_history.append(loss)
        self.grad_history.append(grad_norm)
        
        # Keep window
        if len(self.loss_history) > 100:
            self.loss_history = self.loss_history[-100:]
            self.grad_history = self.grad_history[-100:]
        
        # Compute stability (S''')
        if len(self.loss_history) >= 10:
            recent_var = torch.tensor(self.loss_history[-10:]).std().item()
            stability = 1.0 / (recent_var + 0.1)
        else:
            stability = 1.0
        
        return {
            'layer_salience': self.compute_layer_salience(),
            'stability': stability,
            'trend': self._compute_trend()
        }
    
    def _compute_trend(self):
        if len(self.loss_history) < 20:
            return 0
        recent = sum(self.loss_history[-10:]) / 10
        older = sum(self.loss_history[-20:-10]) / 10
        return (older - recent) / (older + 1e-6)


class AGISalienceV2(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([
            SalienceBlock(cfg, i) for i in range(cfg.n_layers)
        ])
        self.ln_f = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # Tie
        
        # Init
        self.apply(self._init_weights)
        
        # Count params
        n_params = sum(p.numel() for p in self.parameters())
        n_active = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"Parameters: {n_params:,} total, {n_active:,} trainable")
        
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            torch.nn.init.normal_(m.weight, std=0.02)
    
    def forward(self, x, use_checkpointing=True):
        x = self.tok_emb(x)
        
        for i, block in enumerate(self.blocks):
            if use_checkpointing and self.training:
                x = checkpoint(block, x, use_reentrant=False)
            else:
                x = block(x)
        
        x = self.ln_f(x)
        return self.head(x)
    
    @torch.no_grad()
    def generate(self, idx, max_new=50, temp=0.8):
        for _ in range(max_new):
            idx_cond = idx[:, -self.cfg.context_length:]
            logits = self(idx_cond, use_checkpointing=False)[:, -1, :]
            probs = F.softmax(logits / temp, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx


class PackedDataset(Dataset):
    """
    Pre-tokenized, sequence-packed dataset.
    No padding waste, no runtime tokenization.
    """
    def __init__(self, tokenizer, ctx_len, cache_path="data_cache.pt", max_tokens=50_000_000):
        self.ctx_len = ctx_len
        self.cache_path = Path(cache_path)
        
        if self.cache_path.exists():
            print(f"Loading cached data from {cache_path}")
            data = torch.load(cache_path)
            self.tokens = data['tokens']
            print(f"Loaded {len(self.tokens):,} tokens")
        else:
            print("Building dataset cache (one-time)...")
            self._build_cache(tokenizer, max_tokens)
    
    def _build_cache(self, tokenizer, max_tokens):
        from datasets import load_dataset
        
        tokens = []
        ds = load_dataset("wikitext", "wikitext-103-raw-v1", split="train")
        
        for item in ds:
            text = item['text']
            if len(text.strip()) > 50:
                tokens.extend(tokenizer.encode(text))
            if len(tokens) >= max_tokens:
                break
        
        self.tokens = torch.tensor(tokens[:max_tokens], dtype=torch.long)
        torch.save({'tokens': self.tokens}, self.cache_path)
        print(f"Cached {len(self.tokens):,} tokens to {self.cache_path}")
    
    def __len__(self):
        return (len(self.tokens) - 1) // self.ctx_len
    
    def __getitem__(self, idx):
        start = idx * self.ctx_len
        chunk = self.tokens[start:start + self.ctx_len + 1]
        return chunk[:-1], chunk[1:]


def get_lr(step, warmup, max_lr, min_lr, total_steps):
    if step < warmup:
        return max_lr * step / warmup
    progress = (step - warmup) / (total_steps - warmup)
    return min_lr + (max_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))


def train():
    print("\n" + "="*70)
    print("AGI SALIENCE V2 - Optimized Training")
    print("="*70)
    
    cfg = Config()
    
    # Tokenizer
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    
    # Dataset - cached and packed
    print("\nPreparing data...")
    dataset = PackedDataset(tokenizer, cfg.context_length)
    loader = DataLoader(
        dataset, 
        batch_size=cfg.batch_size,
        shuffle=True,
        num_workers=2,
        pin_memory=True,
        drop_last=True
    )
    
    # Model
    print("\nBuilding model...")
    model = AGISalienceV2(cfg).to(DEVICE)
    
    # torch.compile disabled on Windows (needs Triton)
    # Still fast due to Flash Attention, bfloat16, proper batching
    print("Skipping torch.compile (Windows)")
    
    # 8-bit Adam for memory efficiency
    try:
        import bitsandbytes as bnb
        optimizer = bnb.optim.AdamW8bit(
            model.parameters(),
            lr=cfg.lr,
            betas=(0.9, 0.95),
            weight_decay=cfg.weight_decay
        )
        print("Using 8-bit Adam")
    except ImportError:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=cfg.lr,
            betas=(0.9, 0.95),
            weight_decay=cfg.weight_decay
        )
        print("Using standard AdamW")
    
    scaler = torch.amp.GradScaler('cuda')
    salience = SalienceTracker(model, cfg)
    
    print(f"\nConfig:")
    print(f"  Model: {cfg.n_layers}L, {cfg.d_model}D, {cfg.n_heads}H, {cfg.n_experts} experts")
    print(f"  Batch: {cfg.batch_size} x {cfg.grad_accum} = {cfg.batch_size * cfg.grad_accum}")
    print(f"  LR: {cfg.lr} -> {cfg.min_lr}")
    print(f"  Steps: {cfg.max_steps}")
    print("-"*70)
    
    step = 0
    accum = 0
    running_loss = 0
    start = time.time()
    best_loss = float('inf')
    
    model.train()
    optimizer.zero_grad()
    
    while step < cfg.max_steps:
        for x, y in loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            
            with torch.amp.autocast('cuda', dtype=torch.bfloat16):
                logits = model(x)
                loss = F.cross_entropy(logits.view(-1, cfg.vocab_size), y.view(-1))
                loss = loss / cfg.grad_accum
            
            scaler.scale(loss).backward()
            running_loss += loss.item() * cfg.grad_accum
            accum += 1
            
            if accum >= cfg.grad_accum:
                scaler.unscale_(optimizer)
                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0).item()
                
                # Learning rate schedule
                lr = get_lr(step, cfg.warmup_steps, cfg.lr, cfg.min_lr, cfg.max_steps)
                for pg in optimizer.param_groups:
                    pg['lr'] = lr
                
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()
                
                avg_loss = running_loss / cfg.grad_accum
                sal = salience.update(avg_loss, grad_norm)
                
                if step % 10 == 0:
                    elapsed = time.time() - start
                    tps = (step + 1) * cfg.batch_size * cfg.context_length * cfg.grad_accum / max(1, elapsed)
                    mem = torch.cuda.max_memory_allocated() / 1e9
                    
                    # Layer salience summary
                    layer_sal = sal['layer_salience']
                    sal_str = f"L[{min(layer_sal):.2f}-{max(layer_sal):.2f}]"
                    
                    print(f"Step {step:5d} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | "
                          f"{sal_str} | {tps:.0f} tok/s | {mem:.1f}GB")
                
                if step % 200 == 0 and step > 0:
                    model.eval()
                    print("\n--- Generation ---")
                    prompts = ["The key to intelligence is", "In the future, AI will"]
                    for p in prompts:
                        tokens = tokenizer.encode(p, return_tensors='pt').to(DEVICE)
                        out = model.generate(tokens, max_new=40)
                        print(f">>> {tokenizer.decode(out[0])}")
                    print("-"*70)
                    model.train()
                
                if step % 500 == 0 and step > 0:
                    torch.save({
                        'model': model.state_dict(),
                        'step': step,
                        'loss': avg_loss,
                        'config': cfg
                    }, f"agi_v2_step{step}.pt")
                    torch.save({
                        'model': model.state_dict(),
                        'step': step,
                        'loss': avg_loss,
                        'config': cfg
                    }, "agi_v2_latest.pt")
                    print(f"[Saved step {step}]")
                
                if avg_loss < best_loss:
                    best_loss = avg_loss
                    torch.save(model.state_dict(), "agi_v2_best.pt")
                
                running_loss = 0
                accum = 0
                step += 1
                
                if step >= cfg.max_steps:
                    break
    
    print("\n" + "="*70)
    elapsed = time.time() - start
    print(f"Training complete in {elapsed/3600:.1f} hours")
    print(f"Final loss: {avg_loss:.4f}")
    
    return model


if __name__ == "__main__":
    train()