# createsonline/ml/training.py
"""
CREATESONLINE Training Infrastructure — Pure NumPy

Optimizers, learning rate schedulers, data loaders,
training loops, and model checkpointing.
"""

import math
import time
import json
import os
import copy
from typing import List, Dict, Optional, Callable, Tuple, Any
import numpy as np


# =====================================================
# OPTIMIZERS
# =====================================================

class SGD:
    """Stochastic Gradient Descent with momentum"""
    
    def __init__(self, params: List[np.ndarray], lr: float = 0.01,
                 momentum: float = 0.0, weight_decay: float = 0.0):
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.params = params
        self.velocities = [np.zeros_like(p) for p in params]
        self.step_count = 0
    
    def step(self, grads: List[np.ndarray]):
        """Update parameters"""
        self.step_count += 1
        for i, (param, grad) in enumerate(zip(self.params, grads)):
            if self.weight_decay > 0:
                grad = grad + self.weight_decay * param
            if self.momentum > 0:
                self.velocities[i] = self.momentum * self.velocities[i] + grad
                param -= self.lr * self.velocities[i]
            else:
                param -= self.lr * grad
    
    def zero_grad(self, grads: List[np.ndarray]):
        for g in grads:
            g[:] = 0


class Adam:
    """
    Adam optimizer (Kingma & Ba, 2014).
    The standard optimizer for transformers.
    """
    
    def __init__(self, params: List[np.ndarray], lr: float = 0.001,
                 beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8,
                 weight_decay: float = 0.0):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.params = params
        
        self.m = [np.zeros_like(p) for p in params]  # First moment
        self.v = [np.zeros_like(p) for p in params]  # Second moment
        self.step_count = 0
    
    def step(self, grads: List[np.ndarray]):
        """Update parameters using Adam"""
        self.step_count += 1
        t = self.step_count
        
        for i, (param, grad) in enumerate(zip(self.params, grads)):
            if self.weight_decay > 0:
                grad = grad + self.weight_decay * param
            
            # Update moments
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * grad**2
            
            # Bias correction
            m_hat = self.m[i] / (1 - self.beta1**t)
            v_hat = self.v[i] / (1 - self.beta2**t)
            
            # Update
            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
    
    def zero_grad(self, grads: List[np.ndarray]):
        for g in grads:
            g[:] = 0


class AdamW:
    """
    AdamW optimizer (Loshchilov & Hutter, 2019).
    Decoupled weight decay — the preferred optimizer for LLM training.
    """
    
    def __init__(self, params: List[np.ndarray], lr: float = 0.001,
                 beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8,
                 weight_decay: float = 0.01):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.params = params
        
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.step_count = 0
    
    def step(self, grads: List[np.ndarray]):
        """Update parameters using AdamW"""
        self.step_count += 1
        t = self.step_count
        
        for i, (param, grad) in enumerate(zip(self.params, grads)):
            # Decoupled weight decay (applied BEFORE moment update)
            if self.weight_decay > 0:
                param -= self.lr * self.weight_decay * param
            
            # Update moments
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * grad**2
            
            # Bias correction
            m_hat = self.m[i] / (1 - self.beta1**t)
            v_hat = self.v[i] / (1 - self.beta2**t)
            
            # Update
            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
    
    def zero_grad(self, grads: List[np.ndarray]):
        for g in grads:
            g[:] = 0


# =====================================================
# LEARNING RATE SCHEDULERS
# =====================================================

class WarmupScheduler:
    """Linear warmup then constant learning rate"""
    
    def __init__(self, optimizer, warmup_steps: int = 1000):
        self.optimizer = optimizer
        self.warmup_steps = warmup_steps
        self.base_lr = optimizer.lr
        self.step_count = 0
    
    def step(self):
        self.step_count += 1
        if self.step_count <= self.warmup_steps:
            self.optimizer.lr = self.base_lr * (self.step_count / self.warmup_steps)
        else:
            self.optimizer.lr = self.base_lr
    
    def get_lr(self) -> float:
        return self.optimizer.lr


