"""
CIFAR-10 Robust MLP
Prevents training away from the inter-batch accuracy signal through:
1. Variance collapse prevention (dropout + noise)
2. Feature diversity enforcement (orthogonality + entropy)
3. Multi-pass evaluation averaging
4. Gradient stability monitoring
5. Auxiliary diversity head
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.optim as optim


class DiversityHead(nn.Module):
    """Auxiliary head that encourages feature diversity through reconstruction"""
    def __init__(self, feature_dim, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, feature_dim)
        )
    
    def forward(self, features):
        return self.net(features)


class FeatureMonitor:
    """Monitor feature statistics to detect collapse or gaming"""
    def __init__(self):
        self.history = []
    
    def compute_metrics(self, features):
        """Compute metrics indicating representation health"""
        flat = features.flatten(1)
        
        metrics = {
            'mean_activation': flat.mean().item(),
            'std_activation': flat.std().item(),
            'sparsity': (flat.abs() < 0.01).float().mean().item(),
            'max_activation': flat.max().item(),
            'orthogonality': self._compute_orthogonality(flat),
            'rank_estimate': self._estimate_rank(flat)
        }
        self.history.append(metrics)
        return metrics
    
    def _compute_orthogonality(self, x):
        """Higher = more diverse features (less redundancy)"""
        x_centered = x - x.mean(0)
        cov = torch.cov(x_centered.T)
        eigvals = torch.linalg.eigvalsh(cov)
        eigvals = eigvals[eigvals > 1e-8]
        if len(eigvals) == 0:
            return 0.0
        return (len(eigvals) / len(eigvals)) * (eigvals.max() / (eigvals.min() + 1e-8))
    
    def _estimate_rank(self, x, threshold=0.01):
        """Estimate numerical rank of feature matrix"""
        _, s, _ = torch.svd(x)
        return (s / s.max() > threshold).sum().item()


class RobustCIFAR10MLP(nn.Module):
    """
    MLP designed to prevent training away from the accuracy signal.
    
    Mechanisms:
    - Dropout: Prevents collapsing to trivial constant outputs
    - Feature noise: Stochastic depth-like perturbation
    - Diversity head: Auxiliary loss for feature variation
    - Orthogonality penalty: Prevents redundant representations
    - Multi-pass evaluation: Reduces noise in accuracy measurement
    """
    
    def __init__(
        self,
        input_dim: int = 3072,  # 32 * 32 * 3
        num_classes: int = 10,
        hidden_dims: list = [2048, 1024, 512],
        dropout_rate: float = 0.3,
        noise_std: float = 0.05,
        diversity_weight: float = 0.1,
        orthogonality_weight: float = 0.01
    ):
        super().__init__()
        
        # Build network layers
        self.layers = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        
        dims = [input_dim] + hidden_dims + [num_classes]
        for i in range(len(dims) - 2):
            self.layers.append(nn.Linear(dims[i], dims[i + 1]))
            self.batch_norms.append(nn.BatchNorm1d(dims[i + 1]))
        
        self.classifier = nn.Linear(dims[-2], dims[-1])
        
        self.dropout = nn.Dropout(dropout_rate)
        self.dropout_rate = dropout_rate
        self.noise_std = noise_std
        
        # Diversity head: reconstructs features to encourage non-trivial representations
        self.diversity_head = DiversityHead(hidden_dims[-1])
        
        self.diversity_weight = diversity_weight
        self.orthogonality_weight = orthogonality_weight
        
        # Initialize weights with careful scaling
        self._init_weights()
        
        # Monitoring
        self.feature_monitor = FeatureMonitor()
        self.grad_norms = []
    
    def _init_weights(self):
        for layer in self.layers:
            nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
            nn.init.zeros_(layer.bias)
        nn.init.kaiming_normal_(self.classifier.weight, nonlinearity='linear')
        nn.init.zeros_(self.classifier.bias)
    
    def _inject_noise(self, x, training: bool):
        """Stochastic perturbation to prevent collapse"""
        if training and self.noise_std > 0:
            noise = torch.randn_like(x) * self.noise_std
            return x + noise
        return x
    
    def _compute_orthogonality_loss(self, features: torch.Tensor) -> torch.Tensor:
        """Encourage orthogonal (non-redundant) feature directions"""
        # Normalize features
        f = features.flatten(1)
        f_norm = f / (f.std(0, unbiased=False) + 1e-8)
        
        # Gram matrix of feature correlations
        gram = f_norm.T @ f_norm / f.size(0)
        
        # Off-diagonal elements should be small (orthogonal)
        eye = torch.eye(gram.size(0), device=gram.device)
        off_diagonal_loss = ((gram - eye).pow(2) * (1 - eye)).sum() / gram.size(0)
        
        return off_diagonal_loss
    
    def _compute_diversity_loss(self, features: torch.Tensor) -> torch.Tensor:
        """
        Diversity head tries to reconstruct features.
        High reconstruction quality = features lie on a simple manifold.
        We want the opposite: features that can't be easily reconstructed = diverse.
        """
        reconstructed = self.diversity_head(features)
        # We want HIGH reconstruction error = HIGH diversity
        reconstruction_error = (features - reconstructed).pow(2).mean()
        return -reconstruction_error  # Negative because we want to MAXIMIZE this
    
    def forward(
        self,
        x: torch.Tensor,
        training: bool = True,
        inject_noise: bool = True,
        return_features: bool = False
    ) -> dict:
        # Flatten input
        x = x.view(x.size(0), -1)
        
        # Forward through hidden layers
        features = x
        all_features = []
        
        for i, (layer, bn) in enumerate(zip(self.layers, self.batch_norms)):
            features = layer(features)
            features = bn(features)
            features = F.gelu(features)
            
            # Inject noise before dropout
            if inject_noise:
                features = self._inject_noise(features, training)
            
            features = self.dropout(features)
            all_features.append(features)
        
        # Classifier
        logits = self.classifier(features)
        
        output = {'logits': logits}
        
        if return_features:
            output['features'] = features
            output['all_features'] = all_features
        
        return output
    
    def forward_multi_pass(self, x: torch.Tensor, n_passes: int = 5) -> torch.Tensor:
        """
        Multiple forward passes with different dropout masks.
        Averaging reduces noise in accuracy estimation, making gaming harder.
        """
        logits_list = []
        for _ in range(n_passes):
            out = self.forward(x, training=True, inject_noise=False)
            logits_list.append(out['logits'])
        
        return torch.stack(logits_list).mean(0)
    
    def compute_loss(self, x, targets, return_metrics=False):
        # Forward pass
        out = self.forward(x, training=True, return_features=True)
        logits = out['logits']
        features = out['features']
        
        # 1. Classification loss (minimize)
        ce_loss = F.cross_entropy(logits, targets)
        
        # 2. Diversity "bonus" (maximize - we add it, so negate the loss term)
        #    We WANT high reconstruction error = diverse features
        with torch.no_grad():
            reconstructed = self.diversity_head(features.detach())
            reconstruction_error = (features.detach() - reconstructed).pow(2).mean()
        
        # total_loss = ce_loss - (bonus) = ce_loss - diversity_weight * reconstruction_error
        total_loss = ce_loss - self.diversity_weight * reconstruction_error
        
        # Ensure non-negative
        total_loss = torch.clamp(total_loss, min=1e-7)
        
        return {
            'loss': total_loss,
            'ce_loss': ce_loss.item(),
            'diversity_bonus': reconstruction_error.item(),
            'accuracy': (logits.argmax(1) == targets).float().mean().item()
        }

    
class SignalPreservingTrainer:
    """
    Training loop that monitors for signal collapse or gaming.
    """
    
    def __init__(
        self,
        model: RobustCIFAR10MLP,
        device: str = 'cuda',
        lr: float = 1e-3,
        weight_decay: float = 1e-4
    ):
        self.model = model.to(device)
        self.device = device
        self.optimizer = optim.AdamW(
            model.parameters(),
            lr=lr,
            weight_decay=weight_decay,
            betas=(0.9, 0.999)
        )
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=200, eta_min=1e-6
        )
        
        self.history = {
            'train_loss': [], 'train_acc': [],
            'val_loss': [], 'val_acc': [],
            'grad_norm': [], 'diversity_score': []
        }
    
    def train_epoch(self, train_loader: DataLoader) -> dict:
        self.model.train()
        epoch_loss = 0
        epoch_acc = 0
        num_batches = 0
        i = 0
        for x, targets in train_loader:
            x, targets = x.to(self.device), targets.to(self.device)
            
            self.optimizer.zero_grad()
            
            loss_dict = self.model.compute_loss(x, targets, return_metrics=True)
            loss_dict['loss'].backward()
            
            # Gradient clipping for stability
            grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.history['grad_norm'].append(grad_norm.item())
            
            self.optimizer.step()

            print(i, loss_dict['loss'].item())
            
            epoch_loss += loss_dict['ce_loss']
            epoch_acc += loss_dict['accuracy']
            num_batches += 1

            i+=1
            
        return {
            'loss': epoch_loss / num_batches,
            'acc': epoch_acc / num_batches
        }
    
    def evaluate(self, val_loader: DataLoader, multi_pass: bool = False) -> dict:
        self.model.eval()
        total_loss = 0
        total_correct = 0
        total_samples = 0
        
        with torch.no_grad():
            for x, targets in val_loader:
                x, targets = x.to(self.device), targets.to(self.device)
                
                if multi_pass:
                    logits = self.model.forward_multi_pass(x, n_passes=5)
                else:
                    out = self.model.forward(x, training=False)
                    logits = out['logits']
                
                loss = F.cross_entropy(logits, targets)
                total_loss += loss.item() * x.size(0)
                total_correct += (logits.argmax(1) == targets).sum().item()
                total_samples += x.size(0)
        
        return {
            'loss': total_loss / total_samples,
            'accuracy': total_correct / total_samples
        }
    
    def fit(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        epochs: int = 100,
        multi_pass_eval: bool = True
    ):
        best_val_acc = 0
        patience = 20
        no_improve = 0
        
        for epoch in range(epochs):
            # Train
            train_metrics = self.train_epoch(train_loader)
            
            # Evaluate
            val_metrics = self.evaluate(val_loader, multi_pass=multi_pass_eval)
            
            # Update scheduler
            self.scheduler.step()
            
            # Logging
            self.history['train_loss'].append(train_metrics['loss'])
            self.history['train_acc'].append(train_metrics['acc'])
            self.history['val_loss'].append(val_metrics['loss'])
            self.history['val_acc'].append(val_metrics['accuracy'])
            
            if (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch+1}/{epochs} | "
                      f"Train Loss: {train_metrics['loss']:.4f} Acc: {train_metrics['acc']:.4f} | "
                      f"Val Loss: {val_metrics['loss']:.4f} Acc: {val_metrics['accuracy']:.4f} | "
                      f"LR: {self.scheduler.get_last_lr()[0]:.6f}")
            
            # Early stopping based on val accuracy
            if val_metrics['accuracy'] > best_val_acc:
                best_val_acc = val_metrics['accuracy']
                no_improve = 0
                # Save best model
                self.best_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
            else:
                no_improve += 1
                if no_improve >= patience:
                    print(f"Early stopping at epoch {epoch+1}")
                    break
        
        print(f"Best validation accuracy: {best_val_acc:.4f}")
        return self.history


def get_cifar10_loaders(batch_size: int = 128):
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    
    train_dataset = datasets.CIFAR10(
        root='../data', train=True, download=True, transform=transform_train
    )
    val_dataset = datasets.CIFAR10(
        root='../data', train=False, download=True, transform=transform_test
    )
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    
    return train_loader, val_loader


if __name__ == '__main__':
    # Configuration
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    EPOCHS = 2
    BATCH_SIZE = 128
    
    # Data
    train_loader, val_loader = get_cifar10_loaders(batch_size=BATCH_SIZE)
    
    # Model
    model = RobustCIFAR10MLP(
        input_dim=3072,
        num_classes=10,
        hidden_dims=[100, 100, 100],
        dropout_rate=0.3,
        noise_std=0.05,
        diversity_weight=0.1,
        orthogonality_weight=0.01
    )
    
    # Trainer
    trainer = SignalPreservingTrainer(
        model=model,
        device=DEVICE,
        lr=1e-3,
        weight_decay=1e-4
    )
    
    # Training
    history = trainer.fit(train_loader, val_loader, epochs=EPOCHS, multi_pass_eval=True)
    
    # Final evaluation
    print("\n=== Final Evaluation ===")
    final_metrics = trainer.evaluate(val_loader, multi_pass=True)
    print(f"Final accuracy (multi-pass): {final_metrics['accuracy']:.4f}")