"""
CCT-ODE Escape Variable Machine Learning
Proof of Concept: MLP on CIFAR-10

Tests 4 methods:
1. Standard SGD (baseline)
2. Escape Type 1: Latent Dimension Injection
3. Escape Type 2: Loss Landscape Deformation (ζ-deformation)
4. Escape Type 3: Complex-Valued Parameter Escape

Author: CCT Framework Implementation
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
import os

# ============================================================================
# CONFIGURATION
# ============================================================================
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")

BATCH_SIZE = 128
EPOCHS = 30
LR = 0.01
MOMENTUM = 0.9

# CCT Configuration
ENTROPY_THRESHOLD = 1e-4
TRAP_STEPS_THRESHOLD = 5  # How many flat steps before escape
ESCAPE_MAGNITUDE = 0.5    # How strong the escape is

# ============================================================================
# DATASET LOADING
# ============================================================================
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = datasets.CIFAR10(root='../data', train=True, transform=transform, download=True)
test_dataset = datasets.CIFAR10(root='../data', train=False, transform=transform, download=True)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

NUM_CLASSES = 10

# ============================================================================
# MODELS
# ============================================================================

class StandardMLP(nn.Module):
    """Baseline MLP: No escape variables."""
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3 * 32 * 32, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, NUM_CLASSES)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x


class LatentInjectionMLP(nn.Module):
    """Type 1: Latent dimension injection on trapped layers."""
    def __init__(self, latent_dim=32):
        super().__init__()
        self.latent_dim = latent_dim
        self.fc1 = nn.Linear(3 * 32 * 32, 512)
        self.fc2 = nn.Linear(512 + latent_dim, 256)
        self.fc3 = nn.Linear(256, NUM_CLASSES)
        self.latent_gate = nn.Parameter(torch.ones(1))
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x, inject_latent=False):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        
        batch_size = x.size(0)
        if inject_latent:
            # Inject escape latent dimension
            z = torch.randn(batch_size, self.latent_dim, device=x.device) * self.latent_gate
        else:
            # Inject zero latent dimension to maintain shape
            z = torch.zeros(batch_size, self.latent_dim, device=x.device)
            
        x = torch.cat([x, z], dim=1)
        
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x


class DeformationMLP(nn.Module):
    """Type 2: Loss landscape deformation via ζ-parameter."""
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3 * 32 * 32, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, NUM_CLASSES)
        self.dropout = nn.Dropout(0.2)
        
        # ζ (zeta) escape parameter - modifies gradient direction
        # Set requires_grad=False to manage it manually and avoid divergence
        self.zeta = nn.Parameter(torch.tensor(0.0), requires_grad=False)
        self.zeta_momentum = 0.0
    
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x
    
    def get_deformed_loss(self, loss, grad_norm):
        """L_deform = L + ζ * C(θ) where C is curvature."""
        # Clamp curvature term to prevent explosion
        curvature_term = torch.clamp(grad_norm ** 2, 0, 100.0)
        return loss + self.zeta * curvature_term


class ComplexMLP(nn.Module):
    """Type 3: Complex-valued parameter escape via imaginary perturbation."""
    def __init__(self):
        super().__init__()
        self.fc1_real = nn.Linear(3 * 32 * 32, 512)
        self.fc1_imag = nn.Linear(3 * 32 * 32, 512)
        self.fc2_real = nn.Linear(512, 256)
        self.fc2_imag = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, NUM_CLASSES)
        self.dropout = nn.Dropout(0.2)
        # Set requires_grad=False to manage it manually
        self.zeta_imag = nn.Parameter(torch.tensor(0.0), requires_grad=False)
    
    def forward(self, x, use_complex_escape=False):
        x = x.view(x.size(0), -1)
        
        # Real path
        h1_real = F.relu(self.fc1_real(x))
        
        if use_complex_escape and self.zeta_imag > 0:
            # Imaginary escape path - inject via imaginary weights
            h1_imag = F.relu(self.fc1_imag(x))
            # Complex combination: sqrt(real² + imag²) with ζ weighting
            h1 = torch.sqrt(h1_real**2 + (self.zeta_imag * h1_imag)**2 + 1e-8)
        else:
            h1 = h1_real
        
        h1 = self.dropout(h1)
        h2 = F.relu(self.fc2_real(h1))
        h2 = self.dropout(h2)
        out = self.fc3(h2)
        return out

# ============================================================================
# CCT ENTROPY MONITOR
# ============================================================================

class CCTEntropyMonitor:
    """Tracks entropy (uncertainty) of gradient landscape to detect traps."""
    
    def __init__(self, window_size=10):
        self.window_size = window_size
        self.grad_history = []
        self.entropy_history = []
    
    def update(self, grad_norm):
        if not np.isfinite(grad_norm):
            return self.entropy_history[-1] if self.entropy_history else float('inf')
            
        self.grad_history.append(grad_norm)
        if len(self.grad_history) > self.window_size:
            self.grad_history.pop(0)
        
        # Compute entropy proxy: variance in gradient magnitudes
        if len(self.grad_history) >= 2:
            # Avoid warnings if history is constant or has inf
            with np.errstate(all='ignore'):
                variance = np.var(self.grad_history)
                entropy = np.log(variance + 1e-8)
        else:
            entropy = float('inf')
        
        self.entropy_history.append(entropy)
        return entropy
    
    def is_trapped(self, threshold=ENTROPY_THRESHOLD):
        """Low entropy variance = trapped state."""
        if len(self.grad_history) < self.window_size:
            return False
        
        # Trapped if gradients are consistently small (low variance = stuck)
        mean_grad = np.mean(self.grad_history)
        return mean_grad < threshold
    
    def get_trap_count(self):
        """Count consecutive near-zero gradient steps."""
        count = 0
        for g in reversed(self.grad_history):
            if g < ENTROPY_THRESHOLD:
                count += 1
            else:
                break
        return count

# ============================================================================
# TRAINING FUNCTIONS
# ============================================================================

def train_epoch(model, loader, optimizer, device, method='standard', entropy_monitor=None):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)
        
        # Check for trap using CCT entropy monitor
        inject_latent = False
        use_complex_escape = False
        apply_deformation = False
        
        if entropy_monitor and method != 'standard':
            trap_count = entropy_monitor.get_trap_count()
            if trap_count >= TRAP_STEPS_THRESHOLD:
                if method == 'latent':
                    inject_latent = True
                    print(f"    [CCT ESCAPE] Trap detected! Injecting latent dimension.")
                elif method == 'deformation':
                    apply_deformation = True
                    print(f"    [CCT ESCAPE] Trap detected! Deforming loss landscape.")
                elif method == 'complex':
                    use_complex_escape = True
                    print(f"    [CCT ESCAPE] Trap detected! Complex escape activated.")

        if batch_idx % 100 == 0:
            print(f"    Batch {batch_idx:3d}/{len(loader)}")

        for i in range(10):
            # For deformation method: we need the gradient norm from the previous step
            # before zeroing it out.
            grad_norm_prev = torch.tensor(1.0, device=device)
            if method == 'deformation':
                grads = [p.grad.reshape(-1) for p in model.parameters() if p.grad is not None]
                if grads:
                    grad_norm_prev = torch.norm(torch.cat(grads))

            # Forward pass
            optimizer.zero_grad()
            
            if method == 'latent':
                output = model(data, inject_latent=inject_latent)
            elif method == 'complex':
                output = model(data, use_complex_escape=use_complex_escape)
            else:
                output = model(data)
            
            loss = F.cross_entropy(output, target)
            
            # Deformation method: modify loss with ζ
            if method == 'deformation' and hasattr(model, 'get_deformed_loss'):
                loss = model.get_deformed_loss(loss, grad_norm_prev)
            
            # Backward pass
            loss.backward()
            
            # Update escape parameters manually (not by optimizer)
            with torch.no_grad():
                trap_count = entropy_monitor.get_trap_count() if entropy_monitor else 0
                is_stuck = trap_count >= TRAP_STEPS_THRESHOLD
                
                if method == 'deformation' and hasattr(model, 'zeta'):
                    if is_stuck:
                        # Increase ζ to deform landscape
                        model.zeta.data = torch.clamp(model.zeta.data + 0.1 * ESCAPE_MAGNITUDE, -2.0, 2.0)
                    else:
                        # Decay ζ toward zero when escaping
                        model.zeta.data *= 0.95
                
                elif method == 'complex' and hasattr(model, 'zeta_imag'):
                    if is_stuck:
                        # Increase imaginary component injection
                        model.zeta_imag.data = torch.clamp(model.zeta_imag.data + 0.1 * ESCAPE_MAGNITUDE, 0.0, 1.0)
                    else:
                        # Decay toward real-only
                        model.zeta_imag.data *= 0.95
            
            optimizer.step()
            if i == 0:        
                if batch_idx % 100 == 0:
                    if method == 'deformation':
                        print(f"      [Step 0] Loss: {loss.item():.4f} | GradNorm: {grad_norm_prev.item():.4f} | Zeta: {model.zeta.item():.4f}")
                    elif method == 'complex':
                        print(f"      [Step 0] Loss: {loss.item():.4f} | ZetaImag: {model.zeta_imag.item():.4f}")
                
                total_loss += loss.item()
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()
                total += target.size(0)
            
        # Update entropy monitor
        if entropy_monitor:
            grads = [p.grad.reshape(-1) for p in model.parameters() if p.grad is not None]
            if grads:
                grad_norm = torch.norm(torch.cat(grads)).item()
                entropy_monitor.update(grad_norm)
    
    return total_loss / len(loader), correct / total


def evaluate(model, loader, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)
            
            total_loss += loss.item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
    
    return total_loss / len(loader), correct / total

# ============================================================================
# MAIN EXPERIMENT
# ============================================================================

def run_experiment():
    results = {}
    
    methods = [
        #('Standard SGD', 'standard', StandardMLP().to(DEVICE)),
        #('Type 1: Latent Injection', 'latent', LatentInjectionMLP().to(DEVICE)),
        ('Type 2: ζ-Deformation', 'deformation', DeformationMLP().to(DEVICE)),
        ('Type 3: Complex Escape', 'complex', ComplexMLP().to(DEVICE)),
    ]
    
    for name, method, model in methods:
        print(f"\n{'='*60}")
        print(f"Training: {name}")
        print(f"{'='*60}")
        
        optimizer = optim.SGD(model.parameters(), lr=LR, momentum=MOMENTUM)
        entropy_monitor = CCTEntropyMonitor()
        
        train_losses = []
        train_accs = []
        test_losses = []
        test_accs = []
        escape_events = []
        
        for epoch in range(EPOCHS):
            train_loss, train_acc = train_epoch(
                model, train_loader, optimizer, DEVICE, 
                method=method, entropy_monitor=entropy_monitor
            )
            test_loss, test_acc = evaluate(model, test_loader, DEVICE)
            
            train_losses.append(train_loss)
            train_accs.append(train_acc)
            test_losses.append(test_loss)
            test_accs.append(test_acc)
            
            # Track escape events for deformation method
            if method == 'deformation' and hasattr(model, 'zeta'):
                escape_events.append(model.zeta.item())
            
            print(f"Epoch {epoch+1:2d}/{EPOCHS} | "
                  f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | "
                  f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}")
        
        results[name] = {
            'method': method,
            'train_losses': train_losses,
            'train_accs': train_accs,
            'test_losses': test_losses,
            'test_accs': test_accs,
            'final_test_acc': test_accs[-1],
            'escape_events': escape_events if method == 'deformation' else None,
        }
        
        print(f"\nFinal Test Accuracy: {test_accs[-1]:.4f}")
    
    return results

# ============================================================================
# VISUALIZATION
# ============================================================================

def plot_results(results):
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    colors = ['blue', 'green', 'red', 'purple']
    
    # Plot 1: Training Loss
    ax1 = axes[0, 0]
    for i, (name, data) in enumerate(results.items()):
        ax1.plot(data['train_losses'], label=name, color=colors[i], alpha=0.8)
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Training Loss')
    ax1.set_title('Training Loss Comparison')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Training Accuracy
    ax2 = axes[0, 1]
    for i, (name, data) in enumerate(results.items()):
        ax2.plot(data['train_accs'], label=name, color=colors[i], alpha=0.8)
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Training Accuracy')
    ax2.set_title('Training Accuracy Comparison')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # Plot 3: Test Accuracy
    ax3 = axes[1, 0]
    for i, (name, data) in enumerate(results.items()):
        ax3.plot(data['test_accs'], label=name, color=colors[i], alpha=0.8)
    ax3.set_xlabel('Epoch')
    ax3.set_ylabel('Test Accuracy')
    ax3.set_title('Test Accuracy Comparison')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    
    # Plot 4: Final Comparison Bar
    ax4 = axes[1, 1]
    names = list(results.keys())
    final_accs = [results[n]['final_test_acc'] for n in names]
    bars = ax4.bar(names, final_accs, color=colors, alpha=0.8)
    ax4.set_ylabel('Final Test Accuracy')
    ax4.set_title('Final Test Accuracy Comparison')
    ax4.set_ylim(0, 1)
    for bar, acc in zip(bars, final_accs):
        ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01, 
                 f'{acc:.3f}', ha='center', va='bottom', fontsize=10)
    ax4.tick_params(axis='x', rotation=15)
    ax4.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig('cct_mlp_cifar10_results.png', dpi=150, bbox_inches='tight')
    print("\nPlot saved to: cct_mlp_cifar10_results.png")
    plt.show()

# ============================================================================
# ENTRY POINT
# ============================================================================

if __name__ == '__main__':
    print("="*60)
    print("CCT-ODE Escape Variable ML Training")
    print("Proof of Concept: MLP on CIFAR-10")
    print("="*60)
    
    # Run experiments
    results = run_experiment()
    
    # Plot results
    plot_results(results)
    
    # Print summary
    print("\n" + "="*60)
    print("FINAL RESULTS SUMMARY")
    print("="*60)
    
    for name, data in results.items():
        print(f"{name:30s}: Test Accuracy = {data['final_test_acc']:.4f}")
    
    best_method = max(results.items(), key=lambda x: x[1]['final_test_acc'])
    print(f"\nBest Method: {best_method[0]} with {best_method[1]['final_test_acc']:.4f} accuracy")