#!/usr/bin/env python3
"""
PASM-MLP Proof-of-Concept
----------------------------
A Probabilistic-Assembly MLP for CIFAR-10.
Each layer is a "joint" that forks into K particular solutions by perturbing
weights from a base (stationary) distribution.  The final readout collapses
the K trajectories using the geometric-mean Truth operator:
    log P_truth = mean_k log_softmax(logits_k)  (over K particular solutions)
                - logsumexp( ... )
This implements the PASM product formula in differentiable form.
"""

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

# -----------------------------------------------------------------------------
# Hyperparameters
# -----------------------------------------------------------------------------
BATCH_SIZE      = 100
EPOCHS          = 5
LR              = 1e-3
N_PARTICULAR    = 50 #3      # K particular solutions per joint (like K universes)
NOISE_SCALE     = 0.08       # std of weight perturbation at each joint
HIDDEN_DIMS     = [100] # Slightly larger for CIFAR
DEVICE          = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -----------------------------------------------------------------------------
# PASM Layer: Probability in the Joints
# -----------------------------------------------------------------------------
class PASMLinear(nn.Module):
    """
    A Linear layer that treats its weight matrix as a probabilistic field.
    On every forward pass it samples K particular weight/bias realizations
    (particular solutions) and computes K parallel outputs.
    Input:  (B, in)   at layer 0, or (K, B, in) from a previous PASM joint.
    Output: (K, B, out)
    """
    def __init__(self, in_features: int, out_features: int,
                 n_particular: int = N_PARTICULAR, noise_scale: float = NOISE_SCALE):
        super().__init__()
        self.in_features  = in_features
        self.out_features = out_features
        self.n_particular = n_particular
        self.noise_scale  = noise_scale

        # Stationary base weights (the "Law" in ODE-CCT)
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.01)
        self.bias   = nn.Parameter(torch.zeros(out_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # If input is still 2D, expand into particular-solution space
        if x.dim() == 2:
            if self.n_particular == 1:
                return F.linear(x, self.weight, self.bias)
            x = x.unsqueeze(0).expand(self.n_particular, -1, -1)   # (K, B, in)

        assert x.dim() == 3 and x.size(0) == self.n_particular, \
            "PASMLinear expects (K,B,in) with K == n_particular"

        K = self.n_particular

        # --- Probability injection at the joint ---------------------------
        # Sample K particular realizations around the stationary weight field.
        # This is the "probability in the joints" — each realization is one
        # particular solution of the function pattern at this batch iteration.
        W_noise = torch.randn_like(self.weight.unsqueeze(0).expand(K, -1, -1)) * self.noise_scale
        b_noise = torch.randn_like(self.bias.unsqueeze(0).expand(K, -1))       * self.noise_scale

        W_k = self.weight.unsqueeze(0) + W_noise   # (K, out, in)
        b_k = self.bias.unsqueeze(0)   + b_noise     # (K, out)
        # ------------------------------------------------------------------

        # Parallel forward pass of all K particular solutions
        out = torch.einsum('koi,kbi->kbo', W_k, x) + b_k.unsqueeze(1)  # (K, B, out)
        return out


# -----------------------------------------------------------------------------
# Truth Collapse (PASM MEAS / Final Formula)
# -----------------------------------------------------------------------------
class TruthCollapse(nn.Module):
    """
    Collapses K particular logit fields into a single log-probability field.
    Implements the product-formula limit:
        log P_truth(s) = (1/K) * sum_k log P_k(s)  -  logsumexp(...)
    where P_k(s) = softmax(logits_k)[s].
    This is the geometric-mean consensus across particular solutions.
    """
    def __init__(self, eps: float = 1e-8):
        super().__init__()
        self.eps = eps

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        # logits: (K, B, C)
        log_P = F.log_softmax(logits, dim=-1)          # (K, B, C)
        log_geo_mean = log_P.mean(dim=0)               # (B, C)

        # Normalize so that sum_s exp(log_collapse[s]) == 1
        log_collapse = log_geo_mean - torch.logsumexp(log_geo_mean, dim=-1, keepdim=True)
        return log_collapse                             # (B, C)  -> log-probabilities


# -----------------------------------------------------------------------------
# PASM-MLP Model
# -----------------------------------------------------------------------------
class PASMMLP(nn.Module):
    def __init__(self, input_dim: int = 3072, num_classes: int = 10):
        super().__init__()
        dims = [input_dim] + HIDDEN_DIMS + [num_classes]

        self.joints = nn.ModuleList()
        for i in range(len(dims) - 1):
            is_last = (i == len(dims) - 2)
            self.joints.append(PASMLinear(dims[i], dims[i + 1]))
            if not is_last:
                self.joints.append(nn.ReLU())   # elementwise; handles (K,B,H) natively

        self.collapse = TruthCollapse()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)          # flatten CIFAR (B, 3072)
        for joint in self.joints:
            x = joint(x)                   # shape becomes / stays (K, B, H)
        return self.collapse(x)            # (B, 10) log-probs


# -----------------------------------------------------------------------------
# Baseline MLP (deterministic, same architecture size for comparison)
# -----------------------------------------------------------------------------
class BaselineMLP(nn.Module):
    def __init__(self, input_dim: int = 3072, num_classes: int = 10):
        super().__init__()
        layers = []
        dims = [input_dim] + HIDDEN_DIMS + [num_classes]
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i + 1]))
            if i < len(dims) - 2:
                layers.append(nn.ReLU())
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)
        return F.log_softmax(self.net(x), dim=-1)   # log-probs for NLLLoss


