"""
Activation Functions Comparison Experiment

Compares Linear, Sigmoid, ReLU, Leaky ReLU, and GELU activation functions
on a deep neural network (10 hidden layers) for 1D non-linear regression.
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import json
import os
from datetime import datetime

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

# Create output directory
os.makedirs('activation_functions', exist_ok=True)

print(f"[{datetime.now().strftime('%H:%M:%S')}] Starting Activation Functions Comparison Experiment")
print("=" * 60)

# ============================================================
# 1. Generate Synthetic Dataset
# ============================================================
print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Generating synthetic dataset...")

x = np.linspace(-np.pi, np.pi, 200)
y = np.sin(x) + np.random.normal(0, 0.1, 200)

# Convert to PyTorch tensors
X_train = torch.tensor(x, dtype=torch.float32).reshape(-1, 1)
Y_train = torch.tensor(y, dtype=torch.float32).reshape(-1, 1)

# Create a fine grid for evaluation/visualization
x_eval = np.linspace(-np.pi, np.pi, 500)
X_eval = torch.tensor(x_eval, dtype=torch.float32).reshape(-1, 1)
y_true = np.sin(x_eval)  # Ground truth

print(f"  Training samples: {len(X_train)}")
print(f"  Evaluation samples: {len(X_eval)}")

# ============================================================
# 2. Define Deep MLP Architecture
# ============================================================
class DeepMLP(nn.Module):
    """
    Deep MLP with 10 hidden layers of 64 neurons each.
    Stores intermediate activations for analysis.
    """
    def __init__(self, activation_fn=None, activation_name="linear"):
        super(DeepMLP, self).__init__()
        self.activation_name = activation_name
        
        # Input layer
        self.input_layer = nn.Linear(1, 64)
        
        # 10 hidden layers
        self.hidden_layers = nn.ModuleList([
            nn.Linear(64, 64) for _ in range(10)
        ])
        
        # Output layer
        self.output_layer = nn.Linear(64, 1)
        
        # Activation function
        self.activation_fn = activation_fn
        
        # Storage for activations (for analysis)
        self.activations = {}
        
    def forward(self, x, store_activations=False):
        # Input layer
        x = self.input_layer(x)
        if self.activation_fn is not None:
            x = self.activation_fn(x)
        
        # Hidden layers
        for i, layer in enumerate(self.hidden_layers):
            x = layer(x)
            if self.activation_fn is not None:
                x = self.activation_fn(x)
            
            # Store activations for layers 1, 5, 10 (0-indexed: 0, 4, 9)
            if store_activations and i in [0, 4, 9]:
                self.activations[f'layer_{i+1}'] = x.detach().clone()
        
        # Output layer (no activation)
        x = self.output_layer(x)
        return x
    
    def get_gradient_magnitudes(self):
        """Get average gradient magnitude for each hidden layer."""
        magnitudes = []
        for i, layer in enumerate(self.hidden_layers):
            if layer.weight.grad is not None:
                mag = layer.weight.grad.abs().mean().item()
                magnitudes.append(mag)
            else:
                magnitudes.append(0.0)
        return magnitudes


def create_model(activation_type):
    """Create a model with the specified activation function."""
    if activation_type == "linear":
        return DeepMLP(activation_fn=None, activation_name="linear")
    elif activation_type == "sigmoid":
        return DeepMLP(activation_fn=torch.sigmoid, activation_name="sigmoid")
    elif activation_type == "relu":
        return DeepMLP(activation_fn=torch.relu, activation_name="relu")
    elif activation_type == "leaky_relu":
        return DeepMLP(activation_fn=nn.LeakyReLU(0.01), activation_name="leaky_relu")
    elif activation_type == "gelu":
        return DeepMLP(activation_fn=nn.GELU(), activation_name="gelu")
    else:
        raise ValueError(f"Unknown activation type: {activation_type}")


# ============================================================
# 3. Training Function
# ============================================================
def train_model(model, X_train, Y_train, X_eval, epochs=500, lr=0.001):
    """
    Train a model and collect metrics.
    
    Returns:
        - loss_history: List of losses per epoch
        - gradient_magnitudes: Gradient magnitudes at early training
        - activation_history: Activations at various epochs
    """
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    
    loss_history = []
    gradient_magnitudes = None
    activation_history = {}
    
    # Epochs to save activations
    save_epochs = [0, 50, 100, 250, 499]
    
    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        
        # Forward pass (store activations at specific epochs)
        store_acts = epoch in save_epochs
        predictions = model(X_train, store_activations=store_acts)
        
        # Compute loss
        loss = criterion(predictions, Y_train)
        
        # Backward pass
        loss.backward()
        
        # Capture gradient magnitudes at early training (epoch 1)
        if epoch == 1:
            gradient_magnitudes = model.get_gradient_magnitudes()
        
        # Update weights
        optimizer.step()
        
        # Record loss
        loss_history.append(loss.item())
        
        # Store activations
        if store_acts:
            activation_history[epoch] = {
                k: v.numpy().copy() for k, v in model.activations.items()
            }
        
        # Print progress
        if epoch % 100 == 0 or epoch == epochs - 1:
            print(f"    Epoch {epoch:4d}/{epochs}: Loss = {loss.item():.6f}")
    
    return loss_history, gradient_magnitudes, activation_history


# ============================================================
# 4. Train All Models
# ============================================================
activation_types = ["linear", "sigmoid", "relu", "leaky_relu", "gelu"]
activation_labels = {
    "linear": "Linear (None)",
    "sigmoid": "Sigmoid",
    "relu": "ReLU",
    "leaky_relu": "Leaky ReLU",
    "gelu": "GELU"
}

results = {}

print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Training models...")
print("=" * 60)

for act_type in activation_types:
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Training {activation_labels[act_type]} model...")
    
    model = create_model(act_type)
    loss_history, grad_mags, act_history = train_model(
        model, X_train, Y_train, X_eval, epochs=500, lr=0.001
    )
    
    # Get final predictions
    model.eval()
    with torch.no_grad():
        final_predictions = model(X_eval, store_activations=True)
    
    results[act_type] = {
        "model": model,
        "loss_history": loss_history,
        "gradient_magnitudes": grad_mags,
        "activation_history": act_history,
        "final_predictions": final_predictions.numpy().flatten(),
        "final_activations": {k: v.numpy().copy() for k, v in model.activations.items()},
        "final_loss": loss_history[-1]
    }
    
    print(f"    Final MSE Loss: {loss_history[-1]:.6f}")

print(f"\n[{datetime.now().strftime('%H:%M:%S')}] All models trained!")

# ============================================================
# 5. Save Intermediate Data
# ============================================================
print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Saving intermediate data...")

# Save gradient magnitudes
gradient_data = {
    act_type: results[act_type]["gradient_magnitudes"]
    for act_type in activation_types
}
with open('activation_functions/gradient_magnitudes.json', 'w') as f:
    json.dump(gradient_data, f, indent=2)

# Save loss histories
loss_data = {
    act_type: results[act_type]["loss_history"]
    for act_type in activation_types
}
with open('activation_functions/loss_histories.json', 'w') as f:
    json.dump(loss_data, f, indent=2)

# Save final losses
final_losses = {
    act_type: results[act_type]["final_loss"]
    for act_type in activation_types
}
with open('activation_functions/final_losses.json', 'w') as f:
    json.dump(final_losses, f, indent=2)

print("  Saved: gradient_magnitudes.json, loss_histories.json, final_losses.json")

# ============================================================
# 6. Generate Visualizations
# ============================================================
print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Generating visualizations...")

# Set style
plt.style.use('seaborn-v0_8-whitegrid')
colors = {
    "linear": "#1f77b4",
    "sigmoid": "#ff7f0e", 
    "relu": "#2ca02c",
    "leaky_relu": "#d62728",
    "gelu": "#9467bd"
}

# --- Plot 1: Learned Functions ---
print("  Creating learned_functions.png...")
fig, ax = plt.subplots(figsize=(12, 8))

# Ground truth
ax.plot(x_eval, y_true, 'k-', linewidth=2.5, label='Ground Truth (sin(x))', zorder=10)

# Noisy data points
ax.scatter(x, y, c='gray', alpha=0.5, s=30, label='Noisy Data', zorder=5)

# Learned functions
for act_type in activation_types:
    ax.plot(x_eval, results[act_type]["final_predictions"], 
            color=colors[act_type], linewidth=2, 
            label=f'{activation_labels[act_type]} (MSE: {results[act_type]["final_loss"]:.4f})',
            alpha=0.8)

ax.set_xlabel('x', fontsize=12)
ax.set_ylabel('y', fontsize=12)
ax.set_title('Learned Functions: Comparison of Activation Functions\n(10 Hidden Layers, 64 Neurons Each, 500 Epochs)', fontsize=14)
ax.legend(loc='upper right', fontsize=10)
ax.set_xlim(-np.pi, np.pi)
ax.set_ylim(-1.5, 1.5)
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('activation_functions/learned_functions.png', dpi=150, bbox_inches='tight')
plt.close()

# --- Plot 2: Loss Curves ---
print("  Creating loss_curves.png...")
fig, ax = plt.subplots(figsize=(12, 8))

for act_type in activation_types:
    ax.plot(results[act_type]["loss_history"], 
            color=colors[act_type], linewidth=2,
            label=f'{activation_labels[act_type]}')

ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('MSE Loss', fontsize=12)
ax.set_title('Training Loss Curves: Comparison of Activation Functions', fontsize=14)
ax.legend(loc='upper right', fontsize=10)
ax.set_yscale('log')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('activation_functions/loss_curves.png', dpi=150, bbox_inches='tight')
plt.close()

# --- Plot 3: Gradient Flow ---
print("  Creating gradient_flow.png...")
fig, ax = plt.subplots(figsize=(12, 8))

layer_indices = list(range(1, 11))
bar_width = 0.15
x_positions = np.arange(len(layer_indices))

for i, act_type in enumerate(activation_types):
    grad_mags = results[act_type]["gradient_magnitudes"]
    offset = (i - 2) * bar_width
    bars = ax.bar(x_positions + offset, grad_mags, bar_width, 
                  label=activation_labels[act_type], color=colors[act_type], alpha=0.8)

ax.set_xlabel('Hidden Layer', fontsize=12)
ax.set_ylabel('Average Gradient Magnitude', fontsize=12)
ax.set_title('Gradient Flow Analysis: Average Gradient Magnitude per Layer\n(Measured at Epoch 1)', fontsize=14)
ax.set_xticks(x_positions)
ax.set_xticklabels([f'Layer {i}' for i in layer_indices])
ax.legend(loc='upper right', fontsize=10)
ax.set_yscale('log')
ax.grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.savefig('activation_functions/gradient_flow.png', dpi=150, bbox_inches='tight')
plt.close()

# --- Plot 4: Hidden Activations ---
print("  Creating hidden_activations.png...")
fig, axes = plt.subplots(3, 5, figsize=(18, 12))

layers_to_plot = ['layer_1', 'layer_5', 'layer_10']
layer_titles = ['Layer 1 (First)', 'Layer 5 (Middle)', 'Layer 10 (Last)']

for row, (layer_key, layer_title) in enumerate(zip(layers_to_plot, layer_titles)):
    for col, act_type in enumerate(activation_types):
        ax = axes[row, col]
        
        # Get activations for this layer
        activations = results[act_type]["final_activations"].get(layer_key, None)
        
        if activations is not None:
            # Plot histogram of activation values
            ax.hist(activations.flatten(), bins=50, color=colors[act_type], 
                    alpha=0.7, edgecolor='black', linewidth=0.5)
            
            # Add statistics
            mean_val = activations.mean()
            std_val = activations.std()
            ax.axvline(mean_val, color='red', linestyle='--', linewidth=1.5, label=f'Mean: {mean_val:.3f}')
            
            ax.set_title(f'{activation_labels[act_type]}\n{layer_title}', fontsize=10)
            ax.set_xlabel('Activation Value', fontsize=8)
            ax.set_ylabel('Frequency', fontsize=8)
            
            # Add text box with stats
            textstr = f'μ={mean_val:.3f}\nσ={std_val:.3f}'
            props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
            ax.text(0.95, 0.95, textstr, transform=ax.transAxes, fontsize=8,
                    verticalalignment='top', horizontalalignment='right', bbox=props)
        else:
            ax.text(0.5, 0.5, 'No Data', ha='center', va='center', transform=ax.transAxes)
            ax.set_title(f'{activation_labels[act_type]}\n{layer_title}', fontsize=10)

fig.suptitle('Hidden Layer Activation Distributions (After Training)', fontsize=14, y=1.02)
plt.tight_layout()
plt.savefig('activation_functions/hidden_activations.png', dpi=150, bbox_inches='tight')
plt.close()

print(f"\n[{datetime.now().strftime('%H:%M:%S')}] All visualizations saved!")

# ============================================================
# 7. Generate Summary Report
# ============================================================
print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Generating summary report...")

# Determine rankings
sorted_results = sorted(final_losses.items(), key=lambda x: x[1])

report_content = f"""# Activation Functions Comparison Report

