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

# ----------------------------
# 1. Define the MLP model
# ----------------------------
class MLP(nn.Module):
    def __init__(self, input_size=784, hidden1=256, hidden2=128, num_classes=10):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden1)
        self.fc2 = nn.Linear(hidden1, hidden2)
        self.fc3 = nn.Linear(hidden2, num_classes)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# ----------------------------
# 2. Helper functions for constraint
# ----------------------------
def flatten_params(model):
    """Return a flat vector of all parameters."""
    return torch.cat([p.data.view(-1) for p in model.parameters()])

def unflatten_params(model, flat_vec):
    """Assign a flat vector back to the model's parameters."""
    idx = 0
    for p in model.parameters():
        numel = p.numel()
        p.data.copy_(flat_vec[idx:idx+numel].view(p.shape))
        idx += numel

def project_to_sphere(model, W0, radius=math.pi):
    """
    Project the model's parameters onto the sphere
    ||W - W0||_2 = radius.
    """
    W_flat = flatten_params(model)
    diff = W_flat - W0
    norm = torch.norm(diff)
    if norm > 1e-12:
        new_flat = W0 + radius * diff / norm
    else:
        # If diff is zero (should not happen after initialization),
        # add a random small perturbation along a random direction.
        rand_dir = torch.randn_like(W0)
        rand_dir = rand_dir / torch.norm(rand_dir)
        new_flat = W0 + radius * rand_dir
    unflatten_params(model, new_flat)

# ----------------------------
# 3. Training and evaluation
# ----------------------------
def train(model, device, train_loader, optimizer, epoch, W0, radius):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

        # --- Enforce the π-constraint after each update ---
        project_to_sphere(model, W0, radius)

        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
    return accuracy

def main():
    # ----------------------------
    # 4. Setup data and device
    # ----------------------------
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean/std
    ])

    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=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

    # ----------------------------
    # 5. Create model and initialize
    # ----------------------------
    model = MLP().to(device)
    # Use a standard initialization (Kaiming normal for ReLU)
    def init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
            nn.init.zeros_(m.bias)
    model.apply(init_weights)

    # Save the initial parameter vector W0
    W0 = flatten_params(model).clone().detach()
    # Ensure the initial distance is already π (we'll project once)
    project_to_sphere(model, W0, radius=math.pi)

    # Re-flatten to get the adjusted W0? Actually W0 should remain the initial center.
    # But we projected the model to be on the sphere centered at W0, so now the distance is π.
    # We keep W0 as the initial (pre-projection) values? The theory says W0 is the initial
    # weights before training; we want the constraint ||W - W0|| = π. So we set W0 to the
    # original random weights, and then we project after each step. We already projected once
    # to ensure we start on the sphere.

    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # ----------------------------
    # 6. Train and evaluate
    # ----------------------------
    num_epochs = 1
    for epoch in range(1, num_epochs + 1):
        train(model, device, train_loader, optimizer, epoch, W0, radius=math.pi)
        accuracy = test(model, device, test_loader)

    # Final test accuracy with high precision
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
    final_acc = correct / total
    print(f"\nFinal Test Accuracy: {final_acc:.10f}")  # 10 decimal places

    # Also print the constraint violation
    W_flat = flatten_params(model)
    diff_norm = torch.norm(W_flat - W0).item()
    print(f"Constraint ||W - W0||_2 = {diff_norm:.10f} (should be {math.pi:.10f})")

if __name__ == "__main__":
    main()