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

# ═══════════════════════════════════════════════════════════════
# FIELD LINEAR ALGEBRA - Core Operations
# ═══════════════════════════════════════════════════════════════

class FieldLinearAlgebra:
    """
    Field representation where features m of matrix (n, m) are sources of energy.
    Enables operations between matrices with different feature dimensions.
    """
    
    @staticmethod
    def field_mul(A: torch.Tensor, B: torch.Tensor, W_A: torch.Tensor, W_B: torch.Tensor) -> torch.Tensor:
        """
        Field multiplication: combines two matrices into unified field representation.
        
        Y = A @ W_A + B @ W_B
        
        Args:
            A: (n, m1) first feature matrix
            B: (n, m2) second feature matrix  
            W_A: (m1, k) learned projection for A
            W_B: (m2, k) learned projection for B
            
        Returns:
            Y: (n, k) combined field representation
        """
        return A @ W_A + B @ W_B
    
    @staticmethod
    def field_add(A: torch.Tensor, B: torch.Tensor, W_A: torch.Tensor, W_B: torch.Tensor) -> torch.Tensor:
        """
        Field addition: combines information preserving both sources.
        """
        return A @ W_A + B @ W_B
    
    @staticmethod
    def field_scale(A: torch.Tensor, W: torch.Tensor, alpha: float) -> torch.Tensor:
        """Scale field strength."""
        return alpha * (A @ W)

# ═══════════════════════════════════════════════════════════════
# MODEL ARCHITECTURES
# ═══════════════════════════════════════════════════════════════

class StandardMLP(nn.Module):
    """Standard MLP for comparison baseline."""
    def __init__(self, input_dim=784, hidden1=256, hidden2=128, num_classes=10):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden1)
        self.fc2 = nn.Linear(hidden1, hidden2)
        self.fc3 = nn.Linear(hidden2, num_classes)
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x


class FieldMLP(nn.Module):
    """
    MLP using Field Linear Algebra.
    Splits features into two sources, projects to unified field.
    """
    def __init__(self, input_dim=784, split_dim=392, field_dim=256, hidden=128, num_classes=10):
        super().__init__()
        # Feature source A: first half of features
        # Feature source B: second half of features
        
        self.split_dim = split_dim
        
        # Learned projection matrices (m1 + m2 -> k)
        # W_A: (392, field_dim), W_B: (392, field_dim)
        # Combined parameter count: 392 * 256 * 2 = 200,704
        
        self.W_A = nn.Parameter(torch.randn(split_dim, field_dim) * 0.01)
        self.W_B = nn.Parameter(torch.randn(split_dim, field_dim) * 0.01)
        
        self.field_proj = nn.Linear(field_dim, hidden)
        self.fc_out = nn.Linear(hidden, num_classes)
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        x = x.view(x.size(0), -1)  # (batch, 784)
        
        # Split into two feature sources
        A = x[:, :self.split_dim]      # (batch, 392)
        B = x[:, self.split_dim:]      # (batch, 392)
        
        # Field multiplication: combine sources in learned field space
        # Y = A @ W_A + B @ W_B  -> (batch, 256)
        y = A @ self.W_A + B @ self.W_B
        
        y = F.relu(self.field_proj(y))
        y = self.dropout(y)
        y = self.fc_out(y)
        return y
    
    def get_field_dim(self):
        """Return the combined field dimension."""
        return self.W_A.shape[1]


class FieldMLP_v2(nn.Module):
    """
    Field MLP with residual connections and normalization.
    """
    def __init__(self, input_dim=784, split_dim=392, field_dim=512, hidden=256, num_classes=10):
        super().__init__()
        self.split_dim = split_dim
        
        # Learned projections for field multiplication
        self.W_A = nn.Parameter(torch.randn(split_dim, field_dim) * 0.02)
        self.W_B = nn.Parameter(torch.randn(split_dim, field_dim) * 0.02)
        
        # Optional: learned scaling for each source
        self.alpha = nn.Parameter(torch.ones(1))
        self.beta = nn.Parameter(torch.ones(1))
        
        self.norm = nn.LayerNorm(field_dim)
        self.fc1 = nn.Linear(field_dim, hidden)
        self.fc2 = nn.Linear(hidden, num_classes)
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        
        A = x[:, :self.split_dim]
        B = x[:, self.split_dim:]
        
        # Field multiplication with learnable scaling
        y = self.alpha * (A @ self.W_A) + self.beta * (B @ self.W_B)
        y = self.norm(y)
        
        y = F.gelu(self.fc1(y))
        y = self.dropout(y)
        y = self.fc2(y)
        return y


class FieldLinearBlock(nn.Module):
    """Single field linear transformation block."""
    def __init__(self, dim_a, dim_b, field_dim):
        super().__init__()
        self.W_A = nn.Parameter(torch.randn(dim_a, field_dim) * 0.02)
        self.W_B = nn.Parameter(torch.randn(dim_b, field_dim) * 0.02)
        self.norm = nn.LayerNorm(field_dim)
        
    def forward(self, A, B):
        return self.norm(A @ self.W_A + B @ self.W_B)


# ═══════════════════════════════════════════════════════════════
# TRAINING & EVALUATION
# ═══════════════════════════════════════════════════════════════