class CosineScheduler:
    """Cosine annealing learning rate with optional warmup"""
    
    def __init__(self, optimizer, total_steps: int, warmup_steps: int = 0,
                 min_lr: float = 0.0):
        self.optimizer = optimizer
        self.total_steps = total_steps
        self.warmup_steps = warmup_steps
        self.base_lr = optimizer.lr
        self.min_lr = min_lr
        self.step_count = 0
    
    def step(self):
        self.step_count += 1
        if self.step_count <= self.warmup_steps:
            # Linear warmup
            self.optimizer.lr = self.base_lr * (self.step_count / max(1, self.warmup_steps))
        else:
            # Cosine decay
            progress = (self.step_count - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)
            self.optimizer.lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))
    
    def get_lr(self) -> float:
        return self.optimizer.lr


class StepScheduler:
    """Step-wise learning rate decay"""
    
    def __init__(self, optimizer, step_size: int = 1000, gamma: float = 0.1):
        self.optimizer = optimizer
        self.step_size = step_size
        self.gamma = gamma
        self.base_lr = optimizer.lr
        self.step_count = 0
    
    def step(self):
        self.step_count += 1
        if self.step_count > 0 and self.step_count % self.step_size == 0:
            self.optimizer.lr *= self.gamma
    
    def get_lr(self) -> float:
        return self.optimizer.lr


# =====================================================
# DATA LOADER
# =====================================================

class DataLoader:
    """
    Mini-batch data loader with shuffling.
    
    Usage:
        loader = DataLoader(data, batch_size=32, shuffle=True)
        for batch in loader:
            # batch is a dict of numpy arrays
            x, y = batch['input'], batch['target']
    """
    
    def __init__(self, data: Dict[str, np.ndarray], batch_size: int = 32,
                 shuffle: bool = True, drop_last: bool = False):
        self.data = data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last
        
        # Get dataset size from first array
        first_key = next(iter(data))
        self.n_samples = len(data[first_key])
    
    def __len__(self) -> int:
        if self.drop_last:
            return self.n_samples // self.batch_size
        return math.ceil(self.n_samples / self.batch_size)
    
    def __iter__(self):
        indices = np.arange(self.n_samples)
        if self.shuffle:
            np.random.shuffle(indices)
        
        for start in range(0, self.n_samples, self.batch_size):
            end = start + self.batch_size
            if self.drop_last and end > self.n_samples:
                break
            batch_indices = indices[start:end]
            batch = {key: arr[batch_indices] for key, arr in self.data.items()}
            yield batch


