#!/usr/bin/env python3
"""
Graviton-CIFAR: Stable Physics-Inspired Neural Network
=======================================================
Fixed issues:
1. Cross-product destabilizing gradients → reduced q strength
2. Gravitational attention saturating → proper normalization
3. Manifold projection causing collapse → gentler projection
4. Missing gradient clipping
"""

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

# ============================================================================
# STABILIZATION FIXES
# ============================================================================

def stable_normalize(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Normalization with clamp to prevent explosion."""
    norm = torch.norm(x, p=2, dim=dim, keepdim=True)
    return x / (norm.clamp(min=1e-6, max=10.0))

class SoftmaxAttention(nn.Module):
    """Stable baseline attention for comparison."""
    def __init__(self, embed_dim: int, num_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor:
        if context is None:
            context = x
            
        batch, seq_len, _ = x.shape
        ctx_len = context.shape[1]
        
        Q = self.q_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k_proj(context).reshape(batch, ctx_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v_proj(context).reshape(batch, ctx_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Stable softmax attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        
        out = torch.matmul(attn, V).transpose(1, 2).reshape(batch, seq_len, self.embed_dim)
        return self.out_proj(out)


class GravitationalAttention(nn.Module):
    """
    FIXED: Gravitational attention with proper normalization.
    Uses 1/r² but properly scaled to prevent vanishing/exploding.
    """
    def __init__(self, embed_dim: int, num_heads: int = 8):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        # Learnable scale for gravitational strength (was G constant)
        self.gravity_scale = nn.Parameter(torch.tensor(1.0))
        
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor:
        if context is None:
            context = x
            
        batch, seq_len, _ = x.shape
        ctx_len = context.shape[1]
        
        Q = self.q_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)
        K = self.k_proj(context).reshape(batch, ctx_len, self.num_heads, self.head_dim)
        V = self.v_proj(context).reshape(batch, ctx_len, self.num_heads, self.head_dim)
        
        # Compute distances with clamping for stability
        # dist = ||q - k||² for gravitational attraction
        Q_expanded = Q.unsqueeze(2)  # [B, S, 1, H, D]
        K_expanded = K.unsqueeze(1)  # [B, 1, C, H, D]
        
        # Squared distance
        dist_sq = torch.sum((Q_expanded - K_expanded) ** 2, dim=-1)  # [B, S, C, H]
        
        # FIX: Clamp distance to prevent explosion when r→0
        dist_sq = dist_sq.clamp(min=0.01, max=100.0)
        
        # Gravitational attraction: scale / r²
        # Using sqrt for gentler 1/r instead of 1/r²
        gravity_scores = self.gravity_scale / (dist_sq + 1.0)  # +1.0 for additional stability
        
        # Normalize per query position (like softmax)
        gravity_scores = gravity_scores / (gravity_scores.sum(dim=2, keepdim=True) + 1e-8)
        
        # Apply attention
        attn = self.dropout(gravity_scores)
        out = torch.einsum('bsch,bchd->bshd', attn, V)
        out = out.reshape(batch, seq_len, self.embed_dim)
        
        return self.out_proj(out)


class LorentzGradientLayer(nn.Module):
    """
    FIXED: Lorentz-inspired gradient modifier with reduced impact.
    F = q * (v × B) but scaled down to prevent destabilization.
    """
    def __init__(self, embed_dim: int):
        super().__init__()
        self.embed_dim = embed_dim
        
        # FIX: Much smaller initial q value
        self.q = nn.Parameter(torch.tensor(0.001))  # Was 0.01, now 0.001
        
        # Gentle B-field projection
        self.b_proj = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.Tanh(),
            nn.Linear(embed_dim // 2, embed_dim)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply gentle Lorentz-like transformation.
        Instead of cross-product, use rotation-like operation.
        """
        B = self.b_proj(x)
        
        # Simple rotation: x + q * (B - x) = (1-q)*x + q*B
        # This is like a gentle interpolation, stable and interpretable
        out = x + self.q * (B - x)
        
        return out


class ManifoldProjection(nn.Module):
    """
    FIXED: GM=v²r constraint with gentle normalization.
    Projects to sphere but with gradual transitions.
    """
    def __init__(self, radius: float = 1.0, projection_ratio: float = 0.1):
        super().__init__()
        self.radius = radius
        self.projection_ratio = projection_ratio  # How strongly to apply
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = torch.norm(x, p=2, dim=-1, keepdim=True)
        
        # Gentle projection: only partially project
        projected = (x / (norm + 1e-6)) * self.radius
        
        # Interpolate: x_new = (1-ratio)*x + ratio*projected
        # This prevents sudden jumps
        return x + self.projection_ratio * (projected - x)


# ============================================================================
# STABILIZED GRAVITON MODEL
# ============================================================================

class GravitonCIFAR(nn.Module):
    """
    Stabilized CIFAR classifier using physics concepts.
    - Gravitational attention for feature aggregation
    - Lorentz-inspired residual connections  
    - Gentle manifold projection (no collapse)
    """
    
    def __init__(self, num_classes: int = 10, embed_dim: int = 128, 
                 num_heads: int = 4, use_physics: bool = True):
        super().__init__()
        self.use_physics = use_physics
        self.embed_dim = embed_dim
        
        # Convolutional feature extraction
        self.features = nn.Sequential(
            # Block 1: 32x32 -> 16x16
            nn.Conv2d(3, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.1),
            
            # Block 2: 16x16 -> 8x8
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.GELU(),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.1),
            
            # Block 3: 8x8 -> 4x4
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.GELU(),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.1),
            
            # Block 4: 4x4 -> 2x2
            nn.Conv2d(256, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.GELU(),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.1),
        )
        
        # Convert to sequence for attention
        self.to_embed = nn.Linear(512, embed_dim)
        
        # Attention mechanism
        if use_physics:
            self.attention = GravitationalAttention(embed_dim, num_heads)
        else:
            self.attention = SoftmaxAttention(embed_dim, num_heads)
        
        # Lorentz-inspired layer
        self.lorentz = LorentzGradientLayer(embed_dim) if use_physics else nn.Identity()
        
        # Manifold projection (gentle)
        self.manifold = ManifoldProjection(radius=1.0, projection_ratio=0.05) if use_physics else nn.Identity()
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Extract features: [B, 512, 2, 2]
        feat = self.features(x)
        
        # Convert to sequence: [B, 4, 512]
        batch = feat.shape[0]
        seq = feat.reshape(batch, 512, -1).transpose(1, 2)  # [B, 4, 512]
        
        # Project to embed dim: [B, 4, embed_dim]
        seq = self.to_embed(seq)
        
        # Apply attention
        seq = self.attention(seq) + seq  # Residual
        
        # Apply Lorentz transformation
        seq = self.lorentz(seq) + seq  # Residual
        
        # Gentle manifold projection
        seq = self.manifold(seq)
        
        # Global average pooling + classify
        pooled = seq.mean(dim=1)  # [B, embed_dim]
        
        return self.classifier(pooled)


