import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# -------------------------------------------------------------------
# 1. Define attractor positions (one per digit) on a circle
# -------------------------------------------------------------------
num_classes = 10
attractors = torch.tensor([
    [np.cos(2*np.pi*i/num_classes), np.sin(2*np.pi*i/num_classes)]
    for i in range(num_classes)
], dtype=torch.float32)   # shape: (10, 2)

# -------------------------------------------------------------------
# 2. Vector field: attraction toward the nearest attractor
#    dz/dt = sum_i (z - a_i) * exp(-||z-a_i||²/(2σ²)) / σ²
# -------------------------------------------------------------------
def vector_field(z, attractors, sigma=0.5):
    """
    z: (batch, 2)
    returns: (batch, 2)  – time derivative
    """
    diff = z[:, None, :] - attractors[None, :, :]          # (B, K, 2)
    dist_sq = torch.sum(diff**2, dim=-1)                   # (B, K)
    weights = torch.exp(-dist_sq / (2 * sigma**2))         # (B, K)
    forces = torch.sum(diff * weights[:, :, None] / (sigma**2), dim=1)
    return forces

# -------------------------------------------------------------------
# 3. ODE integrator (fixed number of Euler steps)
# -------------------------------------------------------------------
def integrate(z0, attractors, steps=20, dt=0.1):
    z = z0
    for _ in range(steps):
        z = z + dt * vector_field(z, attractors)
    return z

# -------------------------------------------------------------------
# 4. Encoder: CNN -> 2D latent code
# -------------------------------------------------------------------
class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        self.fc = nn.Linear(128 * 7 * 7, 2)   # after two stride‑2 convs: 28→14→7

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

# -------------------------------------------------------------------
# 5. Full XYFlow classifier: encoder + integration + logits
# -------------------------------------------------------------------
class XYFlowClassifier(nn.Module):
    def __init__(self, attractors, steps=20, dt=0.1):
        super().__init__()
        self.encoder = Encoder()
        # attractors are fixed (non‑trainable)
        self.register_buffer('attractors', attractors)
        self.steps = steps
        self.dt = dt

    def forward(self, x):
        z0 = self.encoder(x)                     # (B, 2)
        zT = integrate(z0, self.attractors, self.steps, self.dt)
        # logits: negative squared distance to each attractor
        diff = zT[:, None, :] - self.attractors[None, :, :]   # (B, K, 2)
        dist_sq = torch.sum(diff**2, dim=-1)                  # (B, K)
        return -dist_sq   # closer → higher logit

# -------------------------------------------------------------------
# 6. Training setup
# -------------------------------------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = torchvision.datasets.MNIST('../data', train=True, download=True, transform=transform)
test_dataset  = torchvision.datasets.MNIST('../data', train=False, download=True, transform=transform)

batch_size = 128
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader  = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

model = XYFlowClassifier(attractors, steps=20, dt=0.1).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# -------------------------------------------------------------------
# 7. Train for a few epochs
# -------------------------------------------------------------------
epochs = 5
for epoch in range(epochs):
    model.train()
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}")

# -------------------------------------------------------------------
# 8. Test accuracy
# -------------------------------------------------------------------
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        logits = model(images)
        preds = torch.argmax(logits, dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
print(f"Test accuracy: {100.*correct/total:.2f}%")