class TextDataset:
    """
    Prepare text corpus for language model training.
    
    Usage:
        dataset = TextDataset(tokenizer, texts, seq_length=128)
        loader = DataLoader(dataset.get_data(), batch_size=16)
    """
    
    def __init__(self, tokenizer, texts: List[str], seq_length: int = 128):
        self.seq_length = seq_length
        
        # Tokenize all texts and concatenate
        all_ids = []
        for text in texts:
            ids = tokenizer.encode(text, add_special=False)
            all_ids.extend(ids)
        
        self.all_ids = np.array(all_ids, dtype=np.int64)
        
        # Create input-target pairs
        n_sequences = max(0, (len(self.all_ids) - 1) // seq_length)
        self.n_sequences = n_sequences
    
    def get_data(self) -> Dict[str, np.ndarray]:
        """Get input/target pairs for training"""
        inputs = []
        targets = []
        
        for i in range(self.n_sequences):
            start = i * self.seq_length
            end = start + self.seq_length
            inputs.append(self.all_ids[start:end])
            targets.append(self.all_ids[start + 1:end + 1])
        
        if not inputs:
            # Return at least one sequence
            pad_len = self.seq_length
            padded = np.zeros(pad_len + 1, dtype=np.int64)
            padded[:min(len(self.all_ids), pad_len + 1)] = self.all_ids[:pad_len + 1]
            inputs = [padded[:pad_len]]
            targets = [padded[1:pad_len + 1]]
        
        return {
            'input': np.array(inputs, dtype=np.int64),
            'target': np.array(targets, dtype=np.int64),
        }


# =====================================================
# GRADIENT CLIPPING
# =====================================================

def clip_grad_norm(grads: List[np.ndarray], max_norm: float = 1.0) -> float:
    """
    Clip gradients by global norm.
    Returns the original norm before clipping.
    """
    total_norm = 0.0
    for g in grads:
        total_norm += np.sum(g ** 2)
    total_norm = math.sqrt(total_norm)
    
    if total_norm > max_norm:
        scale = max_norm / (total_norm + 1e-8)
        for g in grads:
            g *= scale
    
    return total_norm


# =====================================================
# MODEL CHECKPOINTING
# =====================================================

class Checkpoint:
    """
    Save and load model checkpoints.
    
    Saves model parameters, optimizer state, and training metadata.
    """
    
    @staticmethod
    def save(filepath: str, model, optimizer=None, epoch: int = 0,
             step: int = 0, loss: float = 0.0, metadata: Dict = None):
        """Save a checkpoint"""
        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else ".", exist_ok=True)
        
        checkpoint = {
            'params': [p.copy() for p in model.parameters()],
            'epoch': epoch,
            'step': step,
            'loss': loss,
            'metadata': metadata or {},
            'timestamp': time.time(),
        }
        
        if optimizer is not None:
            checkpoint['optimizer_state'] = {
                'lr': optimizer.lr,
                'step_count': optimizer.step_count,
            }
            if hasattr(optimizer, 'm'):
                checkpoint['optimizer_state']['m'] = [m.copy() for m in optimizer.m]
                checkpoint['optimizer_state']['v'] = [v.copy() for v in optimizer.v]
        
        np.savez(filepath, **{f'param_{i}': p for i, p in enumerate(checkpoint['params'])},
                 meta=np.array([epoch, step, loss]),
                 n_params=np.array([len(checkpoint['params'])]))
        
        # Save metadata as JSON
        meta_path = filepath + '.meta.json'
        meta = {
            'epoch': epoch,
            'step': step,
            'loss': loss,
            'timestamp': checkpoint['timestamp'],
            'metadata': metadata or {},
        }
        with open(meta_path, 'w') as f:
            json.dump(meta, f, indent=2)
    
    @staticmethod
    def load(filepath: str, model, optimizer=None) -> Dict:
        """Load a checkpoint into model"""
        data = np.load(filepath + '.npz', allow_pickle=True)
        n_params = int(data['n_params'][0])
        meta = data['meta']
        
        params = model.parameters()
        for i in range(min(n_params, len(params))):
            key = f'param_{i}'
            if key in data:
                params[i][:] = data[key]
        
        info = {
            'epoch': int(meta[0]),
            'step': int(meta[1]),
            'loss': float(meta[2]),
        }
        
        # Load JSON metadata if exists
        meta_path = filepath + '.meta.json'
        if os.path.exists(meta_path):
            with open(meta_path, 'r') as f:
                info['metadata'] = json.load(f)
        
        return info


# =====================================================
# TRAINING LOOP
# =====================================================

