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
        
        # Instead of raw weights, we store parameters for functions
        # Each connection (out, in) has a 'Function Type' and 'Function Params'
        # params: [out, in, 3] -> e.g., (amplitude, frequency/rate, phase/offset)
        self.func_params = nn.Parameter(torch.randn(out_features, in_features, 3) * 0.1)
        self.bias_particles = nn.Parameter(torch.randn(out_features, 5) * 0.1)

    def forward(self, x):
        # x shape: [batch, in_features, 5]
        # self.func_params shape: [out_features, in_features, 3]
        
        # 1. We treat 'weights' as functions acting on the input particles.
        # We decompose the input x into its collapsed value Psi(P) for the function's core.
        psi_x = x.sum(dim=-1) # [batch, in_features]
        
        # We define our 'Weight Subroutines' acting on the input
        # Weight i,j is not a number, but a function: W(psi_x)
        
        # Functional Weights:
        # Param 0: Amplitude (a), Param 1: Frequency/Scale (w), Param 2: Phase (p)
        a = self.func_params[..., 0] # [out, in]
        w = self.func_params[..., 1] # [out, in]
        p = self.func_params[..., 2] # [out, in]
        
        # Interaction 1: Linear Weighting (Classical)
        # W_linear = a * psi_x
        lin_interaction = torch.einsum('bn,on->bo', psi_x, a)
        
        # Interaction 2: Oscillatory Weighting (The Wave Function Weight)
        # W_osc = a * sin(w * psi_x + p)
        # We broadcast psi_x [b, n] and w [o, n] to get [b, o, n]
        # Then sum over n.
        wave_input = (w.unsqueeze(0) * psi_x.unsqueeze(1)) + p.unsqueeze(0)
        osc_interaction = torch.sum(a.unsqueeze(0) * torch.sin(wave_input), dim=-1)
        
        # Interaction 3: Soft-Exponential Weighting (The Decay Weight)
        # W_exp = a * exp(-|w * psi_x|)
        exp_interaction = torch.sum(a.unsqueeze(0) * torch.exp(-torch.abs(w.unsqueeze(0) * psi_x.unsqueeze(1))), dim=-1)

        # 2. Map these functional results into the PNS multi-field state
        # we assign each functional interaction to a specific field
        fields = [
            lin_interaction,          # Linear Field (L)
            lin_interaction**2,       # Quadratic Field (Q)
            osc_interaction,         # Oscillatory Field (O)
            exp_interaction,         # Exponential Field (E)
            torch.mean(lin_interaction, dim=1, keepdim=True).expand_as(lin_interaction) # Rest Mass (rho)
        ]
        
        # Stack fields back into [batch, out_features, 5]
        out_particles = torch.stack(fields, dim=-1)
        
        # Add Bias (DC Voltage Source)
        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()