#!/usr/bin/env python3
"""
Gravity & Lorentz Mechanics for Machine Learning
==================================================
Physics-inspired neural network using:
- Gravitational attention (1/r² attraction)
- Lorentz force updates (v × B cross-product)
- GM = v²r manifold constraint (singularity prevention)
"""

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
import numpy as np
from typing import Tuple, Optional

# ============================================================================
# PHYSICS CONSTANTS & HELPERS
# ============================================================================

class PhysicsConfig:
    """Configuration for physics-inspired components."""
    G = 1.0              # Gravitational constant (learnable)
    q = 0.01             # Charge / coupling strength for Lorentz force
    eps = 1e-6           # Singularity prevention
    manifold_radius = 1.0  # Target manifold radius for GM=v²r constraint
    use_riemannian = True  # Whether to enforce manifold constraint


def safe_normalize(x: torch.Tensor, dim: int = -1, eps: float = 1e-8) -> torch.Tensor:
    """Safe normalization with numerical stability."""
    norm = torch.norm(x, p=2, dim=dim, keepdim=True)
    return x / (norm + eps)


def project_to_manifold(x: torch.Tensor, radius: float = 1.0) -> torch.Tensor:
    """
    Project embeddings onto a Riemannian manifold.
    GM = v²r → keep ||embedding||² × ||gradient|| = constant
    """
    norm = torch.norm(x, p=2, dim=-1, keepdim=True)
    # Project to sphere of given radius
    return (x / (norm + PhysicsConfig.eps)) * radius


# ============================================================================
# GRAVITATIONAL ATTENTION MECHANISM
# ============================================================================

