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.
    Fixed: Dynamic trajectory buffer initialization to handle varying batch sizes.
    """
    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

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

        # The 'Memory Kernel' M(t)
        self.memory_kernel = nn.Parameter(torch.randn(memory_size, 1))
        
        # The 'Resonance' term
        self.resonance = nn.Parameter(torch.zeros(out_features))

        # Initialize as None; we will create the buffer on the first forward pass
        self.trajectory_buffer = None

    def forward(self, x):
        batch_size = x.size(0)
        
        # 1. Standard projection (The snapshot)
        current_proj = torch.matmul(x, self.weight) + self.bias
        
        # 2. Dynamic Trajectory Buffer Management
        # Initialize buffer on the first pass to match the current batch size
        if self.trajectory_buffer is None or self.trajectory_buffer.size(0) != batch_size:
            # Create a zero buffer of shape [batch, memory_size, out_features]
            self.trajectory_buffer = torch.zeros(batch_size, self.memory_size, self.out_features, device=x.device)
        
        # Update Trajectory: Slide window left and append the current projection to the end
        # [batch, 1:mem_size, out_feat] + [batch, 1, out_feat]
        self.trajectory_buffer = torch.cat([
            self.trajectory_buffer[:, 1:, :], 
            current_proj.unsqueeze(1)
        ], dim=1)
        
        # 3. Causal Integration (The 'Star Product')
        # bmm performs batch matrix multiplication
        # trajectory_buffer: [batch, mem_size, out_features]
        # memory_kernel: [mem_size, 1] (expanded to [batch, mem_size, 1])
        memory_contribution = torch.bmm(
            self.trajectory_buffer, 
            self.memory_kernel.unsqueeze(0).expand(batch_size, -1, -1)
        )
        memory_contribution = memory_contribution.squeeze(2) # [batch, out_features]

        # 4. Combine: Snapshot + Memory + Resonance
        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 (MNIST) -> 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):
        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
    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()
    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"\nTest Accuracy: {100. * correct / len(test_loader.dataset):.2f}%")

if __name__ == "__main__":
    run_experiment()