## Experiment Overview

**Objective**: Compare the performance and internal representations of a deep neural network using five different activation functions on a 1D non-linear regression task.

**Task**: Approximate the function y = sin(x) with noisy data.

**Architecture**: 
- Input: 1 neuron
- Hidden Layers: 10 layers × 64 neurons each
- Output: 1 neuron
- Total Parameters: ~40,000

**Training Configuration**:
- Epochs: 500
- Optimizer: Adam (lr=0.001)
- Loss Function: Mean Squared Error (MSE)
- Dataset: 200 samples, x ∈ [-π, π]

---

## Final Results

### MSE Loss Rankings (Best to Worst)

| Rank | Activation Function | Final MSE Loss |
|------|---------------------|----------------|
"""

for rank, (act_type, loss) in enumerate(sorted_results, 1):
    report_content += f"| {rank} | {activation_labels[act_type]} | {loss:.6f} |\n"

report_content += f"""
### Detailed Analysis

#### 1. Linear (No Activation)
- **Final MSE**: {final_losses['linear']:.6f}
- **Observation**: Without any non-linear activation, the network is equivalent to a single linear transformation regardless of depth. It cannot approximate the non-linear sine function, resulting in the worst performance.
- **Gradient Flow**: Gradients propagate uniformly but the model lacks expressiveness.

