"""
Fast, accurate MNIST HPN classifier.

Fixes:
  1. Single hidden layer (two HPN layers with k=3 kill gradients via w^4 scaling).
  2. Small first-layer init (w=0.05) lets BatchNorm amplify the weak w^2 gradients.
  3. Larger head init (w=0.15) gives healthy logits.
  4. Algebraic shortcut kept (x^k @ (w^k/v^{k-1})^T) — the bug was init/depth, not math.
"""

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


# ───────────────────────── HPN Layer ─────────────────────────
class HPNLayer(nn.Module):
    """
    Homogeneous Power Node layer.
    Forward: y = x^k @ (w^k / v^{k-1})^T + c
    """
    def __init__(self, in_features: int, out_features: int, k: int = 3):
        super().__init__()
        self.k = k
        self.w = nn.Parameter(torch.randn(out_features, in_features))
        self.v_raw = nn.Parameter(torch.zeros(out_features, in_features))
        self.c = nn.Parameter(torch.zeros(out_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # v > 0, with a floor to avoid division by zero
        v = F.softplus(self.v_raw) + 1e-4
        eff = self.w.pow(self.k) / v.pow(self.k - 1)   # [out, in]
        return x.pow(self.k) @ eff.t() + self.c


# ───────────────────────── Model ─────────────────────────
class HPNNet(nn.Module):
    def __init__(self, k: int = 3, hidden: int = 1024):
        super().__init__()
        self.fc1 = HPNLayer(784, hidden, k=k)
        self.bn1 = nn.BatchNorm1d(hidden)
        self.fc2 = HPNLayer(hidden, 10, k=k)

        # Layer-aware init:
        #   - First layer: very small w so BN amplifies the weak w^2 gradients.
        #   - Head: larger w so logits start with healthy variance (~2-3).
        with torch.no_grad():
            nn.init.normal_(self.fc1.w, 0.0, 0.05)
            nn.init.normal_(self.fc2.w, 0.0, 0.15)

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


# ───────────────────────── Data ─────────────────────────
def get_loaders(batch_size: int = 128):
    train_tf = transforms.Compose([
        transforms.RandomRotation(5),
        transforms.RandomAffine(degrees=0, translate=(0.05, 0.05)),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    test_tf = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_ds = datasets.MNIST("../data", train=True, download=True, transform=train_tf)
    test_ds  = datasets.MNIST("../data", train=False, download=True, transform=test_tf)

    kwargs = {"num_workers": 2, "pin_memory": True, "persistent_workers": True}
    if not torch.cuda.is_available():
        kwargs = {}

    return (
        DataLoader(train_ds, batch_size=batch_size, shuffle=True, **kwargs),
        DataLoader(test_ds,  batch_size=512, shuffle=False, **kwargs),
    )


# ───────────────────────── Train + Test ─────────────────────────
def run(k: int = 3, epochs: int = 15, lr: float = 1e-2, hidden: int = 1024):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device} | HPN degree k={k} | hidden={hidden}")

    torch.backends.cudnn.benchmark = True

    train_loader, test_loader = get_loaders()
    model = HPNNet(k=k, hidden=hidden).to(device)

    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr, epochs=epochs, steps_per_epoch=len(train_loader)
    )
    criterion = nn.CrossEntropyLoss()

    for epoch in range(1, epochs + 1):
        model.train()
        total_loss = 0.0
        for x, y in train_loader:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            optimizer.zero_grad()
            logits = model(x)
            loss = criterion(logits, y)
            loss.backward()
            optimizer.step()
            scheduler.step()
            total_loss += loss.item() * x.size(0)

        avg_loss = total_loss / len(train_loader.dataset)

        model.eval()
        correct = 0
        with torch.no_grad():
            for x, y in test_loader:
                x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
                pred = model(x).argmax(dim=1)
                correct += (pred == y).sum().item()
        acc = 100.0 * correct / len(test_loader.dataset)
        print(f"Epoch {epoch:02d}/{epochs}  loss={avg_loss:.4f}  test_acc={acc:.2f}%")

    print("Done.")
    return model


if __name__ == "__main__":
    # k=3 is required to break the ~91% ceiling (odd k restores signed weights).
    run(k=3, epochs=15, hidden=1024)