"""
Minimalist MNIST classifier using Homogeneous Power Node (HPN) layers.

HPN layer (degree k):  y_j = Σ_i (w_ji · x_i)^k / (v_ji)^{k-1}  +  c_j

We use k=2 (the Area-Node / quadratic RAN form):
    y_j = Σ_i (w_ji · x_i)² / v_ji  +  c_j

PyTorch autograd handles all powers and divisions natively.
"""

import torch
import torch.nn as nn
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 of degree k.

        y_j = Σ_i (w_ji · x_i)^k / (v_ji)^{k-1}  +  c_j

    Parameters
    ----------
    in_features  : int
    out_features : int
    k            : int   (power degree, default 2)
    """
    def __init__(self, in_features: int, out_features: int, k: int = 2):
        super().__init__()
        self.k = k
        self.in_features = in_features
        self.out_features = out_features
        # numerator weights — initialise with He-style scaling
        self.w = nn.Parameter(torch.randn(out_features, in_features) * (1.0 / in_features ** 0.5))
        # denominator parameters — initialised to 1, kept positive via softplus
        self.v_raw = nn.Parameter(torch.zeros(out_features, in_features))   # softplus(0)≈0.693
        # scalar bias
        self.c = nn.Parameter(torch.zeros(out_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [batch, in_features]
        v = torch.nn.functional.softplus(self.v_raw)              # [out, in]  always > 0
        # compute (w_ji · x_i)^k for every (batch, out, in)
        wx = torch.einsum("oi,bi->boi", self.w, x)                # [batch, out, in]
        if self.k == 2:
            numerator = wx * wx                                  # avoids complex for k=2
        else:
            numerator = torch.pow(wx, self.k)
        denominator = torch.pow(v, self.k - 1)                   # v^{k-1}
        terms = numerator / denominator                           # [batch, out, in]
        y = terms.sum(dim=2) + self.c                             # [batch, out]
        return y


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)               # flatten 28×28 → 784
        x = self.fc1(x)                         # HPN collapse → logits
        x = torch.relu(x)                        # non-linearity between layers
        x = self.fc2(x)                         # final logits
        return x


# ───────────────────────── Data ─────────────────────────
def get_loaders(batch_size: int = 128):
    tf = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),   # MNIST mean/std
    ])
    train_ds = datasets.MNIST("../data", train=True,  download=True, transform=tf)
    test_ds  = datasets.MNIST("../data", train=False, download=True, transform=tf)
    return (
        DataLoader(train_ds, batch_size=batch_size, shuffle=True),
        DataLoader(test_ds,  batch_size=256,       shuffle=False),
    )


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

    train_loader, test_loader = get_loaders()
    model = HPNNet(k=k).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(1, epochs + 1):
        # ---- Train ----
        model.train()
        total_loss = 0.0
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            logits = model(x)                       # collapsed HPN values as logits
            loss = criterion(logits, y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * x.size(0)
        avg_loss = total_loss / len(train_loader.dataset)

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

    print("Done.")
    return model


if __name__ == "__main__":
    # k=2 → Area-Node (quadratic RAN):  Σ (w·x)² / v  +  c
    # Try k=3 for the cubic Fermat-node form as well!
    run(k=2, epochs=5)