# ============================================================================
# TRAINING WITH STABILITY
# ============================================================================

def train_model(model, trainloader, testloader, device, epochs=50, lr=0.001):
    """Training with gradient clipping for stability."""
    
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    results = {'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': []}
    
    for epoch in range(epochs):
        # Training
        model.train()
        running_loss = 0.0
        correct, total = 0, 0
        
        for inputs, targets in trainloader:
            inputs, targets = inputs.to(device), targets.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            
            # FIX: Gradient clipping to prevent explosion
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        scheduler.step()
        
        train_loss = running_loss / len(trainloader)
        train_acc = 100. * correct / total
        
        # Evaluation
        model.eval()
        test_loss, test_correct, test_total = 0.0, 0, 0
        
        with torch.no_grad():
            for inputs, targets in testloader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                test_total += targets.size(0)
                test_correct += predicted.eq(targets).sum().item()
        
        test_loss = test_loss / len(testloader)
        test_acc = 100. * test_correct / test_total
        
        results['train_loss'].append(train_loss)
        results['train_acc'].append(train_acc)
        results['test_loss'].append(test_loss)
        results['test_acc'].append(test_acc)
        
        print(f"Epoch {epoch+1:2d}/{epochs} | "
              f"Train: {train_loss:.4f}/{train_acc:.1f}% | "
              f"Test: {test_loss:.4f}/{test_acc:.1f}%")
    
    return results


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}\n")
    
    # Data
    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))
    ])
    
    trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform_train)
    testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test)
    
    trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)
    testloader = DataLoader(testset, batch_size=128, shuffle=False, num_workers=4)
    
    # Train physics model
    print("="*60)
    print("Training PHYSICS-MODEL (Gravitational + Lorentz)")
    print("="*60)
    
    model_physics = GravitonCIFAR(num_classes=10, embed_dim=128, num_heads=4, use_physics=True)
    results_p = train_model(model_physics, trainloader, testloader, device, epochs=30, lr=0.001)
    
    # Train baseline for comparison
    print("\n" + "="*60)
    print("Training BASELINE (Standard Attention)")
    print("="*60)
    
    model_baseline = GravitonCIFAR(num_classes=10, embed_dim=128, num_heads=4, use_physics=False)
    results_b = train_model(model_baseline, trainloader, testloader, device, epochs=30, lr=0.001)
    
    # Summary
    print("\n" + "="*60)
    print("RESULTS SUMMARY")
    print("="*60)
    best_physics = max(results_p['test_acc'])
    best_baseline = max(results_b['test_acc'])
    
    print(f"Physics Model - Best Test Acc: {best_physics:.2f}%")
    print(f"Baseline     - Best Test Acc: {best_baseline:.2f}%")
    
    if hasattr(model_physics, 'attention') and hasattr(model_physics.attention, 'gravity_scale'):
        print(f"Learned gravity scale: {model_physics.attention.gravity_scale.item():.4f}")
    
    # Save
    torch.save({
        'physics_results': results_p,
        'baseline_results': results_b
    }, 'graviton_comparison.pth')
    
    return results_p, results_b


if __name__ == '__main__':
    main()