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
# ----------------------------
def flatten_params(model):
    return torch.cat([p.data.view(-1) for p in model.parameters()])

def unflatten_params(model, flat_vec):
    idx = 0
    for p in model.parameters():
        numel = p.numel()
        p.data.copy_(flat_vec[idx:idx + numel].view(p.shape))
        idx += numel

def count_params(model):
    return sum(p.numel() for p in model.parameters())

def project_to_sphere(model, W0, radius):
    """Hard projection: rescale (W - W0) to have norm = 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:
        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)

def project_grad_to_tangent(model, W0):
    """
    FIX #1: Riemannian gradient projection.
    Before the optimizer step, project each parameter's gradient onto the
    tangent space of the sphere at the current point W.
    Tangent space at W: { v : (W - W0) · v = 0 }
    Projection: grad_tangent = grad - ((W-W0) · grad / ||(W-W0)||^2) * (W-W0)
    """
    W_flat = flatten_params(model)
    diff = W_flat - W0  # shape [p]
    norm_sq = (diff * diff).sum()

    # Collect flat gradient
    grads = []
    for p in model.parameters():
        if p.grad is not None:
            grads.append(p.grad.data.view(-1))
        else:
            grads.append(torch.zeros(p.numel(), device=W0.device))
    flat_grad = torch.cat(grads)

    # Project out the radial component
    if norm_sq > 1e-12:
        radial_coeff = (diff * flat_grad).sum() / norm_sq
        flat_grad_tangent = flat_grad - radial_coeff * diff
    else:
        flat_grad_tangent = flat_grad

    # Write back
    idx = 0
    for p in model.parameters():
        numel = p.numel()
        if p.grad is not None:
            p.grad.data.copy_(flat_grad_tangent[idx:idx + numel].view(p.shape))
        idx += numel

# ----------------------------
# 3. Training and evaluation
# ----------------------------
def train(model, device, train_loader, optimizer, epoch, W0, radius, project_every=10):
    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()

        # FIX #1: Project gradient to tangent space before optimizer step
        project_grad_to_tangent(model, W0)

        optimizer.step()

        # FIX #3: Project to sphere less frequently (every N batches)
        # This lets Adam accumulate meaningful momentum between projections
        if (batch_idx + 1) % project_every == 0:
            project_to_sphere(model, W0, radius)

        if batch_idx % 100 == 0:
            W_flat = flatten_params(model)
            dist = torch.norm(W_flat - W0).item()
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\t'
                  f'Loss: {loss.item():.6f}\t||W-W0||: {dist:.4f}')

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: Avg loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
    return accuracy

def main():
    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,))
    ])

    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)

    model = MLP().to(device)

    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)

    W0 = flatten_params(model).clone().detach()
    total_params = count_params(model)

    # FIX #2: Scale radius to match the natural weight norm.
    # Kaiming init gives ||W|| ~ sqrt(p). We use pi * sqrt(p) as the radius
    # so the constraint doesn't crush the weights into a tiny shell.
    # Alternatively set radius = torch.norm(W0).item() to start at natural scale.
    natural_norm = torch.norm(W0).item()
    radius = natural_norm  # start at the natural scale; pi encodes the shape, not the scale
    print(f"Total params: {total_params}, natural norm: {natural_norm:.4f}, radius: {radius:.4f}")

    # Project once to start exactly on the sphere
    project_to_sphere(model, W0, radius)

    # Use Adam with a moderate LR; SGD+momentum also works well here
    optimizer = optim.Adam(model.parameters(), lr=5e-4)

    # FIX #4: Train for enough epochs (10, not 1)
    num_epochs = 10
    for epoch in range(1, num_epochs + 1):
        train(model, device, train_loader, optimizer, epoch, W0, radius, project_every=10)
        accuracy = test(model, device, test_loader)

    # Final report
    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}")

    W_flat = flatten_params(model)
    diff_norm = torch.norm(W_flat - W0).item()
    print(f"Constraint ||W - W0||_2 = {diff_norm:.10f} (target radius = {radius:.10f})")

if __name__ == "__main__":
    main()