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

# -------------------------------
# 1. Load MNIST
# -------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.view(-1))
])

train_full = datasets.MNIST('../data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('../data', train=False, download=True, transform=transform)
train_set, _ = random_split(train_full, [60000, len(train_full)-60000])

batch_size = 100
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# -------------------------------
# 2. Define learnable parameters w1, w2
# -------------------------------
w1 = nn.Parameter(0.01 * torch.randn(784, 100, device=device))
w2 = nn.Parameter(0.01 * torch.randn(784, 100, device=device))

# Optimizer for w1, w2 (using SGD, same learning rate 0.1 as original)
opt_w = optim.SGD([w1, w2], lr=0.001)

# -------------------------------
# 3. MLP classifier (same as before)
# -------------------------------
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 100),
            nn.ReLU(),
            nn.Linear(100, 10)
        )
    def forward(self, x):
        return self.net(x)

mlp = MLP().to(device)
opt_classifier = optim.SGD(mlp.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Helper: pseudo-inverse
def pinv(A):
    return torch.linalg.pinv(A)

# -------------------------------
# 4. Training loop
# -------------------------------
iteration = 0
while True:
    for X_batch, yt in train_loader:
        X_batch = X_batch.to(device)
        yt = yt.to(device)

        # --- Step 1: Update w1, w2 to minimise reconstruction error ---
        opt_w.zero_grad()

        b = X_batch @ w1                 # (batch, 100)
        w2_pinv = pinv(w2)               # (100, 784)
        a = b @ w2_pinv @ w2             # (batch, 100)

        reconstruction_loss = torch.norm(a - b, p='fro')**2
        reconstruction_loss.backward()
        opt_w.step()

        # --- Step 2: Train MLP on transformed features (every 10 batches) ---
        if iteration % 10 == 0:
            with torch.no_grad():
                forward = X_batch @ w1 @ pinv(w2)   # (batch, 784)

            opt_classifier.zero_grad()
            outputs = mlp(forward)
            class_loss = criterion(outputs, yt)
            class_loss.backward()
            opt_classifier.step()

            # Print stats
            pred = outputs.argmax(dim=1)
            acc = (pred == yt).float().mean().item()
            print(f"Iter {iteration:5d} | recon loss: {reconstruction_loss.item():.4f} | class loss: {class_loss.item():.4f} | batch acc: {acc:.4f}")

        iteration += 1
        if iteration >= 2000:
            break
    else:
        continue
    break