"""
PlainMLP vs ResMLP Comparison on Distant Identity Task (Final Version)

This experiment demonstrates the vanishing gradient problem in deep networks
and how residual connections solve it.

Key Design Choices:
1. PlainMLP: Standard x = ReLU(Linear(x)) - suffers from vanishing gradients
2. ResMLP: x = x + ReLU(Linear(x)) with zero-initialized bias and small weight scale
   - This allows the network to start as near-identity and learn deviations
   - Gradients can flow through the skip connection even when residual branch is small

The "Distant Identity" task (Y=X) is particularly revealing because:
- ResMLP can trivially solve it by zeroing the residual branch (identity shortcut)
- PlainMLP must learn a complex function composition to approximate identity
- With ReLU, PlainMLP can never perfectly learn identity (negative values are zeroed)
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple
import json
import os

# Set random seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)

# Configuration
NUM_LAYERS = 20
HIDDEN_DIM = 64
NUM_SAMPLES = 1024
TRAINING_STEPS = 500
LEARNING_RATE = 1e-3
BATCH_SIZE = 64

print(f"[Config] Layers: {NUM_LAYERS}, Hidden Dim: {HIDDEN_DIM}")
print(f"[Config] Samples: {NUM_SAMPLES}, Steps: {TRAINING_STEPS}, LR: {LEARNING_RATE}")


class PlainMLP(nn.Module):
    """Plain MLP: x = ReLU(Linear(x)) for each layer
    
    This architecture suffers from:
    1. Vanishing gradients - gradients must flow through all layers multiplicatively
    2. Information loss - ReLU zeros negative values at each layer
    3. Complex optimization - must learn exact function composition for identity
    """
    
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList()
        for _ in range(num_layers):
            layer = nn.Linear(dim, dim)
            # Kaiming He initialization
            nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
            nn.init.zeros_(layer.bias)
            self.layers.append(layer)
        self.activation = nn.ReLU()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = self.activation(layer(x))
        return x


class ResMLP(nn.Module):
    """Residual MLP: x = x + ReLU(Linear(x)) for each layer
    
    Key advantages:
    1. Identity shortcut - gradients flow directly to early layers via skip connection
    2. Residual learning - network learns deviation from identity, not full mapping
    3. For identity task - optimal solution is to zero the residual branch
    
    Uses small weight initialization (scaled by 1/sqrt(num_layers)) to:
    - Start near-identity behavior
    - Prevent activation explosion
    - Allow gradual learning of residuals
    """
    
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList()
        self.num_layers = num_layers
        
        for _ in range(num_layers):
            layer = nn.Linear(dim, dim)
            # Small initialization for residual branch
            # This ensures the network starts close to identity
            nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
            layer.weight.data *= 1.0 / np.sqrt(num_layers)  # Scale down weights
            nn.init.zeros_(layer.bias)
            self.layers.append(layer)
        self.activation = nn.ReLU()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = x + self.activation(layer(x))  # Residual connection
        return x


def generate_identity_data(num_samples: int, dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Generate synthetic data where Y = X, with X ~ U(-1, 1)"""
    X = torch.empty(num_samples, dim).uniform_(-1, 1)
    Y = X.clone()
    return X, Y


def train_model(model: nn.Module, X: torch.Tensor, Y: torch.Tensor, 
                steps: int, lr: float, batch_size: int) -> List[float]:
    """Train model and record loss at each step"""
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    losses = []
    
    num_samples = X.shape[0]
    
    for step in range(steps):
        # Random batch sampling
        indices = torch.randint(0, num_samples, (batch_size,))
        batch_x = X[indices]
        batch_y = Y[indices]
        
        # Forward pass
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        losses.append(loss.item())
        
        if step % 100 == 0:
            print(f"  Step {step}/{steps}, Loss: {loss.item():.6f}")
    
    return losses


