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

# ==========================================
# 1. RANDOM PROJECTION ENCODER (Symmetry-Preserving)
# ==========================================
class RandomProjectionEncoder(nn.Module):
    """
    The 'Blind Probe'. Projects high-dim data to low-dim space 
    preserving distances. No training required.
    """
    def __init__(self, in_features, out_features):
        super(RandomProjectionEncoder, self).__init__()
        # Generate a random matrix from a Gaussian distribution
        # This matrix is FIXED. It is the 'Stationary' projection.
        proj_matrix = torch.randn(in_features, out_features) / np.sqrt(out_features)
        self.register_buffer('proj_matrix', proj_matrix)

    def forward(self, x):
        # Simple matrix multiply: [batch, 784] @ [784, 64] -> [batch, 64]
        return torch.matmul(x, self.proj_matrix)

# ==========================================
# 2. CAUSAL-RESONANCE DOT PRODUCT LAYER
# ==========================================
class CRDPLinear(nn.Module):
    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

        self.weight = nn.Parameter(torch.randn(in_features, out_features) * 0.01)
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.memory_kernel = nn.Parameter(torch.randn(1, memory_size))
        self.resonance = nn.Parameter(torch.zeros(out_features))
        self.trajectory_buffer = None

    def forward(self, x):
        batch_size = x.size(0)
        current_proj = torch.matmul(x, self.weight) + self.bias
        
        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)
        
        updated_buffer = torch.cat([
            self.trajectory_buffer[:, 1:, :].detach(), 
            current_proj.unsqueeze(1)
        ], dim=1)
        self.trajectory_buffer = updated_buffer
        
        kernel_expanded = self.memory_kernel.unsqueeze(0).expand(batch_size, -1, -1)
        memory_contribution = torch.bmm(kernel_expanded, self.trajectory_buffer).squeeze(1)

        return current_proj + memory_contribution + self.resonance

# ==========================================
# 3. THE INTEGRATED ARCHITECTURE
# ==========================================
class RandomCCT_MLP(nn.Module):
    def __init__(self):
        super(RandomCCT_MLP, self).__init__()
        # STEP 1: Fixed Random Projection (784 -> 64)
        # This is the 'Symmetry-Preserving' Stationary layer.
        self.encoder = RandomProjectionEncoder(784, 64)
        
        # STEP 2: CRDP Layers (Working in the 64-dim projected space)
        self.layer1 = CRDPLinear(64, 32)
        self.layer2 = CRDPLinear(32, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(-1, 784)
        
        # 1. Project to low-dim space (Zero Weights, Zero Gradient)
        x = self.encoder(x) 
        
        # 2. Process via Causal Resonance
        x = self.relu(self.layer1(x))
        x = self.layer2(x)
        return x

# ==========================================
# 4. RUN EXPERIMENT
# ==========================================
def run_experiment():
    batch_size, learning_rate, epochs = 64, 0.001, 5

    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_loader = DataLoader(datasets.MNIST('../data', train=True, download=True, transform=transform), batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(datasets.MNIST('../data', train=False, transform=transform), batch_size=batch_size, shuffle=False)

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

    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for data, target in 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}")

    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"\nTest Accuracy with Random Projection: {100. * correct / len(test_loader.dataset):.2f}%")

if __name__ == "__main__":
    run_experiment()