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

# -------------------------------------------------------------------------
# The SpectralAtom Layer
# -------------------------------------------------------------------------
class SpectralAtom(nn.Module):
    """
    Implementation of an 'Atom' in the Spectral Paradigm.
    An atom consists of a learnable spectral filter (Complex Weights).
    Computation: Y(f) = X(f) * W(f)
    """
    def __init__(self, channels):
        super(SpectralAtom, self).__init__()
        # We represent the complex weights as two real tensors (real and imaginary)
        # The filter size matches the input signal size (28x28 for MNIST)
        self.weight_real = nn.Parameter(torch.randn(channels, 28, 28) * 0.01)
        self.weight_imag = nn.Parameter(torch.randn(channels, 28, 28) * 0.01)

    def forward(self, x_fft):
        # x_fft is expected to be a complex tensor
        # Pointwise multiplication in frequency domain
        w_complex = torch.complex(self.weight_real, self.weight_imag)
        return x_fft * w_complex

# -------------------------------------------------------------------------
# The SpectralComputingNetwork
# -------------------------------------------------------------------------
class SpectralNet(nn.Module):
    def __init__(self, num_classes=10):
        super(SpectralNet, self).__init__()
        # We create 10 'Atoms', one for each digit class.
        # Each atom will attempt to resonate with the frequencies of its class.
        self.atoms = SpectralAtom(num_classes)
        
        # Final projection to flatten the resulting "Cloud" into a class score
        # In a pure spectral model, we look for the dominant mode (max intensity).
        self.classifier = nn.Linear(num_classes * 28 * 28, num_classes)

    def forward(self, x):
        # 1. PROMPT (Nucleus) -> FFT
        # Convert image to complex and perform 2D FFT
        x = x.float()
        x_fft = torch.fft.fft2(x) 

        # 2. ATOM FIELD (Cloud)
        # x_fft shape: [batch, 1, 28, 28]
        # We unsqueeze to allow each atom to operate on the signal
        # Result: [batch, 10, 28, 28]
        cloud_fft = self.atoms(x_fft.unsqueeze(1))

        # 3. EMIT LIGHT (IFFT)
        # Transform each class-specific spectral cloud back to spatial domain
        cloud_spatial = torch.fft.ifft2(cloud_fft)
        
        # We care about the magnitude (energy) of the reconstructed signal
        # Truth is the constructive interference magnitude
        magnitude = torch.abs(cloud_spatial) 
        
        # 4. TRUTH EXTRACTION
        # Flatten and project to final class logits
        flat_magnitude = magnitude.view(magnitude.size(0), -1)
        return self.classifier(flat_magnitude)

# -------------------------------------------------------------------------
# Training and Testing Logic
# -------------------------------------------------------------------------
def train_spectral_mnist():
    # Hyperparameters
    batch_size = 64
    epochs = 2
    lr = 0.001

    # Data Loading
    transform = transforms.Compose([transforms.ToTensor()])
    train_set = datasets.MNIST(root='../data', train=True, download=True, transform=transform)
    test_set = datasets.MNIST(root='../data', train=False, download=True, transform=transform)
    
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SpectralNet().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    print(f"Training Spectral Computer on {device}...")

    for _ in range(5):
        for epoch in range(epochs):
            model.train()
            total_loss = 0
            for batch_idx, (data, target) in enumerate(train_loader):
                data, target = data.to(device), target.to(device)
                for i in range(10):
                    optimizer.zero_grad()
                    output = model(data)
                    loss = criterion(output, target)
                    loss.backward()
                    optimizer.step()
                    if i==0:    
                        total_loss += loss.item()
            
            print(f"Epoch {epoch+1}/{epochs} - Loss: {total_loss/len(train_loader):.4f}")

        # Evaluation
        model.eval()
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()

        accuracy = 100. * correct / len(test_loader.dataset)
        print(f"\nFinal Truth Extraction Accuracy: {accuracy:.2f}%")

if __name__ == "__main__":
    train_spectral_mnist()