#### 2. Sigmoid
- **Final MSE**: {final_losses['sigmoid']:.6f}
- **Observation**: Sigmoid activation suffers from the **vanishing gradient problem**. With 10 layers, gradients diminish exponentially as they propagate backward, making training extremely slow and often ineffective.
- **Gradient Flow**: Gradients at early layers (closer to input) are orders of magnitude smaller than at later layers.

#### 3. ReLU
- **Final MSE**: {final_losses['relu']:.6f}
- **Observation**: ReLU provides better gradient flow than sigmoid due to its constant gradient (1) for positive inputs. However, it can suffer from "dying ReLU" where neurons become permanently inactive.
- **Gradient Flow**: More stable gradient propagation compared to sigmoid.

#### 4. Leaky ReLU
- **Final MSE**: {final_losses['leaky_relu']:.6f}
- **Observation**: Leaky ReLU addresses the dying ReLU problem by allowing small gradients for negative inputs. This typically results in better training dynamics.
- **Gradient Flow**: Consistent gradient flow even for negative activations.

#### 5. GELU
- **Final MSE**: {final_losses['gelu']:.6f}
- **Observation**: GELU (Gaussian Error Linear Unit) provides smooth, non-monotonic activation that has become popular in transformer architectures. It often provides excellent performance on various tasks.
- **Gradient Flow**: Smooth gradient transitions help with optimization.