class Trainer:
    """
    Complete training loop for transformer models.
    
    Usage:
        trainer = Trainer(model, optimizer, scheduler=scheduler)
        history = trainer.train(train_loader, epochs=10, val_loader=val_loader)
    """
    
    def __init__(self, model, optimizer, scheduler=None,
                 max_grad_norm: float = 1.0,
                 checkpoint_dir: str = "checkpoints",
                 log_interval: int = 10):
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.max_grad_norm = max_grad_norm
        self.checkpoint_dir = checkpoint_dir
        self.log_interval = log_interval
        
        self.history: Dict[str, List[float]] = {
            'train_loss': [],
            'val_loss': [],
            'learning_rate': [],
            'grad_norm': [],
        }
        self.global_step = 0
        self.best_val_loss = float('inf')
    
    def train_step(self, batch: Dict[str, np.ndarray]) -> float:
        """Single training step"""
        input_ids = batch['input']
        target_ids = batch['target']
        
        # Forward
        logits = self.model.forward(input_ids)
        loss = self.model.compute_loss(logits, target_ids)
        
        # Backward
        self.model.backward(logits, target_ids)
        
        # Gradient clipping
        grads = self.model.gradients()
        grad_norm = clip_grad_norm(grads, self.max_grad_norm)
        
        # Optimizer step
        self.optimizer.step(grads)
        self.optimizer.zero_grad(grads)
        
        # Scheduler step
        if self.scheduler:
            self.scheduler.step()
        
        self.global_step += 1
        self.history['grad_norm'].append(grad_norm)
        self.history['learning_rate'].append(self.optimizer.lr)
        
        return loss
    
    def evaluate(self, val_loader: DataLoader) -> float:
        """Evaluate on validation set"""
        total_loss = 0.0
        n_batches = 0
        
        for batch in val_loader:
            logits = self.model.forward(batch['input'])
            loss = self.model.compute_loss(logits, batch['target'])
            total_loss += loss
            n_batches += 1
        
        return total_loss / max(1, n_batches)
    
    def train(self, train_loader: DataLoader, epochs: int = 10,
              val_loader: DataLoader = None,
              save_every: int = 0,
              early_stopping_patience: int = 0,
              callback: Callable = None) -> Dict[str, List[float]]:
        """
        Full training loop.
        
        Args:
            train_loader: Training data loader
            epochs: Number of training epochs
            val_loader: Optional validation loader
            save_every: Save checkpoint every N epochs (0 = disabled)
            early_stopping_patience: Stop if val_loss doesn't improve for N epochs (0 = disabled)
            callback: Optional callback fn(epoch, step, loss)
        
        Returns:
            Training history dict
        """
        patience_counter = 0
        
        print(f"Training started | {self.model.num_parameters():,} parameters")
        print(f"{'='*60}")
        
        for epoch in range(1, epochs + 1):
            epoch_loss = 0.0
            n_batches = 0
            epoch_start = time.time()
            
            for batch_idx, batch in enumerate(train_loader):
                loss = self.train_step(batch)
                epoch_loss += loss
                n_batches += 1
                
                if self.log_interval > 0 and self.global_step % self.log_interval == 0:
                    avg_loss = epoch_loss / n_batches
                    lr = self.optimizer.lr
                    print(f"  Step {self.global_step:>6d} | Loss: {loss:.4f} | Avg: {avg_loss:.4f} | LR: {lr:.2e}")
            
            # Epoch summary
            avg_train_loss = epoch_loss / max(1, n_batches)
            self.history['train_loss'].append(avg_train_loss)
            epoch_time = time.time() - epoch_start
            
            summary = f"Epoch {epoch}/{epochs} | Train Loss: {avg_train_loss:.4f} | Time: {epoch_time:.1f}s"
            
            # Validation
            if val_loader:
                val_loss = self.evaluate(val_loader)
                self.history['val_loss'].append(val_loss)
                summary += f" | Val Loss: {val_loss:.4f}"
                
                # Best model tracking
                if val_loss < self.best_val_loss:
                    self.best_val_loss = val_loss
                    patience_counter = 0
                    if self.checkpoint_dir:
                        Checkpoint.save(
                            os.path.join(self.checkpoint_dir, "best_model"),
                            self.model, self.optimizer, epoch, self.global_step, val_loss,
                            metadata={'type': 'best'}
                        )
                        summary += " *best*"
                else:
                    patience_counter += 1
            
            print(summary)
            
            # Callback
            if callback:
                callback(epoch, self.global_step, avg_train_loss)
            
            # Checkpointing
            if save_every > 0 and epoch % save_every == 0 and self.checkpoint_dir:
                Checkpoint.save(
                    os.path.join(self.checkpoint_dir, f"checkpoint_epoch_{epoch}"),
                    self.model, self.optimizer, epoch, self.global_step, avg_train_loss
                )
            
            # Early stopping
            if early_stopping_patience > 0 and patience_counter >= early_stopping_patience:
                print(f"Early stopping at epoch {epoch} (patience={early_stopping_patience})")
                break
        
        print(f"{'='*60}")
        print(f"Training complete | Final Loss: {self.history['train_loss'][-1]:.4f}")
        
        return self.history