class ActivationGradientHook:
    """Hook to capture activations and gradients at each layer"""
    
    def __init__(self):
        self.activations: List[torch.Tensor] = []
        self.gradients: List[torch.Tensor] = []
        self.handles = []
    
    def register_hooks(self, model: nn.Module):
        """Register forward and backward hooks on each layer"""
        for layer in model.layers:
            handle_fwd = layer.register_forward_hook(self._forward_hook)
            handle_bwd = layer.register_full_backward_hook(self._backward_hook)
            self.handles.extend([handle_fwd, handle_bwd])
    
    def _forward_hook(self, module, input, output):
        self.activations.append(output.detach().clone())
    
    def _backward_hook(self, module, grad_input, grad_output):
        self.gradients.append(grad_output[0].detach().clone())
    
    def clear(self):
        self.activations = []
        self.gradients = []
    
    def remove_hooks(self):
        for handle in self.handles:
            handle.remove()
        self.handles = []
    
    def get_activation_stats(self) -> Tuple[List[float], List[float]]:
        """Get mean and std of activations for each layer"""
        means = [act.mean().item() for act in self.activations]
        stds = [act.std().item() for act in self.activations]
        return means, stds
    
    def get_gradient_norms(self) -> List[float]:
        """Get L2 norm of gradients for each layer (in forward order)"""
        norms = [grad.norm(2).item() for grad in reversed(self.gradients)]
        return norms


def analyze_final_state(model: nn.Module, dim: int, batch_size: int = 64) -> Dict:
    """Perform forward/backward pass and capture activation/gradient stats"""
    hook = ActivationGradientHook()
    hook.register_hooks(model)
    
    # Generate new random batch
    X_test = torch.empty(batch_size, dim).uniform_(-1, 1)
    Y_test = X_test.clone()
    
    # Forward pass
    model.zero_grad()
    output = model(X_test)
    loss = nn.MSELoss()(output, Y_test)
    
    # Backward pass
    loss.backward()
    
    # Get statistics
    act_means, act_stds = hook.get_activation_stats()
    grad_norms = hook.get_gradient_norms()
    
    hook.remove_hooks()
    
    return {
        'activation_means': act_means,
        'activation_stds': act_stds,
        'gradient_norms': grad_norms,
        'final_loss': loss.item()
    }


def plot_training_loss(plain_losses: List[float], res_losses: List[float], save_path: str):
    """Plot training loss curves for both models"""
    fig, ax = plt.subplots(figsize=(10, 6))
    steps = range(len(plain_losses))
    
    ax.plot(steps, plain_losses, label='PlainMLP (20 layers)', color='#e74c3c', 
            alpha=0.8, linewidth=2)
    ax.plot(steps, res_losses, label='ResMLP (20 layers)', color='#3498db', 
            alpha=0.8, linewidth=2)
    
    ax.set_xlabel('Training Steps', fontsize=12)
    ax.set_ylabel('MSE Loss', fontsize=12)
    ax.set_title('Training Loss: PlainMLP vs ResMLP on Identity Task (Y = X)', fontsize=14)
    ax.legend(fontsize=11, loc='upper right')
    ax.grid(True, alpha=0.3)
    ax.set_yscale('log')
    
    # Add final loss annotations
    final_plain = plain_losses[-1]
    final_res = res_losses[-1]
    
    # Text box with final results
    textstr = f'Final Loss:\n  PlainMLP: {final_plain:.4f}\n  ResMLP: {final_res:.4f}\n  Improvement: {final_plain/final_res:.1f}x'
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
    ax.text(0.02, 0.02, textstr, transform=ax.transAxes, fontsize=10,
            verticalalignment='bottom', bbox=props)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"[Plot] Saved training loss plot to {save_path}")