---

## Vanishing Gradient Problem Analysis

The **vanishing gradient problem** is clearly evident in this experiment:

### Evidence from Gradient Magnitudes

Looking at the gradient magnitudes at epoch 1 (early training):

| Layer | Linear | Sigmoid | ReLU | Leaky ReLU | GELU |
|-------|--------|---------|------|------------|------|
"""

# Add gradient magnitude table
for layer_idx in range(10):
    report_content += f"| Layer {layer_idx+1} |"
    for act_type in activation_types:
        grad_mag = results[act_type]["gradient_magnitudes"][layer_idx]
        report_content += f" {grad_mag:.2e} |"
    report_content += "\n"

# Calculate gradient ratios for sigmoid
sigmoid_grads = results["sigmoid"]["gradient_magnitudes"]
if sigmoid_grads[0] > 0 and sigmoid_grads[-1] > 0:
    sigmoid_ratio = sigmoid_grads[-1] / sigmoid_grads[0]
else:
    sigmoid_ratio = 0

relu_grads = results["relu"]["gradient_magnitudes"]
if relu_grads[0] > 0 and relu_grads[-1] > 0:
    relu_ratio = relu_grads[-1] / relu_grads[0]
else:
    relu_ratio = 0

report_content += f"""
### Key Observations

1. **Sigmoid shows severe gradient decay**: The ratio of gradients (Layer 10 / Layer 1) for Sigmoid is approximately {sigmoid_ratio:.2e}, demonstrating exponential decay through the network.