class GravitationalAttention(nn.Module):
    """
    Attention mechanism based on gravitational attraction: F = GMm/r²
    
    Instead of softmax similarity, we use 1/r² attraction between query-key pairs.
    This creates a physically-motivated attention with inherent inductive biases.
    """
    
    def __init__(self, embed_dim: int, num_heads: int = 8, G: float = 1.0):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
        
        # Gravitational constant (learnable for flexibility)
        self.G = nn.Parameter(torch.tensor(G), requires_grad=True)
        
        # Projections
        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)
        
        # Learnable bias term for stability
        self.bias = nn.Parameter(torch.zeros(1))
        
    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x: [batch, seq_len, embed_dim] - query/source
            context: [batch, ctx_len, embed_dim] - key/value context (if None, use x)
        """
        if context is None:
            context = x
            
        batch, seq_len, _ = x.shape
        ctx_len = context.shape[1]
        
        # Project to Q, K, V
        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)
        
        # Compute Euclidean distances between Q and K
        # distances: [batch, num_heads, seq_len, ctx_len]
        Q_expanded = Q.unsqueeze(3)  # [B, H, S, 1, D]
        K_expanded = K.unsqueeze(2)  # [B, H, 1, C, D]
        distances = torch.norm(Q_expanded - K_expanded, p=2, dim=-1)  # r
        
        # Gravitational attraction: F = GM/r²
        # Using G from config and adding bias for stability
        gravity_scores = self.G / (distances ** 2 + PhysicsConfig.eps + self.bias)
        
        # Normalize to weights (like softmax)
        attention_weights = gravity_scores / (gravity_scores.sum(dim=-1, keepdim=True) + PhysicsConfig.eps)
        
        # Apply attention to values
        # V: [B, H, C, D]
        attended = torch.matmul(attention_weights, V)
        
        # Reshape output
        attended = attended.transpose(1, 2).reshape(batch, seq_len, self.embed_dim)
        
        return self.out_proj(attended)


# ============================================================================
# LORENTZ FORCE GRADIENT UPDATE
# ============================================================================

class LorentzGradientModifier(nn.Module):
    """
    Implements Lorentz force: F = q(E + v × B)
    
    The cross-product v × B creates updates perpendicular to both
    velocity (gradient direction) and B-field (hidden state).
    This geometric structure stabilizes training dynamics.
    """
    
    def __init__(self, embed_dim: int, q: float = 0.01):
        super().__init__()
        self.embed_dim = embed_dim
        self.q = nn.Parameter(torch.tensor(q), requires_grad=True)
        
        # B-field generator from rotation (like sun's differential rotation)
        self.b_field_generator = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.Tanh(),
            nn.Linear(embed_dim, embed_dim)
        )
        
        # Rotation matrix generator for differential rotation
        self.rotation_speed = nn.Parameter(torch.ones(1) * 0.1)
        
    def create_b_field(self, hidden_state: torch.Tensor) -> torch.Tensor:
        """
        Generate B-field from hidden state.
        In physics, changing magnetic fields induce electric fields.
        Here, the B-field is derived from hidden representations.
        """
        return self.b_field_generator(hidden_state)
    
    def apply_lorentz_update(
        self, 
        gradients: torch.Tensor, 
        hidden_state: torch.Tensor,
        velocity: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Apply Lorentz force-inspired update.
        
        Args:
            gradients: Standard gradients from backprop (E-field)
            hidden_state: Current hidden representations
            velocity: Movement direction (if None, use gradient direction)
        """
        # v = velocity direction (gradient normalized)
        if velocity is None:
            velocity = F.normalize(gradients.flatten(1), dim=-1).reshape_as(gradients)
        
        # B-field from hidden state
        B = self.create_b_field(hidden_state)
        
        # v × B: cross product for perpendicular update
        # For high-dimensional case, we compute this differently
        # Using batched cross product along feature dimension
        v = velocity.reshape(-1, 3, self.embed_dim // 3)
        b = B.reshape(-1, 3, self.embed_dim // 3)
        cross_term = torch.cross(v, b, dim=1).reshape_as(gradients)
        
        # Lorentz force: F = q * (E + v × B)
        lorentz_force = gradients + self.q * cross_term
        
        return lorentz_force


# ============================================================================
# DIFFERENTIAL ROTATION LAYER (Sun-like B-field generation)
# ============================================================================

class DifferentialRotation(nn.Module):
    """
    Inspired by Sun's differential rotation generating magnetic fields.
    Different attention heads have different "rotation speeds".
    """
    
    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 rotation speeds for each head
        self.head_speeds = nn.Parameter(torch.randn(num_heads) * 0.1)
        
        # Rotation projection
        self.rotation_proj = nn.Linear(embed_dim, embed_dim)
        
    def create_rotation_matrix(self, angle: torch.Tensor, dim: int = 0) -> torch.Tensor:
        """Create rotation matrix for given angle."""
        cos_a = torch.cos(angle)
        sin_a = torch.sin(angle)
        # Simple 2D rotation in the first two dimensions
        zero = torch.zeros_like(angle)
        one = torch.ones_like(angle)
        rotation = torch.stack([
            torch.stack([cos_a, -sin_a, zero]),
            torch.stack([sin_a, cos_a, zero]),
            torch.stack([zero, zero, one])
        ])
        return rotation
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply differential rotation to generate B-field.
        Different heads rotate at different speeds → richer representations.
        """
        batch, seq_len, embed_dim = x.shape
        
        # Split into heads
        x_heads = x.reshape(batch, seq_len, self.num_heads, self.head_dim)
        
        # Apply different rotations to each head
        b_field = torch.zeros_like(x)
        
        for h in range(self.num_heads):
            speed = self.head_speeds[h]
            rotation = self.create_rotation_matrix(speed)
            
            # Apply rotation to head features
            head_features = x_heads[:, :, h, :3]  # Take first 3 dims for rotation
            rotated = torch.einsum('ij,bkj->bki', rotation, head_features)
            
            # Place back
            b_field[:, :, h * self.head_dim:(h + 1) * self.head_dim][:, :, :3] += rotated
            
        return self.rotation_proj(b_field)


# ============================================================================
# GRAVITATIONAL FEATURE EXTRACTION
# ============================================================================

class GravitationalConvBlock(nn.Module):
    """Convolutional block with gravitational attraction features."""
    
    def __init__(self, in_channels: int, out_channels: int, use_gravity: bool = True):
        super().__init__()
        self.use_gravity = use_gravity
        
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)
        
        if use_gravity:
            # Learnable gravitational constant for spatial attention
            self.G_spatial = nn.Parameter(torch.tensor(1.0))
            self.scale = nn.Parameter(torch.ones(1))
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.conv(x)
        out = self.bn(out)
        
        if self.use_gravity:
            # Create gravitational attention map
            # Attraction between spatial positions based on feature similarity
            B, C, H, W = out.shape
            
            # Reshape to [B, C, H*W]
            features = out.reshape(B, C, -1)
            
            # Compute pairwise distances (simplified gravitational attraction)
            # Using cosine distance for efficiency
            features_norm = F.normalize(features, dim=1)
            similarity = torch.matmul(features_norm.transpose(1, 2), features_norm)
            
            # Gravitational: 1/(distance)² → but we have similarity, so invert
            # Higher similarity = closer = stronger attraction
            gravity_weights = self.G_spatial / (1 - similarity + PhysicsConfig.eps)
            gravity_weights = F.softmax(gravity_weights, dim=-1)
            
            # Apply gravitational attention
            out = torch.matmul(out.reshape(B, C, -1), gravity_weights).reshape(B, C, H, W)
            out = out * self.scale
            
        out = F.gelu(out)
        return out


# ============================================================================
# MAIN NETWORK: GRAVITON-CIFAR
# ============================================================================

class GravitonCIFAR(nn.Module):
    """
    CIFAR-10 classifier using gravitational and Lorentz mechanics.
    
    Key components:
    - Gravitational attention for feature aggregation
    - Lorentz force updates for gradient modification
    - Differential rotation for B-field generation
    - GM=v²r manifold projection for stable representations
    """
    
    def __init__(
        self, 
        num_classes: int = 10,
        embed_dim: int = 256,
        num_heads: int = 8,
        use_physics: bool = True
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.use_physics = use_physics
        
        # Initial convolutions to extract features
        self.conv_embed = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            GravitationalConvBlock(64, 128, use_gravity=use_physics),
            nn.MaxPool2d(2),
            GravitationalConvBlock(128, 256, use_gravity=use_physics),
            nn.MaxPool2d(2),
            GravitationalConvBlock(256, 512, use_gravity=use_physics),
            nn.MaxPool2d(2),
        )
        
        # Convert spatial features to sequence for attention
        self.spatial_to_seq = nn.AdaptiveAvgPool2d((4, 4))
        
        # Gravitational attention layer
        self.grav_attention = GravitationalAttention(embed_dim, num_heads) if use_physics else None
        
        # Differential rotation for B-field
        self.differential_rot = DifferentialRotation(embed_dim, num_heads) if use_physics else None
        
        # Lorentz gradient modifier
        self.lorentz_modifier = LorentzGradientModifier(embed_dim) if use_physics else None
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim * 16, 512),  # 4*4*16 = 256 -> but we have embed_dim
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes)
        )
        
        # GM=v²r manifold projection
        self.register_buffer('manifold_radius', torch.tensor(1.0))
        
    def forward(self, x: torch.Tensor, return_attention: bool = False) -> torch.Tensor:
        # Feature extraction
        features = self.conv_embed(x)  # [B, 512, 4, 4]
        
        # Convert to sequence: [B, 16, 512]
        batch, channels, h, w = features.shape
        seq_features = features.reshape(batch, channels, -1).transpose(1, 2)
        
        # Project to embedding dimension
        seq_features = F.linear(seq_features, 
                               torch.eye(self.embed_dim, channels).to(x.device)[:channels, :])
        
        # Apply gravitational attention (physics-inspired)
        if self.use_physics and self.grav_attention is not None:
            # Generate B-field via differential rotation
            b_field = self.differential_rot(seq_features)
            
            # Apply gravitational attention
            attended = self.grav_attention(seq_features)
            
            # Combine with B-field influence
            seq_features = attended + 0.1 * b_field
        
        # GM=v²r constraint: project to manifold
        if self.use_physics:
            seq_features = project_to_manifold(seq_features, self.manifold_radius.item())
        
        # Flatten and classify
        flat = seq_features.reshape(batch, -1)
        
        if return_attention and self.use_physics:
            # Compute attention weights for visualization
            with torch.no_grad():
                distances = torch.cdist(seq_features, seq_features)
                attention = self.grav_attention.G.item() / (distances ** 2 + 1e-6)
                attention = F.softmax(attention, dim=-1)
            return self.classifier(flat), attention
        
        return self.classifier(flat)
    
    def lorentz_gradient_step(self, model_state: dict, gradients: dict, 
                               hidden_states: dict) -> dict:
        """
        Apply Lorentz force to gradient updates.
        This modifies gradients using: F = q(E + v × B)
        """
        if not self.use_physics:
            return gradients
            
        modified_grads = {}
        for name, grad in gradients.items():
            if grad is not None and name in hidden_states:
                h_state = hidden_states[name]
                modified_grads[name] = self.lorentz_modifier.apply_lorentz_update(
                    grad, h_state
                )
            else:
                modified_grads[name] = grad
                
        return modified_grads


# ============================================================================
# PHYSICS-AWARE OPTIMIZER
# ============================================================================

class LorentzOptimizer(optim.Optimizer):
    """
    Custom optimizer applying Lorentz force modifications to gradients.
    
    The update rule: θ = θ - α * (∇E + q * v × B)
    
    Where:
    - ∇E is the standard gradient (E-field)
    - v is the gradient direction (velocity)
    - B is generated from current state
    - q controls cross-product strength
    """
    
    def __init__(self, params, lr=0.001, q=0.01, momentum=0.9):
        defaults = dict(lr=lr, q=q, momentum=momentum)
        super().__init__(params, defaults)
        self.q = q
        
    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
            
        for group in self.param_groups:
            lr = group['lr']
            q = group['q']
            momentum = group['momentum']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                    
                grad = p.grad.data
                state = self.state[p]
                
                # Initialize momentum state
                if len(state) == 0:
                    state['velocity'] = torch.zeros_like(p.data)
                    state['b_field'] = torch.zeros_like(p.data)
                    
                # Update velocity (momentum)
                state['velocity'] = momentum * state['velocity'] + grad
                
                # Generate B-field from current parameters
                state['b_field'] = torch.tanh(p.data)
                
                # Lorentz force: E + v × B
                # Simplified cross product in parameter space
                cross_term = torch.cross(
                    state['velocity'].flatten(),
                    state['b_field'].flatten()
                ).reshape_as(grad)
                
                lorentz_grad = grad + q * cross_term
                
                # Update
                p.data = p.data - lr * lorentz_grad
                
        return loss


# ============================================================================
# TRAINING LOOP
# ============================================================================

def train_graviton(
    model: nn.Module,
    train_loader: DataLoader,
    test_loader: DataLoader,
    device: torch.device,
    epochs: int = 50,
    use_lorentz_optimizer: bool = False
) -> Tuple[list, list, list, list]:
    """Train the Graviton model on CIFAR-10."""
    
    criterion = nn.CrossEntropyLoss()
    
    if use_lorentz_optimizer:
        optimizer = LorentzOptimizer(model.parameters(), lr=0.001, q=0.01)
    else:
        optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)
    
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    train_losses = []
    train_accs = []
    test_losses = []
    test_accs = []
    
    for epoch in range(epochs):
        # Training
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
        train_loss = running_loss / len(train_loader)
        train_acc = 100. * correct / total
        
        # Testing
        model.eval()
        test_loss = 0.0
        test_correct = 0
        test_total = 0
        
        with torch.no_grad():
            for inputs, targets in test_loader:
                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(test_loader)
        test_acc = 100. * test_correct / test_total
        
        scheduler.step()
        
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        test_losses.append(test_loss)
        test_accs.append(test_acc)
        
        print(f"Epoch {epoch+1:3d}/{epochs} | "
              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | "
              f"Test Loss: {test_loss:.4f} Acc: {test_acc:.2f}%")
        
    return train_losses, train_accs, test_losses, test_accs


# ============================================================================
# MAIN
# ============================================================================

def main():
    # Setup
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Data loading
    transform = 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))
    ])
    
    trainset = torchvision.datasets.CIFAR10(
        root='../data', train=True, download=True, transform=transform
    )
    testset = torchvision.datasets.CIFAR10(
        root='../data', train=False, download=True, transform=transform
    )
    
    trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)
    testloader = DataLoader(testset, batch_size=128, shuffle=False, num_workers=4)
    
    print(f"Training samples: {len(trainset)}, Test samples: {len(testset)}")
    
    # Create models for comparison
    print("\n" + "="*60)
    print("Training with Physics (Gravitational Attention + Lorentz Force)")
    print("="*60)
    
    model_physics = GravitonCIFAR(
        num_classes=10,
        embed_dim=256,
        num_heads=8,
        use_physics=True
    ).to(device)
    
    losses, accs, test_losses, test_accs = train_graviton(
        model_physics, trainloader, testloader, device, epochs=30, use_lorentz_optimizer=False
    )
    
    # Save results
    torch.save({
        'model_state': model_physics.state_dict(),
        'train_losses': losses,
        'train_accs': accs,
        'test_losses': test_losses,
        'test_accs': test_accs,
        'G': model_physics.grav_attention.G.item() if model_physics.grav_attention else None
    }, 'graviton_cifar_results.pth')
    
    print("\n" + "="*60)
    print("Physics parameters learned:")
    print(f"  Gravitational constant G: {model_physics.grav_attention.G.item():.4f}")
    print(f"  Lorentz charge q: {model_physics.lorentz_modifier.q.item():.4f}")
    print("="*60)
    
    # Print best results
    best_acc = max(test_accs)
    best_epoch = test_accs.index(best_acc) + 1
    print(f"\nBest test accuracy: {best_acc:.2f}% at epoch {best_epoch}")
    
    return model_physics, losses, test_accs


if __name__ == '__main__':
    model, train_loss, test_acc = main()