import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# ==========================================
# 1. THE CAUSAL-RESONANCE DOT PRODUCT LAYER
# ==========================================
class CRDPLinear(nn.Module):
    """
    Causal-Resonance Dot Product Linear Layer.
    Causal-Resonance Formula:
    _Out = (x @ W) + Sum_{t=1 to N}(MemoryKernel[t] * Trajectory[t]) + Resonance
    """
    def __init__(self, in_features, out_features, memory_size=5):
        super(CRDPLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.memory_size = memory_size

        # Stationary Weights (The 'Laws')
        self.weight = nn.Parameter(torch.randn(in_features, out_features) * 0.01)
        self.bias = nn.Parameter(torch.zeros(out_features))

        # Memory Kernel M(t) - weights for past states
        # Shape: [1, memory_size] to allow batch multiplication
        self.memory_kernel = nn.Parameter(torch.randn(1, memory_size))
        
        # Resonance term - global attractor shift
        self.resonance = nn.Parameter(torch.zeros(out_features))

        # The buffer stores the trajectory of the flow
        self.trajectory_buffer = None

    def forward(self, x):
        batch_size = x.size(0)
        
        # 1. Snapshot Projection (a . b)
        current_proj = torch.matmul(x, self.weight) + self.bias
        
        # 2. Trajectory Buffer Management
        if self.trajectory_buffer is None or self.trajectory_buffer.size(0) != batch_size:
            self.trajectory_buffer = torch.zeros(batch_size, self.memory_size, self.out_features, device=x.device)
        
        # Shift window and append current projection
        # We use .detach() on the buffer to prevent the computation graph 
        # from growing infinitely over every single batch (BPTT limit)
        updated_buffer = torch.cat([
            self.trajectory_buffer[:, 1:, :].detach(), 
            current_proj.unsqueeze(1)
        ], dim=1)
        self.trajectory_buffer = updated_buffer
        
        # 3. Causal Integration (The 'Star Product')
        # Target: [batch, out_features]
        # trajectory_buffer: [batch, memory_size, out_features]
        # memory_kernel: [1, memory_size] -> expand to [batch, 1, memory_size]
        
        kernel_expanded = self.memory_kernel.unsqueeze(0).expand(batch_size, -1, -1)
        
        # Matrix Multiply: [batch, 1, memory_size] @ [batch, memory_size, out_features]
        # Result: [batch, 1, out_features]
        memory_contribution = torch.bmm(kernel_expanded, self.trajectory_buffer)
        memory_contribution = memory_contribution.squeeze(1) # [batch, out_features]

        # 4. Combine: Current Flow + Memory History + Resonance Attractor
        out = current_proj + memory_contribution + self.resonance
        
        return out

# ==========================================
# 2. THE CRDP-MLP ARCHITECTURE
# ==========================================
class CRDP_MLP(nn.Module):
    def __init__(self):
        super(CRDP_MLP, self).__init__()
        # Architecture: 784 -> 128 -> 64 -> 10
        self.layer1 = CRDPLinear(784, 128)
        self.layer2 = CRDPLinear(128, 64)
        self.layer3 = CRDPLinear(64, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        # Flatten MNIST images
        x = x.view(-1, 784)
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.layer3(x)
        return x

# ==========================================
# 3. TRAINING AND TESTING PIPELINE
# ==========================================
def run_experiment():
    # Hyperparameters
    batch_size = 64
    learning_rate = 0.001
    epochs = 5

    # MNIST Dataset loading
    transform = transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('../data', train=False, transform=transform)

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

    model = CRDP_MLP()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Training Loop
    model.train()
    print("Starting Training with CRDP...")
    for epoch in range(epochs):
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}")

    # Testing Loop
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    print(f"\nFinal Test Accuracy: {100. * correct / len(test_loader.dataset):.2f}%")

if __name__ == "__main__":
    run_experiment()