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

# Check for GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ==========================================
# 1. THE COMPLEX-EULER NEURAL-ODE COMPONENT
# ==========================================
class SuperpositionObserverODE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=128, output_dim=10, steps=5, dt=0.2):
        super(SuperpositionObserverODE, self).__init__()
        self.steps = steps
        self.dt = dt
        self.hidden_dim = hidden_dim
        
        # Complex Encoder: Maps real inputs to Amplitude and Phase components (Beam Splitter)
        self.amplitude_encoder = nn.Linear(input_dim, hidden_dim)
        self.phase_encoder = nn.Linear(input_dim, hidden_dim)
        
        # ODE Derivative Function f(z, t): parameterized as complex-valued layers split into real parts
        # This governs the "quiver plot" trajectories of the architecture
        self.real_dynamics = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.imag_dynamics = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        
        # Observer Decoder: Collapses the complex superposition state back to a real classification probability
        self.decoder = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Flatten MNIST images (Batch, 1, 28, 28) -> (Batch, 784)
        x = x.view(x.size(0), -1)
        
        # 1. Complexification (Beam Splitter Transformation)
        # We generate amplitude r (gated via sigmoid to keep bounds stable) and phase theta
        r = torch.sigmoid(self.amplitude_encoder(x))
        theta = torch.tanh(self.phase_encoder(x)) * 3.14159  # Phase bounded between [-pi, pi]
        
        # Initial condition C in the Complex Plane: z = r * (cos(theta) + i*sin(theta))
        real_z = r * torch.cos(theta)
        imag_z = r * torch.sin(theta)
        
        # 2. ODE Trajectory Integration via Complex Euler step (dz/dt = f(z, t))
        for _ in range(self.steps):
            # Compute the complex derivative elements
            # Real and imaginary components cross-couple to model phase rotations naturally
            d_real = self.real_dynamics(real_z) - self.imag_dynamics(imag_z)
            d_imag = self.imag_dynamics(real_z) + self.real_dynamics(imag_z)
            
            # Euler step mutation update
            real_z = real_z + self.dt * d_real
            imag_z = imag_z + self.dt * d_imag
            
        # 3. Superposition Observation / Conditional Collapse
        # Calculate the quantum-inspired probability amplitude |z|^2
        probabilities = (real_z ** 2) + (imag_z ** 2)
        
        # Map collapsed state to final digit classification
        out = self.decoder(probabilities)
        return out

# ==========================================
# 2. DATA PREPARATION
# ==========================================
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='../data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='../data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# ==========================================
# 3. INITIALIZATION & TRAINING LOOP
# ==========================================
model = SuperpositionObserverODE(input_dim=784, hidden_dim=128, output_dim=10, steps=6, dt=0.15).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=0.002, weight_decay=1e-4)

print("Starting Complex Evolution Training...")
epochs = 5

for epoch in range(1, epochs + 1):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 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()
        
        running_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        if batch_idx % 200 == 0:
            print(f"Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] "
                  f"Loss: {loss.item():.4f} | Batch Accuracy: {100. * correct / total:.2f}%")
            
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    print(f"--- Epoch {epoch} Complete | Average Loss: {epoch_loss:.4f} | Training Accuracy: {epoch_acc:.2f}% ---\n")

# ==========================================
# 4. TESTING & EVALUATION
# ==========================================
model.eval()
test_loss = 0
correct = 0

with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        test_loss += criterion(output, target).item()
        _, predicted = output.max(1)
        correct += predicted.eq(target).sum().item()

test_loss /= len(test_loader)
test_accuracy = 100. * correct / len(test_loader.dataset)

print("==========================================")
print(f"FINAL SYSTEM EVALUATION")
print(f"Average Test Loss: {test_loss:.4f}")
print(f"Phase-Aligned Test Accuracy: {test_accuracy:.2f}%")
print("==========================================")