"""
90/10 Hybrid Classifier on MNIST
================================
90%: Holomorphic neural network (general solution)
10%: Sparse anchor memory (particular corrections)
"""

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

# ============================================================
# HYBRID COMPONENTS
# ============================================================

class HolomorphicNetwork(nn.Module):
    """
    90%: A neural network that approximates the decision boundary.
    Uses smooth activation functions to create a holomorphic-like mapping.
    """
    def __init__(self, input_dim, hidden_dim=256, output_dim=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SiLU(),  # Smooth activation (holomorphic approximation)
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    
    def forward(self, x):
        return self.net(x)


class AnchorMemory(nn.Module):
    """
    10%: Sparse memory of training examples for high-confidence correction.
    """
    def __init__(self, n_anchors=100, n_classes=10, feature_dim=784):
        super().__init__()
        self.n_anchors = n_anchors
        self.n_classes = n_classes
        
        # Anchor prototypes (learnable)
        self.anchors = nn.Parameter(
            torch.randn(n_anchors, feature_dim) * 0.1
        )
        # Anchor class logits (one per anchor)
        self.anchor_logits = nn.Parameter(
            torch.randn(n_anchors, n_classes) * 0.1
        )
    
    def forward(self, x):
        """
        Compute anchor-based prediction.
        x: (batch, feature_dim)
        Returns: (batch, n_classes)
        """
        # Compute distances to all anchors
        # (batch, 1, feature_dim) - (1, n_anchors, feature_dim)
        dists = torch.cdist(x.unsqueeze(1), self.anchors.unsqueeze(0))
        dists = dists.squeeze(1)  # (batch, n_anchors)
        
        # Softmax weights (closer anchors have higher weight)
        weights = torch.softmax(-dists / 0.5, dim=1)  # (batch, n_anchors)
        
        # Weighted combination of anchor logits
        anchor_pred = torch.matmul(weights, self.anchor_logits)
        
        return anchor_pred


class HybridClassifier(nn.Module):
    """
    90/10 Hybrid Classifier
    =======================
    90%: Holomorphic network (generalization)
    10%: Anchor memory (specific corrections)
    """
    def __init__(self, input_dim=784, n_anchors=100, n_classes=10, 
                 function_ratio=0.9):
        super().__init__()
        self.function_ratio = function_ratio
        self.anchor_ratio = 1 - function_ratio
        
        # 90% component: Holomorphic network
        self.function_net = HolomorphicNetwork(input_dim, 256, n_classes)
        
        # 10% component: Anchor memory
        self.anchor_mem = AnchorMemory(n_anchors, n_classes, input_dim)
    
    def forward(self, x):
        """
        Combined prediction: 90% function + 10% anchors.
        """
        func_pred = self.function_net(x)
        anchor_pred = self.anchor_mem(x)
        
        return self.function_ratio * func_pred + self.anchor_ratio * anchor_pred


# ============================================================
# TRAINING
# ============================================================

def train_epoch(model, loader, optimizer, criterion, device):
    """Train for one epoch."""
    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)
        data = data.view(data.size(0), -1)  # Flatten images
        
        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)
    
    avg_loss = total_loss / len(loader)
    accuracy = correct / total
    return avg_loss, accuracy


def evaluate(model, loader, criterion, device):
    """Evaluate on test set."""
    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)
            data = data.view(data.size(0), -1)
            
            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)
    
    avg_loss = total_loss / len(loader)
    accuracy = correct / total
    return avg_loss, accuracy


# ============================================================
# MAIN
# ============================================================

def main():
    # Setup
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    
    # Data
    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=128, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
    
    # Model
    model = HybridClassifier(input_dim=784, n_anchors=100, n_classes=10, function_ratio=0.9)
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    print("=" * 60)
    print("90/10 Hybrid Classifier on MNIST")
    print("=" * 60)
    print(f"Training samples: {len(train_dataset)}")
    print(f"Test samples: {len(test_dataset)}")
    print(f"Function ratio: 90%")
    print(f"Anchor ratio: 10%")
    print(f"Anchor points: 100")
    print("=" * 60)
    
    # Training loop
    n_epochs = 10
    best_test_acc = 0
    
    for epoch in range(n_epochs):
        t_start = time.time()
        
        # Train
        train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device)
        
        # Evaluate
        test_loss, test_acc = evaluate(model, test_loader, criterion, device)
        
        elapsed = time.time() - t_start
        
        print(f"\nEpoch {epoch+1}/{n_epochs} ({elapsed:.2f}s)")
        print(f"  Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f}")
        print(f"  Test Loss:  {test_loss:.4f}, Acc: {test_acc:.4f}")
        
        if test_acc > best_test_acc:
            best_test_acc = test_acc
            torch.save(model.state_dict(), 'hybrid_mnist.pth')
            print(f"  ** New best test accuracy: {best_test_acc:.4f}")
    
    # Load best model
    model.load_state_dict(torch.load('hybrid_mnist.pth'))
    
    # Final evaluation
    print("\n" + "=" * 60)
    print("FINAL RESULTS")
    print("=" * 60)
    final_loss, final_acc = evaluate(model, test_loader, criterion, device)
    print(f"  Test Loss:  {final_loss:.4f}")
    print(f"  Test Acc:   {final_acc:.4f}")
    print(f"  Best Acc:   {best_test_acc:.4f}")
    print("=" * 60)
    
    # Show what was learned
    print("\n--- Learned Structure ---")
    print(f"  Function network: {sum(p.numel() for p in model.function_net.parameters()):,} params")
    print(f"  Anchor memory: {sum(p.numel() for p in model.anchor_mem.parameters()):,} params")
    print(f"  Total params: {sum(p.numel() for p in model.parameters()):,}")
    print(f"  Function ratio: {model.function_ratio:.0%}")
    print(f"  Anchor ratio: {model.anchor_ratio:.0%}")
    print("=" * 60)
    print("Hybrid representation: 90% function + 10% values")
    print("=" * 60)


if __name__ == "__main__":
    main()