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

# 1. Novel Model: Radial Manifold Projection Network (Simplified)
class RMPNet(nn.Module):
    def __init__(self):
        super(RMPNet, self).__init__()
        # Encoder to project 28x28 image into a latent "coordinate space"
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 2) # Project to a 2D "Manifold Map" (X, Y)
        )
        
        # The "Circles": Learnable centroids for each digit (10 digits x 2 coordinates)
        # These represent the centers of the circles in your image
        self.centroids = nn.Parameter(torch.randn(10, 2))
        # Learnable "radius" or spread for each manifold
        self.sigmas = nn.Parameter(2*torch.ones(10))

    def forward(self, x):
        # Project image to a point in 2D manifold space
        coords = self.encoder(x) # Shape: [batch, 2]
        
        # Calculate distance from the point to each of the 10 digit-centroids
        # This mimics the "Radial" distance in the image
        dist = torch.cdist(coords, self.centroids) # Shape: [batch, 10]
        
        # RBF Activation: exp(-dist^2 / 2*sigma^2)
        # This creates the "Circle" activation effect
        logits = torch.exp(-(dist**2) / (2 * self.sigmas**2))
        
        return logits

# 2. Setup Training & Testing
def experiment():
    # Data loading
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_set = datasets.MNIST('../data', train=True, download=True, transform=transform)
    test_set = datasets.MNIST('../data', train=False, transform=transform)
    train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=1000, shuffle=False)

    model = RMPNet()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    # Training Loop
    model.train()
    for epoch in range(3): # Short run for experiment
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 200 == 0:
                print(f"Epoch {epoch} Batch {batch_idx} Loss: {loss.item():.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"\nExperiment Result: Test Accuracy: {100. * correct / len(test_set):.2f}%")

if __name__ == "__main__":
    experiment()