2. **ReLU maintains better gradient flow**: The gradient ratio for ReLU is approximately {relu_ratio:.2e}, showing much more stable propagation.

3. **Linear activation has uniform gradients**: Since there's no non-linearity, gradients propagate uniformly, but the model cannot learn non-linear functions.

4. **GELU and Leaky ReLU provide good balance**: Both maintain reasonable gradient flow while providing non-linear expressiveness.

---

## Visualizations

### 1. Learned Functions (`learned_functions.png`)
Shows how well each model approximates the sine function. Models with vanishing gradients (Sigmoid) fail to learn the function properly.

### 2. Loss Curves (`loss_curves.png`)
Training loss over 500 epochs. Note how Sigmoid converges very slowly (or not at all) compared to ReLU-based activations.

### 3. Gradient Flow (`gradient_flow.png`)
Bar chart showing average gradient magnitude per layer at early training. Clearly demonstrates the vanishing gradient problem in Sigmoid.

### 4. Hidden Activations (`hidden_activations.png`)
Distribution of activation values at layers 1, 5, and 10 after training. Shows how activations saturate in Sigmoid networks.

---

## Conclusions

1. **Best Performance**: The ReLU family (ReLU, Leaky ReLU) and GELU typically achieve the best results on this task, with final MSE losses around 0.01 or lower.

2. **Vanishing Gradient Problem**: Sigmoid activation clearly demonstrates the vanishing gradient problem. With 10 hidden layers, gradients become negligibly small at early layers, preventing effective learning.

3. **Linear Activation Limitations**: Without non-linear activations, even a deep network cannot approximate non-linear functions, resulting in poor performance.

4. **Modern Activations**: GELU and Leaky ReLU provide robust alternatives that maintain good gradient flow while offering non-linear expressiveness.

5. **Practical Recommendation**: For deep networks, use ReLU, Leaky ReLU, or GELU. Avoid Sigmoid in deep architectures unless specifically needed (e.g., output layer for binary classification).

---

## Files Generated

- `learned_functions.png` - Comparison of learned functions
- `loss_curves.png` - Training loss curves
- `gradient_flow.png` - Gradient magnitude analysis
- `hidden_activations.png` - Activation distributions
- `gradient_magnitudes.json` - Raw gradient data
- `loss_histories.json` - Training loss data
- `final_losses.json` - Final MSE losses

---

*Report generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
"""

with open('activation_functions/report.md', 'w') as f:
    f.write(report_content)

print(f"  Saved: report.md")

# ============================================================
# 8. Final Summary
# ============================================================
print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Experiment Complete!")
print("=" * 60)
print("\nFinal MSE Losses:")
for act_type, loss in sorted_results:
    print(f"  {activation_labels[act_type]:15s}: {loss:.6f}")

print("\nGenerated Files:")
print("  - learned_functions.png")
print("  - loss_curves.png")
print("  - gradient_flow.png")
print("  - hidden_activations.png")
print("  - report.md")
print("  - gradient_magnitudes.json")
print("  - loss_histories.json")
print("  - final_losses.json")

print(f"\n[{datetime.now().strftime('%H:%M:%S')}] All done!")