"""
Fast, accurate MNIST classifier using Homogeneous Power Node (HPN) layers.

Key improvements:
  1. Algebraic reduction: HPN degree k collapses to a single matmul on x^k,
     eliminating the O(batch*in*out) einsum memory blow-up.
  2. Odd k (default 3) restores signed effective weights, lifting the ~91%
     ceiling that k=2 hits due to non-negative weights.
  3. Modern training: BatchNorm, Dropout, AdamW, OneCycleLR, wider net,
     data augmentation, label smoothing.

HPN layer (degree k):  y_j = Σ_i (w_ji·x_i)^k / (v_ji)^{k-1} + c_j
Optimized:            y = x^k @ (w^k / v^{k-1})^T + c
"""

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 (fast matmul) ─────────────────────────
class HPNLayer(nn.Module):
    """
    Homogeneous Power Node layer of degree k.
    Forward pass reduced to a single matrix multiplication on x^k.
    """
    def __init__(self, in_features: int, out_features: int, k: int = 3):
        super().__init__()
        self.k = k
        # He-style init scaled for the k-th power so that effective weights
        # have variance ~ 1 / fan_in.
        std = (1.0 / in_features) ** (1.0 / (2 * k))
        self.w = nn.Parameter(torch.randn(out_features, in_features) * std)
        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:
        # Safe positive denominator: softplus can approach 0, so add a floor.
        v = F.softplus(self.v_raw) + 1e-4
        # Effective linear weights: A_ji = w_ji^k / v_ji^{k-1}
        eff = self.w.pow(self.k) / v.pow(self.k - 1)   # [out, in]
        # The entire HPN collapses to a matmul on x^k
        return x.pow(self.k) @ eff.t() + self.c


# ───────────────────────── Model ─────────────────────────
class HPNBlock(nn.Module):
    """HPN layer + BatchNorm + ReLU + Dropout."""
    def __init__(self, in_features: int, out_features: int, k: int = 3, dropout: float = 0.2):
        super().__init__()
        self.hpn = HPNLayer(in_features, out_features, k=k)
        self.bn = nn.BatchNorm1d(out_features)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.hpn(x)
        x = self.bn(x)
        x = F.relu(x)
        x = self.dropout(x)
        return x


class HPNNet(nn.Module):
    def __init__(self, k: int = 3, hidden: int = 512, dropout: float = 0.2):
        super().__init__()
        self.block1 = HPNBlock(784, hidden, k=k, dropout=dropout)
        self.block2 = HPNBlock(hidden, hidden, k=k, dropout=dropout)
        self.head = HPNLayer(hidden, 10, k=k)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)
        x = self.block1(x)
        x = self.block2(x)
        return self.head(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} if torch.cuda.is_available() else {}
    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-3, hidden: int = 512, batch_size: int = 128):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device} | HPN degree k={k} | hidden={hidden}")

    # Speed tweaks
    torch.backends.cudnn.benchmark = True
    if hasattr(torch, "set_float32_matmul_precision"):
        torch.set_float32_matmul_precision("high")

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

    if hasattr(torch, "compile"):
        model = torch.compile(model)  # PyTorch 2.x graph fusion

    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(label_smoothing=0.1)

    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 strongly recommended: odd powers keep signed effective weights,
    # breaking the ~91% ceiling that k=2 hits due to non-negative weights.
    run(k=3, epochs=15, hidden=512)