def plot_gradient_magnitudes(plain_grads: List[float], res_grads: List[float], save_path: str):
    """Plot gradient magnitude vs layer depth"""
    fig, ax = plt.subplots(figsize=(10, 6))
    layers = range(1, len(plain_grads) + 1)
    
    ax.plot(layers, plain_grads, 'o-', label='PlainMLP', color='#e74c3c', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    ax.plot(layers, res_grads, 's-', label='ResMLP', color='#3498db', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    
    ax.set_xlabel('Layer Depth (1 = first layer, 20 = last layer)', fontsize=12)
    ax.set_ylabel('Gradient L2 Norm (log scale)', fontsize=12)
    ax.set_title('Gradient Magnitude vs Layer Depth (After 500 Training Steps)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_yscale('log')
    
    # Highlight the gradient difference
    ax.fill_between(layers, plain_grads, res_grads, alpha=0.15, color='gray')
    
    # Add annotation about gradient flow
    ax.annotate('Gradients flow more\nuniformly in ResMLP', 
                xy=(10, res_grads[9]), xytext=(5, res_grads[9]*5),
                fontsize=10, color='#3498db',
                arrowprops=dict(arrowstyle='->', color='#3498db', alpha=0.7))
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"[Plot] Saved gradient magnitude plot to {save_path}")


def plot_activation_means(plain_means: List[float], res_means: List[float], save_path: str):
    """Plot activation mean vs layer depth"""
    fig, ax = plt.subplots(figsize=(10, 6))
    layers = range(1, len(plain_means) + 1)
    
    ax.plot(layers, plain_means, 'o-', label='PlainMLP', color='#e74c3c', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    ax.plot(layers, res_means, 's-', label='ResMLP', color='#3498db', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    
    ax.axhline(y=0, color='gray', linestyle='--', alpha=0.5, linewidth=1)
    
    ax.set_xlabel('Layer Depth', fontsize=12)
    ax.set_ylabel('Activation Mean', fontsize=12)
    ax.set_title('Activation Mean vs Layer Depth (After Training)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"[Plot] Saved activation mean plot to {save_path}")


def plot_activation_stds(plain_stds: List[float], res_stds: List[float], save_path: str):
    """Plot activation std vs layer depth"""
    fig, ax = plt.subplots(figsize=(10, 6))
    layers = range(1, len(plain_stds) + 1)
    
    ax.plot(layers, plain_stds, 'o-', label='PlainMLP', color='#e74c3c', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    ax.plot(layers, res_stds, 's-', label='ResMLP', color='#3498db', 
            markersize=8, linewidth=2, markeredgecolor='white', markeredgewidth=1)
    
    ax.set_xlabel('Layer Depth', fontsize=12)
    ax.set_ylabel('Activation Standard Deviation', fontsize=12)
    ax.set_title('Activation Std vs Layer Depth (After Training)', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    
    # Add annotation about signal preservation
    ax.annotate('ResMLP maintains\nstable activations', 
                xy=(15, res_stds[14]), xytext=(10, res_stds[14]*1.3),
                fontsize=10, color='#3498db',
                arrowprops=dict(arrowstyle='->', color='#3498db', alpha=0.7))
    
    ax.annotate('PlainMLP activations\ndegrade through layers', 
                xy=(18, plain_stds[17]), xytext=(12, plain_stds[17]*0.5),
                fontsize=10, color='#e74c3c',
                arrowprops=dict(arrowstyle='->', color='#e74c3c', alpha=0.7))
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"[Plot] Saved activation std plot to {save_path}")


def main():
    print("=" * 60)
    print("PlainMLP vs ResMLP: Distant Identity Task Experiment")
    print("=" * 60)
    
    # Ensure plots directory exists
    os.makedirs('plots', exist_ok=True)
    
    # Generate synthetic data
    print("\n[1] Generating synthetic identity data...")
    X, Y = generate_identity_data(NUM_SAMPLES, HIDDEN_DIM)
    print(f"  Data shape: X={X.shape}, Y={Y.shape}")
    print(f"  X range: [{X.min():.3f}, {X.max():.3f}]")
    print(f"  Task: Learn Y = X (identity mapping)")
    
    # Initialize models
    print("\n[2] Initializing models...")
    plain_mlp = PlainMLP(HIDDEN_DIM, NUM_LAYERS)
    res_mlp = ResMLP(HIDDEN_DIM, NUM_LAYERS)
    
    plain_params = sum(p.numel() for p in plain_mlp.parameters())
    res_params = sum(p.numel() for p in res_mlp.parameters())
    print(f"  PlainMLP parameters: {plain_params:,}")
    print(f"  ResMLP parameters: {res_params:,}")
    
    # Train PlainMLP
    print("\n[3] Training PlainMLP...")
    plain_losses = train_model(plain_mlp, X, Y, TRAINING_STEPS, LEARNING_RATE, BATCH_SIZE)
    print(f"  Final loss: {plain_losses[-1]:.6f}")
    
    # Train ResMLP
    print("\n[4] Training ResMLP...")
    res_losses = train_model(res_mlp, X, Y, TRAINING_STEPS, LEARNING_RATE, BATCH_SIZE)
    print(f"  Final loss: {res_losses[-1]:.6f}")
    
    # Calculate improvement
    improvement = plain_losses[-1] / res_losses[-1]
    print(f"\n  >>> ResMLP achieves {improvement:.1f}x lower loss than PlainMLP <<<")
    
    # Final state analysis
    print("\n[5] Analyzing final state of trained models...")
    print("  Running forward/backward pass on new random batch...")
    print("  Analyzing PlainMLP...")
    plain_stats = analyze_final_state(plain_mlp, HIDDEN_DIM)
    print("  Analyzing ResMLP...")
    res_stats = analyze_final_state(res_mlp, HIDDEN_DIM)
    
    # Print detailed analysis
    print("\n[6] Detailed Analysis:")
    print("\n  === Loss Comparison ===")
    print(f"  PlainMLP - Initial: {plain_losses[0]:.4f}, Final: {plain_losses[-1]:.4f}")
    print(f"  ResMLP   - Initial: {res_losses[0]:.4f}, Final: {res_losses[-1]:.4f}")
    
    print("\n  === Gradient Flow (L2 norms) ===")
    print(f"  PlainMLP - Layer 1: {plain_stats['gradient_norms'][0]:.2e}, Layer 20: {plain_stats['gradient_norms'][-1]:.2e}")
    print(f"  ResMLP   - Layer 1: {res_stats['gradient_norms'][0]:.2e}, Layer 20: {res_stats['gradient_norms'][-1]:.2e}")
    
    print("\n  === Activation Statistics ===")
    print(f"  PlainMLP - Std range: [{min(plain_stats['activation_stds']):.4f}, {max(plain_stats['activation_stds']):.4f}]")
    print(f"  ResMLP   - Std range: [{min(res_stats['activation_stds']):.4f}, {max(res_stats['activation_stds']):.4f}]")
    
    # Generate plots
    print("\n[7] Generating plots...")
    plot_training_loss(plain_losses, res_losses, 'plots/training_loss.png')
    plot_gradient_magnitudes(plain_stats['gradient_norms'], res_stats['gradient_norms'], 
                            'plots/gradient_magnitude.png')
    plot_activation_means(plain_stats['activation_means'], res_stats['activation_means'],
                         'plots/activation_mean.png')
    plot_activation_stds(plain_stats['activation_stds'], res_stats['activation_stds'],
                        'plots/activation_std.png')
    
    # Save results to JSON
    results = {
        'config': {
            'num_layers': NUM_LAYERS,
            'hidden_dim': HIDDEN_DIM,
            'num_samples': NUM_SAMPLES,
            'training_steps': TRAINING_STEPS,
            'learning_rate': LEARNING_RATE,
            'batch_size': BATCH_SIZE
        },
        'plain_mlp': {
            'final_loss': plain_losses[-1],
            'initial_loss': plain_losses[0],
            'loss_history': plain_losses,
            'gradient_norms': plain_stats['gradient_norms'],
            'activation_means': plain_stats['activation_means'],
            'activation_stds': plain_stats['activation_stds']
        },
        'res_mlp': {
            'final_loss': res_losses[-1],
            'initial_loss': res_losses[0],
            'loss_history': res_losses,
            'gradient_norms': res_stats['gradient_norms'],
            'activation_means': res_stats['activation_means'],
            'activation_stds': res_stats['activation_stds']
        },
        'summary': {
            'loss_improvement': improvement,
            'plain_grad_range': [min(plain_stats['gradient_norms']), max(plain_stats['gradient_norms'])],
            'res_grad_range': [min(res_stats['gradient_norms']), max(res_stats['gradient_norms'])],
            'plain_std_range': [min(plain_stats['activation_stds']), max(plain_stats['activation_stds'])],
            'res_std_range': [min(res_stats['activation_stds']), max(res_stats['activation_stds'])]
        }
    }
    
    with open('results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print("\n[8] Results saved to results.json")
    
    print("\n" + "=" * 60)
    print("Experiment completed successfully!")
    print("=" * 60)
    
    return results


if __name__ == "__main__":
    results = main()