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

# ==========================================
# 0. Device Configuration
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ==========================================
# 1. The Continuous Frequency Architecture
# ==========================================

class ResonanceLayer(nn.Module):
    """
    A 'Resonance Chamber' that filters the input flow.
    Uses 3x3 kernels as 'Tuning Probes' to detect semantic frequencies.
    """
    def __init__(self, in_channels, out_channels):
        super(ResonanceLayer, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        # The flow is passed through the tuning probes
        res = self.conv(x)
        res = self.bn(res)
        # Tanh allows the state to oscillate between -1 (Tension) and 1 (Resolution)
        return torch.tanh(res)

class CCTCollapseLayer(nn.Module):
    """
    Implements Conditional Collapse Theory (CCT).
    Classification only occurs if Semantic Pressure (Amplitude) > threshold.
    """
    def __init__(self, in_features, num_classes, theta_coll=1.0):
        super(CCTCollapseLayer, self).__init__()
        self.linear = nn.Linear(in_features, num_classes)
        self.theta_coll = theta_coll

    def forward(self, x):
        logits = self.linear(x)
        # Semantic Pressure is the magnitude of the final state vector
        pressure = torch.norm(logits, p=2, dim=1, keepdim=True)
        return logits, pressure

class ContinuousFrequencyNet(nn.Module):
    def __init__(self):
        super(ContinuousFrequencyNet, self).__init__()
        # Layer 1: Low-frequency feature detection
        self.res1 = ResonanceLayer(1, 16)
        # Layer 2: High-frequency semantic refinement
        self.res2 = ResonanceLayer(16, 32)
        # Global Average Pooling: Integrating the manifold into a state vector
        self.pool = nn.AdaptiveAvgPool2d(1)
        # CCT Collapse Layer
        self.collapse = CCTCollapseLayer(32, 10)

    def forward(self, x):
        x = self.res1(x)
        x = self.res2(x)
        x = self.pool(x).view(x.size(0), -1)
        logits, pressure = self.collapse(x)
        return logits, pressure

# ==========================================
# 2. Training and Testing Pipeline
# ==========================================

def train():
    # Hyperparameters
    batch_size = 64
    epochs = 5
    lr = 0.001

    # Data Loading (MNIST)
    transform = transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # Send model to CUDA
    model = ContinuousFrequencyNet().to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    print("\nInitiating Resonance Training on CUDA...\n")

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            # Send data and target to CUDA
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            logits, pressure = model(data)
            loss = criterion(logits, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f}")

    # ==========================================
    # 3. Evaluation using CCT Logic
    # ==========================================
    model.eval()
    correct = 0
    unresolved = 0
    total = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            # Send data and target to CUDA
            data, target = data.to(device), target.to(device)
            
            logits, pressure = model(data)
            
            # CCT Condition: Does the system have enough 'Work/Energy' to collapse?
            # We use 2.0 as the 'Pressure' threshold for a hard collapse
            mask = (pressure < 2.0).squeeze() 
            
            pred = logits.argmax(dim=1, keepdim=True)
            
            is_correct = (pred == target).float().squeeze()
            
            # Safety for batch size 1
            if mask.dim() == 0: mask = mask.unsqueeze(0)
            
            # Correct = (Predicted == Target) AND (Not Masked/Unresolved)
            correct += (is_correct * (~mask)).sum().item()
            unresolved += mask.sum().item()
            total += target.size(0)

    print(f"\nFinal Results:")
    print(f"Accuracy (Collapsed): {100 * correct / total:.2f}%")
    print(f"Unresolved (Insufficient Work): {100 * unresolved / total:.2f}%")
    print(f"Total Potential Accuracy: {100 * (correct + (total - correct - unresolved)) / total:.2f}%")

if __name__ == "__main__":
    train()