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

# =============================================================================
# PNS CIRCUIT ELEMENTS
# =============================================================================

class PNSCircuitLayer(nn.Module):
    """
    A Layer that treats data as 'Particle Numbers'.
    Internal state is a multi-field tensor: [batch, channels, fields]
    Fields: 0: Linear (L), 1: Quadratic (Q), 2: Oscillatory (O), 3: Exponential (E), 4: Rest Mass (rho)
    """
    def __init__(self, in_features, out_features):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        
        # The Weight Matrix is also a Particle: each weight has its own field excitations
        # Shape: [out, in, 5] -> 5 fields for each weight
        self.weight_particles = nn.Parameter(torch.randn(out_features, in_features, 5) * 0.1)
        self.bias_particles = nn.Parameter(torch.randn(out_features, 5) * 0.1)

    def forward(self, x):
        # x shape: [batch, in_features, 5]
        # weight_particles shape: [out_features, in_features, 5]
        
        # We compute the coupling for each field independently
        # For each field f in 0..4: output[f] = sum(weight[f] * x[f])
        # Equation 'bn,on->bo' means:
        # b = batch, n = in_features, o = out_features
        
        fields = []
        for f in range(5):
            field_val = torch.einsum('bn,on->bo', x[..., f], self.weight_particles[..., f])
            fields.append(field_val)
            
        # Stack fields back into [batch, out_features, 5]
        out_particles = torch.stack(fields, dim=-1)
        
        # Add Bias (DC Voltage Source) - bias_particles is [out_features, 5]
        out_particles += self.bias_particles
        
        return out_particles

class PNSDiode(nn.Module):
    """
    ReLU equivalent: Only allows 'positive charge' particles to flow.
    """
    def forward(self, x):
        # Collapses the particle to its scalar value to determine the charge sign
        # psi = L + Q + O + E + rho
        psi = x.sum(dim=-1, keepdim=True)
        
        # Diode gating: if psi < 0, the particle is annihilated to vacuum (0)
        gate = (psi > 0).float()
        return x * gate

class PNSImpedanceMatcher(nn.Module):
    """
    LayerNorm equivalent: Normalizes the energy (Quadratic field) to 1.
    """
    def forward(self, x):
        # Normalize using the energy field (index 1)
        energy = torch.sqrt(x[..., 1]**2 + 1e-6)
        return x / energy.unsqueeze(-1)

class PNSCollapse(nn.Module):
    """
    The Measurement Operator: Final collapse of all fields into a scalar.
    """
    def forward(self, x):
        # Psi(P) = sum of all fields
        return x.sum(dim=-1)

# =============================================================================
# FULL CIRCUIT MODEL
# =============================================================================

class PNSClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        # Input: 28*28 = 784. We map each pixel to a Particle (5 fields)
        self.layer1 = PNSCircuitLayer(784, 128)
        self.diode1 = PNSDiode()
        self.matcher1 = PNSImpedanceMatcher()
        
        self.layer2 = PNSCircuitLayer(128, 64)
        self.diode2 = PNSDiode()
        self.matcher2 = PNSImpedanceMatcher()
        
        self.layer3 = PNSCircuitLayer(64, 10)
        self.collapse = PNSCollapse()

    def forward(self, x):
        # 1. Injection: Convert scalar pixels into Particle Numbers
        # Map pixel value to the Linear field, others start at zero
        # x: [batch, 784] -> [batch, 784, 5]
        batch_size = x.shape[0]
        p_in = torch.zeros((batch_size, 784, 5), device=x.device)
        p_in[..., 0] = x # Linear field excitation
        p_in[..., 4] = 1.0 # Every pixel starts with unit rest mass rho=1
        
        # 2. Circuit Flow
        x = self.layer1(p_in)
        x = self.diode1(x)
        x = self.matcher1(x)
        
        x = self.layer2(x)
        x = self.diode2(x)
        x = self.matcher2(x)
        
        x = self.layer3(x)
        
        # 3. Measurement (Collapse)
        return self.collapse(x)

# =============================================================================
# TRAINING LOOP
# =============================================================================

def train_mnist():
    # Data
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_set = datasets.MNIST('./data', train=False, transform=transform)
    train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=1000, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = PNSClassifier().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    model.train()
    for epoch in range(2): # Small epoch count for demonstration
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.view(-1, 784).to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            if batch_idx % 200 == 0:
                print(f"Epoch {epoch} [{batch_idx*64}/{len(train_loader)*64}] Loss: {loss.item():.4f}")
        
    # Evaluation
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.view(-1, 784).to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    print(f"\nTest Accuracy: {100. * correct / len(test_loader.dataset):.2f}%")

if __name__ == "__main__":
    train_mnist()