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, num_atoms, in_channels, img_size):
        super(SpectralAtom, self).__init__()
        # We represent the complex weights as two real tensors (real and imaginary)
        # The filter size matches the input signal size (32x32 for CIFAR-10)
        # Shape: [num_atoms, in_channels, height, width]
        self.weight_real = nn.Parameter(torch.randn(num_atoms, in_channels, img_size, img_size) * 0.01)
        self.weight_imag = nn.Parameter(torch.randn(num_atoms, in_channels, img_size, img_size) * 0.01)

    def forward(self, x_fft):
        # x_fft is expected to be a complex tensor of shape [batch, 1, C, H, W]
        w_complex = torch.complex(self.weight_real, self.weight_imag) # Shape: [num_atoms, C, H, W]
        
        # Pointwise multiplication in frequency domain
        # Broadcasting rules align from the right: 
        # [batch, 1, C, H, W] * [num_atoms, C, H, W] -> [batch, num_atoms, C, H, W]
        return x_fft * w_complex

# -------------------------------------------------------------------------
# The SpectralComputingNetwork
# -------------------------------------------------------------------------
class SpectralNet(nn.Module):
    def __init__(self, num_classes=10, in_channels=3, img_size=32):
        super(SpectralNet, self).__init__()
        # We create 10 'Atoms', one for each digit/object class.
        self.atoms = SpectralAtom(num_classes, in_channels, img_size)
        
        # Final projection to flatten the resulting "Cloud" into a class score
        # Shape after flattening: [batch, num_classes * in_channels * img_size * img_size]
        flat_features = num_classes * in_channels * img_size * img_size
        self.classifier = nn.Linear(flat_features, num_classes)

    def forward(self, x):
        # 1. PROMPT (Nucleus) -> FFT
        # Convert image to complex and perform 2D FFT over the last two dimensions (H, W)
        x = x.float()
        x_fft = torch.fft.fft2(x) 
        
        # 2. ATOM FIELD (Cloud)
        # x_fft shape: [batch, 3, 32, 32]
        # We unsqueeze to allow each atom to operate on the signal -> [batch, 1, 3, 32, 32]
        # Resulting cloud_fft shape: [batch, 10, 3, 32, 32]
        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
        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_cifar10():
    # Hyperparameters
    batch_size = 64
    epochs = 10 # Increased epochs slightly as CIFAR-10 is more complex than MNIST
    lr = 0.001
    
    # Data Loading
    # Note: For RGB images, it's often beneficial to add normalization, 
    # but we keep ToTensor() here to mirror the original script's simplicity.
    transform = transforms.Compose([transforms.ToTensor()])
    
    train_set = datasets.CIFAR10(root='../data', train=True, download=True, transform=transform)
    test_set = datasets.CIFAR10(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")
    
    # Initialize model with CIFAR-10 specific dimensions
    model = SpectralNet(num_classes=10, in_channels=3, img_size=32).to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    print(f"Training Spectral Computer on {device}...")
    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)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            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_cifar10()