import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np

# ------------------------------------------------------------
# Complex-valued linear layer + modulus non-linearity
# ------------------------------------------------------------
class ComplexWaveLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # Trainable complex weight matrix W = real + i*imag
        self.real_weight = nn.Parameter(torch.randn(out_features, in_features) * 0.1)
        self.imag_weight = nn.Parameter(torch.randn(out_features, in_features) * 0.1)

    def forward(self, x):
        """
        x : real tensor of shape (batch, in_features)
        Returns magnitude (modulus) after complex linear transform.
        """
        # Treat input as purely real (imag part = 0)
        complex_input = x.to(dtype=torch.complex64)  # (batch, in_features) real only
        complex_weight = torch.complex(self.real_weight, self.imag_weight)  # (out, in)
        # Complex matrix multiplication: out = input * weight^T
        # Shape: (batch, out_features) complex
        complex_out = torch.matmul(complex_input, complex_weight.T)
        # Modulus = sqrt(Re² + Im²) → wave intensity
        magnitude = torch.abs(complex_out)
        return magnitude

# ------------------------------------------------------------
# Complete model: two complex wave layers (interference cascade)
# ------------------------------------------------------------
class WaveInterferenceMNIST(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=128, num_classes=10):
        super().__init__()
        self.layer1 = ComplexWaveLayer(input_dim, hidden_dim)
        self.layer2 = ComplexWaveLayer(hidden_dim, num_classes)

    def forward(self, x):
        x = x.view(x.size(0), -1)      # flatten (batch, 784)
        x = self.layer1(x)             # (batch, hidden_dim) real, positive
        x = self.layer2(x)             # (batch, 10) real, positive
        # No softmax here – CrossEntropyLoss expects raw logits
        # But we can optionally take log(x+eps) if needed; here raw works.
        return x

# ------------------------------------------------------------
# Load MNIST
# ------------------------------------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

trainset = torchvision.datasets.MNIST(root='../data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='../data', train=False, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

# ------------------------------------------------------------
# Training (both layers are trainable – evolution from frozen wave shaper)
# ------------------------------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = WaveInterferenceMNIST(hidden_dim=128).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in trainloader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader):.4f}")

# ------------------------------------------------------------
# Evaluation
# ------------------------------------------------------------
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"\nTest accuracy: {100 * correct / total:.2f}%")