"""
Singularity Memory ODE (SM-ODE) for CIFAR-10
Based on the theory: Memory stored in shared singularity, not lost between states.
"""

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

# ============================================================
# PART 1: THE SINGULARITY MEMORY MODULE
# ============================================================

class SingularityMemory(nn.Module):
    """
    Stores historical states in a 'singularity' that influences current predictions.
    
    Theory: History is not lost between states - it's stored in a singularity field.
    This allows 'variable determinism' where predictions are constrained by laws
    (standard weights) but influenced by accumulated memory (novel approach).
    """
    
    def __init__(self, feature_dim, memory_dim=256, decay=0.95):
        super().__init__()
        self.feature_dim = feature_dim
        self.memory_dim = memory_dim
        self.decay = decay  # Memory retention factor (closer to 1 = more memory)
        
        # The singularity memory tensor - accumulates all past states
        # Shape: [batch, memory_dim]
        self.register_buffer('singularity', torch.zeros(1, memory_dim))
        
        # Memory transformation: projects features into singularity space
        self.to_singularity = nn.Linear(feature_dim, memory_dim)
        
        # Memory influence gate: how much singularity affects output
        self.memory_gate = nn.Sequential(
            nn.Linear(memory_dim + feature_dim, memory_dim),
            nn.Sigmoid()
        )
        
        # Memory kernel (K): determines resonance/amplification from past
        self.kernel_net = nn.Sequential(
            nn.Linear(memory_dim, memory_dim),
            nn.Tanh()
        )
    
    def reset_singularity(self):
        """Reset singularity memory to zero"""
        self.singularity.fill_(0)

    def compute_memory_bias(self, current_features):        """
        Compute how past states influence current prediction.
        B(S, y) = integral of K(t-tau) * y(tau) dtau
        
        The singularity doesn't just store - it actively influences through kernel.
        """
        batch_size = current_features.size(0)
        
        # Kernel applies resonance weighting to stored memory
        kernel_output = self.kernel_net(self.singularity)
        
        # Memory bias: past states shaped by kernel
        memory_bias = self.singularity + 0.1 * kernel_output
        
        # Expand memory bias to match batch size
        memory_bias_expanded = memory_bias.expand(batch_size, -1)
        
        # Gate determines how much memory influences current features
        gate_input = torch.cat([memory_bias_expanded, current_features], dim=-1)
        gate = self.memory_gate(gate_input)
        
        # Return memory influence in singularity space
        return gate * torch.tanh(memory_bias_expanded)
    
    def update_singularity(self, new_features):
        """
        Update singularity memory with new states.
        S(t) = S(t-1) * decay + new_features
        
        For global singularity, we average updates over the batch.
        """
        new_proj = self.to_singularity(new_features)
        update_val = new_proj.mean(dim=0, keepdim=True)
        self.singularity = self.singularity * self.decay + update_val


# ============================================================
# PART 2: SHARED SINGULARITY NETWORK (Multiple heads share memory)
# ============================================================

class SharedSingularityBlock(nn.Module):
    """
    Multiple processing heads share a common singularity memory.
    
    Theory: Entangled particles share singularity memory (Phi_AB).
    Here, different feature channels share the same memory field.
    """
    
    def __init__(self, channels, memory_dim=128):
        super().__init__()
        self.channels = channels
        self.memory_dim = memory_dim
        
        # Shared singularity for this block
        self.singularity = SingularityMemory(channels, memory_dim)
        
        # Project memory influence back to channel dimension
        self.memory_to_channels = nn.Linear(memory_dim, channels)
        
        # Multiple processing heads (like multiple particles sharing memory)
        self.heads = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(channels, channels, 3, padding=1),
                nn.BatchNorm2d(channels),
                nn.ReLU(inplace=True)
            ) for _ in range(3)  # 3 heads sharing singularity
        ])
        
        # Fusion layer combines head outputs with memory influence
        self.fusion = nn.Sequential(
            nn.Linear(channels * 3 + memory_dim, channels * 2),
            nn.ReLU(),
            nn.Linear(channels * 2, channels)
        )
    
    def forward(self, x):
        batch_size, channels, h, w = x.shape
        
        # Compute head features
        head_outputs = [head(x) for head in self.heads]
        head_concat = torch.cat(head_outputs, dim=1)  # [B, C*3, H, W]
        
        # Pool for singularity interaction
        pooled = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1)  # [B, C]
        
        # Get memory-biased features [B, memory_dim]
        biased_features = self.singularity.compute_memory_bias(pooled)
        
        # Update singularity with current state
        self.singularity.update_singularity(pooled)
        
        # Combine heads with memory bias
        # Project memory back to channels
        memory_influence = self.memory_to_channels(biased_features) # [B, C]
        memory_expanded = memory_influence.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, h, w)
        
        # Fusion combines deterministic (heads) with memory-influenced (singularity)
        fused = self.fusion(torch.cat([
            head_concat.view(batch_size, -1, h * w).mean(-1),  # [B, C*3]
            biased_features  # [B, memory_dim]
        ], dim=-1))
        
        fused = fused.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, h, w)
        
        # Memory and heads both contribute
        return x + 0.2 * fused + 0.1 * torch.tanh(memory_expanded)


# ============================================================
# PART 3: THE FULL SM-ODE NETWORK
# ============================================================