def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(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()
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
    
    return total_loss / len(loader), 100. * correct / total


def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion(output, target)
            
            total_loss += loss.item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
    
    return total_loss / len(loader), 100. * correct / total


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# ═══════════════════════════════════════════════════════════════
# MAIN EXPERIMENT
# ═══════════════════════════════════════════════════════════════

def main():
    # Hyperparameters
    EPOCHS = 15
    BATCH_SIZE = 128
    LR = 0.001
    FIELD_DIM = 256  # k = m1 + m2 = 392 + 392 = 784, but we'll use learned 256
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Load 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)
    
    print(f"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}")
    
    results = {}
    
    # ═══════════════════════════════════════════════════════════
    # Model 1: Standard MLP (Baseline)
    # ═══════════════════════════════════════════════════════════
    print("\n" + "="*60)
    print("Training Standard MLP")
    print("="*60)
    
    standard_model = StandardMLP().to(device)
    print(f"Parameters: {count_parameters(standard_model):,}")
    
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(standard_model.parameters(), lr=LR)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
    
    start_time = time.time()
    for epoch in range(EPOCHS):
        train_loss, train_acc = train_epoch(standard_model, train_loader, criterion, optimizer, device)
        scheduler.step()
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1:2d}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%")
    
    train_time = time.time() - start_time
    test_loss, test_acc = evaluate(standard_model, test_loader, criterion, device)
    
    results['Standard MLP'] = {
        'params': count_parameters(standard_model),
        'train_acc': train_acc,
        'test_acc': test_acc,
        'time': train_time
    }
    print(f"Test Accuracy: {test_acc:.2f}%, Time: {train_time:.2f}s")
    
    # ═══════════════════════════════════════════════════════════
    # Model 2: Field MLP (m1+m2 -> 256 learned projection)
    # ═══════════════════════════════════════════════════════════
    print("\n" + "="*60)
    print("Training Field MLP (256 field dim)")
    print("="*60)
    
    field_model = FieldMLP().to(device)
    print(f"Parameters: {count_parameters(field_model):,}")
    print(f"Field projection: W_A({field_model.W_A.shape}) + W_B({field_model.W_B.shape})")
    
    optimizer = torch.optim.Adam(field_model.parameters(), lr=LR)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
    
    start_time = time.time()
    for epoch in range(EPOCHS):
        train_loss, train_acc = train_epoch(field_model, train_loader, criterion, optimizer, device)
        scheduler.step()
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1:2d}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%")
    
    train_time = time.time() - start_time
    test_loss, test_acc = evaluate(field_model, test_loader, criterion, device)
    
    results['Field MLP (k=256)'] = {
        'params': count_parameters(field_model),
        'train_acc': train_acc,
        'test_acc': test_acc,
        'time': train_time
    }
    print(f"Test Accuracy: {test_acc:.2f}%, Time: {train_time:.2f}s")
    
    # ═══════════════════════════════════════════════════════════
    # Model 3: Field MLP v2 (larger field, residual)
    # ═══════════════════════════════════════════════════════════
    print("\n" + "="*60)
    print("Training Field MLP v2 (512 field dim)")
    print("="*60)
    
    field_model_v2 = FieldMLP_v2().to(device)
    print(f"Parameters: {count_parameters(field_model_v2):,}")
    
    optimizer = torch.optim.Adam(field_model_v2.parameters(), lr=LR)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
    
    start_time = time.time()
    for epoch in range(EPOCHS):
        train_loss, train_acc = train_epoch(field_model_v2, train_loader, criterion, optimizer, device)
        scheduler.step()
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1:2d}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%")
    
    train_time = time.time() - start_time
    test_loss, test_acc = evaluate(field_model_v2, test_loader, criterion, device)
    
    results['Field MLP v2 (k=512)'] = {
        'params': count_parameters(field_model_v2),
        'train_acc': train_acc,
        'test_acc': test_acc,
        'time': train_time
    }
    print(f"Test Accuracy: {test_acc:.2f}%, Time: {train_time:.2f}s")
    
    # ═══════════════════════════════════════════════════════════
    # RESULTS SUMMARY
    # ═══════════════════════════════════════════════════════════
    print("\n" + "="*60)
    print("RESULTS SUMMARY")
    print("="*60)
    print(f"{'Model':<25} {'Params':<12} {'Test Acc':<10} {'Time':<8}")
    print("-"*60)
    for name, res in results.items():
        print(f"{name:<25} {res['params']:<12,} {res['test_acc']:<10.2f} {res['time']:<8.2f}s")
    
    # ═══════════════════════════════════════════════════════════
    # ANALYSIS: Field properties
    # ═══════════════════════════════════════════════════════════
    print("\n" + "="*60)
    print("FIELD ANALYSIS")
    print("="*60)
    
    # Inspect learned field projections
    print("\nField projection weight statistics:")
    print(f"  W_A: mean={field_model.W_A.mean():.4f}, std={field_model.W_A.std():.4f}")
    print(f"  W_B: mean={field_model.W_B.mean():.4f}, std={field_model.W_B.std():.4f}")
    
    # SVD of learned projections (information distribution)
    UA, SA, _ = torch.svd(field_model.W_A)
    UB, SB, _ = torch.svd(field_model.W_B)
    
    print(f"\nSingular value distribution (top 5):")
    print(f"  W_A: {SA[:5].cpu().numpy()}")
    print(f"  W_B: {SB[:5].cpu().numpy()}")
    
    # Variance explained by top-k singular values
    cumvar_A = torch.cumsum(SA**2, dim=0) / torch.sum(SA**2)
    cumvar_B = torch.cumsum(SB**2, dim=0) / torch.sum(SB**2)
    
    print(f"\nCumulative variance explained:")
    for k in [32, 64, 128, 256]:
        print(f"  Top {k}: W_A={cumvar_A[k-1]:.3f}, W_B={cumvar_B[k-1]:.3f}")
    
    return results

if __name__ == "__main__":
    results = main()