# =====================================================
# CONVENIENCE: TRAIN A LANGUAGE MODEL
# =====================================================

def train_language_model(texts: List[str],
                         vocab_size: int = 1000,
                         d_model: int = 128,
                         num_heads: int = 4,
                         num_layers: int = 2,
                         seq_length: int = 64,
                         batch_size: int = 16,
                         epochs: int = 10,
                         lr: float = 3e-4,
                         checkpoint_dir: str = "checkpoints") -> Tuple:
    """
    High-level function to train a language model from scratch.
    
    Args:
        texts: List of training texts
        vocab_size: Target vocabulary size for BPE
        d_model: Model dimension
        num_heads: Number of attention heads
        num_layers: Number of transformer layers
        seq_length: Training sequence length
        batch_size: Training batch size
        epochs: Number of training epochs
        lr: Learning rate
        checkpoint_dir: Where to save checkpoints
    
    Returns:
        (model, tokenizer, history)
    
    Usage:
        from createsonline.ml.training import train_language_model
        
        texts = ["Your training data here...", "More text..."]
        model, tokenizer, history = train_language_model(texts, epochs=20)
        
        # Generate text
        prompt_ids = tokenizer.encode("Hello")
        output_ids = model.generate(np.array([prompt_ids]), max_new_tokens=50)
        print(tokenizer.decode(output_ids[0].tolist()))
    """
    from .transformer import TransformerLM
    from ..nlp.tokenizer import BPETokenizer
    
    print("Step 1: Training BPE tokenizer...")
    tokenizer = BPETokenizer(vocab_size=vocab_size)
    tokenizer.train(texts, verbose=True)
    
    print(f"\nStep 2: Building model (d={d_model}, heads={num_heads}, layers={num_layers})...")
    model = TransformerLM(
        vocab_size=tokenizer.vocab_size,
        d_model=d_model,
        num_heads=num_heads,
        num_layers=num_layers,
        max_seq_len=seq_length * 2,
    )
    print(f"  Parameters: {model.num_parameters():,}")
    
    print("\nStep 3: Preparing dataset...")
    dataset = TextDataset(tokenizer, texts, seq_length=seq_length)
    data = dataset.get_data()
    train_loader = DataLoader(data, batch_size=batch_size, shuffle=True)
    print(f"  Sequences: {len(data['input'])}")
    
    print(f"\nStep 4: Training for {epochs} epochs...")
    optimizer = AdamW(model.parameters(), lr=lr)
    total_steps = len(train_loader) * epochs
    scheduler = CosineScheduler(optimizer, total_steps=total_steps, warmup_steps=min(100, total_steps // 10))
    
    trainer = Trainer(model, optimizer, scheduler, checkpoint_dir=checkpoint_dir)
    history = trainer.train(train_loader, epochs=epochs, save_every=max(1, epochs // 5))
    
    # Save tokenizer alongside model
    if checkpoint_dir:
        tokenizer.save(os.path.join(checkpoint_dir, "tokenizer"))
    
    return model, tokenizer, history
