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

class LinearGraphicalClassifier(torch.nn.Module):
    """
    Linear Graphical Classifier for MNIST in PyTorch.

    Architecture:
        scores = X @ P.T
        probs = softmax(scores)

    Dual‑gradient update (applied manually):
        total_grad = grad_linear + lam * grad_pixel
        where grad_linear  = label_error.T @ X
              grad_pixel   = label_error.T @ X - P * sum(label_error)
    """
    def __init__(self, num_classes=10, input_size=784, lam=0.1, lr=0.01):
        super().__init__()
        self.lam = lam
        # Prototypes initialised uniformly in [0,1] – same image space as data
        self.P = torch.nn.Parameter(
            torch.rand(num_classes, input_size)   # uniform [0,1)
        )
        self.optimizer = optim.SGD(self.parameters(), lr=lr)

    def forward(self, x):
        """Return scores and softmax probabilities."""
        scores = x @ self.P.T          # (N, C)
        probs = F.softmax(scores, dim=1)
        return scores, probs

    def predict(self, x):
        """Return predicted class indices."""
        scores = x @ self.P.T
        return torch.argmax(scores, dim=1)

    def update(self, x, y):
        """
        Perform one batch update using dual‑gradient rule.

        Args:
            x: batch of images, shape (N, 784), values in [0,1]
            y: true labels, shape (N,)
        Returns:
            loss: average cross‑entropy loss for the batch
        """
        self.train()                   # set train mode
        N = x.size(0)

        # Forward pass
        scores, probs = self.forward(x)
        loss = F.cross_entropy(scores, y, reduction='mean')

        # Standard linear gradient via autograd
        self.zero_grad()
        loss.backward()                # now self.P.grad holds dloss/dP (grad_linear)

        # Graphical pixel error term (manually computed)
        with torch.no_grad():
            y_onehot = F.one_hot(y, num_classes=self.P.shape[0]).float()
            label_error = probs - y_onehot          # (N, C)

            # sum_error_per_class = Σ_i label_error[i,:]  → shape (C,)
            sum_error_per_class = label_error.sum(dim=0)   # (C,)

            # grad_pixel = (label_error.T @ X) - P * sum_error_per_class
            # Note: label_error.T @ X  is exactly the same as dloss/dP (the standard gradient)
            grad_pixel = self.P.grad - self.P * sum_error_per_class.unsqueeze(1)

            # Combined gradient
            self.P.grad = self.P.grad + self.lam * grad_pixel

        # SGD step
        self.optimizer.step()

        return loss.item()


# --------------------------
# Data loading & training
# --------------------------
def main():
    # Hyperparameters
    batch_size = 128
    epochs = 20
    lam = 0.005          # reduced graphical weight (was 0.1)
    lr = 0.001           # reduced learning rate

    # MNIST pipeline: flatten, normalize to [0,1]
    transform = transforms.Compose([
        transforms.ToTensor(),                     # [0,1] automatically
        transforms.Lambda(lambda x: x.view(-1))    # flatten 28x28 -> 784
    ])
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_dataset  = datasets.MNIST(root='./data', train=False, download=True, 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
    model = LinearGraphicalClassifier(num_classes=10, input_size=784, lam=lam, lr=lr)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    # Training loop
    for epoch in range(1, epochs+1):
        total_loss = 0.0
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            loss = model.update(x, y)
            total_loss += loss * x.size(0)

        # Evaluate
        model.eval()
        with torch.no_grad():
            # Train accuracy (use a subset for speed if needed, but full set is fine)
            train_correct = 0
            for x, y in train_loader:
                x, y = x.to(device), y.to(device)
                pred = model.predict(x)
                train_correct += (pred == y).sum().item()
            train_acc = train_correct / len(train_dataset)

            test_correct = 0
            for x, y in test_loader:
                x, y = x.to(device), y.to(device)
                pred = model.predict(x)
                test_correct += (pred == y).sum().item()
            test_acc = test_correct / len(test_dataset)

        avg_loss = total_loss / len(train_dataset)
        print(f"Epoch {epoch:2d}/{epochs} | Loss: {avg_loss:.4f} | "
              f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    print("Training finished.")

if __name__ == "__main__":
    main()