class SMODECIFAR10(nn.Module):
    """
    Singularity Memory ODE Network for CIFAR-10
    
    Key innovations:
    1. Memory singularity stores history, doesn't lose it
    2. Multiple blocks share singularity memory (entanglement-like)
    3. Predictions = f(laws) + alpha * memory_bias
    4. Variable determinism: constrained by weights, influenced by memory
    """
    
    def __init__(self, num_classes=10, memory_dim=128):
        super().__init__()
        
        # Initial feature extraction (deterministic laws)
        self.input_conv = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        
        # Shared singularity blocks (memory influence grows with depth)
        self.shared_singularity1 = SharedSingularityBlock(64, memory_dim)
        self.shared_singularity2 = SharedSingularityBlock(128, memory_dim)
        self.shared_singularity3 = SharedSingularityBlock(256, memory_dim)
        
        # Transition layers
        self.transition1 = nn.Sequential(
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        
        self.transition2 = nn.Sequential(
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        
        # Global singularity: final aggregation of all memories
        self.global_singularity = SingularityMemory(256, memory_dim * 2)
        
        # Classifier: combines deterministic predictions with singularity influence
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256 + memory_dim * 2, 256),  # Features + global memory
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x, retain_memory=False):
        # Initial deterministic processing
        x = self.input_conv(x)
        
        # Block 1 with shared singularity
        x = self.shared_singularity1(x)
        
        # Transition
        x = self.transition1(x)
        
        # Block 2 with shared singularity
        x = self.shared_singularity2(x)
        
        # Transition
        x = self.transition2(x)
        
        # Block 3 with shared singularity
        x = self.shared_singularity3(x)
        
        # Global memory aggregation
        pooled = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1)
        memory_features = self.global_singularity.compute_memory_bias(pooled)
        self.global_singularity.update_singularity(pooled)
        
        # Final classification: laws + memory
        pooled = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1)
        combined = torch.cat([pooled, memory_features], dim=-1)
        
        return self.classifier(combined)


# ============================================================
# PART 4: TRAINING WITH MEMORY PERSISTENCE
# ============================================================

def train_with_singularity(model, trainloader, testloader, device, epochs=50):
    """
    Training that leverages singularity memory across batches.
    
    Unlike standard training where each batch is independent,
    here batches benefit from accumulated singularity memory.
    """
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    best_acc = 0
    memory_acc_history = []  # Track how memory improves accuracy
    
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        # Reset singularity only at start of each epoch (memory persists within epoch)
        # This models the theory: singularity stores all time steps
        batch_count = 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()
            
            # Gradient clipping for stability
            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()
            
            batch_count += 1
        
        scheduler.step()
        
        train_acc = 100. * correct / total
        test_acc = evaluate(model, testloader, device)
        
        memory_acc_history.append({
            'epoch': epoch + 1,
            'train_acc': train_acc,
            'test_acc': test_acc,
            'loss': running_loss / len(trainloader)
        })
        
        if test_acc > best_acc:
            best_acc = test_acc
            # Save best model
            torch.save(model.state_dict(), 'smode_cifar10_best.pth')
        
        print(f"Epoch {epoch+1}/{epochs} | Loss: {running_loss/len(trainloader):.4f} | "
              f"Train: {train_acc:.2f}% | Test: {test_acc:.2f}% | Best: {best_acc:.2f}%")
    
    return memory_acc_history, best_acc


def evaluate(model, testloader, device):
    """Evaluate model with memory (test batches still update singularity)"""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    return 100. * correct / total


# ============================================================
# PART 5: BASELINE COMPARISON (NO SINGULARITY)
# ============================================================

class BaselineCIFAR10(nn.Module):
    """Standard ResNet-like architecture without singularity memory"""
    
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(256, num_classes)
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)


# ============================================================
# MAIN TRAINING SCRIPT
# ============================================================

def main():
    # Setup
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # CIFAR-10 Data
    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=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ])
    )
    
    trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
    testloader = DataLoader(testset, batch_size=128, shuffle=False, num_workers=2)
    
    print("=" * 60)
    print("SINGULARITY MEMORY ODE (SM-ODE) CIFAR-10 MODEL")
    print("=" * 60)
    print("\nTheory: Memory stored in singularity, not lost between states")
    print("         -> Predictions use: f(laws) + alpha * memory_bias")
    print("         -> Multiple blocks share singularity (like entanglement)")
    print()
    
    # Train SM-ODE Model
    print("\n--- Training SM-ODE Model (with Singularity Memory) ---\n")
    smode_model = SMODECIFAR10(num_classes=10, memory_dim=128).to(device)
    
    total_params = sum(p.numel() for p in smode_model.parameters())
    print(f"SM-ODE Parameters: {total_params:,}")
    
    smode_history, smode_best = train_with_singularity(
        smode_model, trainloader, testloader, device, epochs=50
    )
    
    # Train Baseline for comparison
    print("\n--- Training Baseline Model (no Singularity) ---\n")
    baseline_model = BaselineCIFAR10(num_classes=10).to(device)
    
    total_params = sum(p.numel() for p in baseline_model.parameters())
    print(f"Baseline Parameters: {total_params:,}")
    
    baseline_history, baseline_best = train_with_singularity(
        baseline_model, trainloader, testloader, device, epochs=50
    )
    
    # Results Summary
    print("\n" + "=" * 60)
    print("RESULTS SUMMARY")
    print("=" * 60)
    print(f"\nSM-ODE Model (Singularity Memory): {smode_best:.2f}%")
    print(f"Baseline Model (No Memory):       {baseline_best:.2f}%")
    print(f"Improvement:                       {smode_best - baseline_best:.2f}%")
    
    # Memory analysis
    print("\n--- Singularity Memory Analysis ---")
    print(f"Memory dimension: 128")
    print(f"Memory decay: 0.95 (high retention)")
    print(f"Shared singularity blocks: 3")
    print(f"Theory: Memory accumulates across batches, not lost")


if __name__ == '__main__':
    main()