# -----------------------------------------------------------------------------
# Training / Evaluation helpers
# -----------------------------------------------------------------------------
def run_epoch(model, loader, optimizer=None):
    is_train = optimizer is not None
    model.train() if is_train else model.eval()

    total_loss = 0.0
    correct = 0
    total = 0

    with torch.set_grad_enabled(is_train):
        for i, (data, target) in enumerate(loader):
            data, target = data.to(DEVICE), target.to(DEVICE)
        
            for _ in range(100):
                log_probs = model(data)

                loss = F.nll_loss(log_probs, target)
                if is_train:
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
            
            total_loss += loss.item() * data.size(0)
            pred = log_probs.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += data.size(0)
            
            if i % 10 == 0:
                print(f"Batch {i}, Loss: {loss.item():.4f}")

    return total_loss / total, correct / total


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------
def main():
    # --- CIFAR-10 loaders ----------------------------------------------------
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])

    train_ds = datasets.CIFAR10(root='../data', train=True, download=True, transform=transform)
    test_ds  = datasets.CIFAR10(root='../data', train=False, download=True, transform=transform)

    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,  num_workers=0)
    test_loader  = DataLoader(test_ds,  batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    # --- PASM-MLP ------------------------------------------------------------
    print("=" * 60)
    print("Training PASM-MLP (CIFAR-10) (K={}, noise={})".format(N_PARTICULAR, NOISE_SCALE))
    print("=" * 60)
    pasm_model = PASMMLP().to(DEVICE)
    pasm_opt   = torch.optim.Adam(pasm_model.parameters(), lr=LR)

    for epoch in range(1, EPOCHS + 1):
        tr_loss, tr_acc = run_epoch(pasm_model, train_loader, pasm_opt)
        te_loss, te_acc = run_epoch(pasm_model, test_loader)
        print(f"Epoch {epoch:02d} | Train Loss: {tr_loss:.4f} Acc: {tr_acc:.4f} | "
              f"Test Loss: {te_loss:.4f} Acc: {te_acc:.4f}")

    # --- Baseline MLP --------------------------------------------------------
    print("\n" + "=" * 60)
    print("Training Baseline MLP (deterministic)")
    print("=" * 60)
    base_model = BaselineMLP().to(DEVICE)
    base_opt   = torch.optim.Adam(base_model.parameters(), lr=LR)

    for epoch in range(1, EPOCHS + 1):
        tr_loss, tr_acc = run_epoch(base_model, train_loader, base_opt)
        te_loss, te_acc = run_epoch(base_model, test_loader)
        print(f"Epoch {epoch:02d} | Train Loss: {tr_loss:.4f} Acc: {tr_acc:.4f} | "
              f"Test Loss: {te_loss:.4f} Acc: {te_acc:.4f}")

    print("\nDone.")


if __name__ == "__main__":
    main()