"""
FUNCTION-ORIENTED MLP (F-MLP)
=============================
Instead of learning static patterns, this architecture learns:
1. Base manifold (what makes a digit a digit)
2. Variation flow (how transformations map through space)
3. Classification as trajectory destination

This is the ODE-CCT model architecture for MNIST.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, TensorDataset
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from scipy import ndimage
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, Optional
from dataclasses import dataclass

# =============================================================================
# CONFIGURATION
# =============================================================================

@dataclass
class FunctionMLPConfig:
    # Model dimensions
    input_size: int = 784
    variation_embedding_dim: int = 64    # How many variation axes we model
    manifold_dim: int = 128              # Base digit manifold dimension
    output_classes: int = 10
    
    # Training
    batch_size: int = 512
    epochs: int = 10
    lr: float = 0.002
    
    # Synthetic generation (supervised)
    n_synthetic_per_class: int = 100_000  # 1M total
    transform_strength_range: Tuple[float, float] = (0.2, 1.5)
    
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'


# =============================================================================
# VARIATION EXTRACTION LAYER
# =============================================================================

class VariationExtractor(nn.Module):
    """
    FUNCTION 1: Intra-Class Variation Handler
    
    Extracts the "variation parameters" from an input.
    Instead of classifying directly, it asks:
    "What transformation was applied to produce this?"
    
    This is the inverse of the generation process.
    """
    
    def __init__(self, input_dim=784, variation_dim=64):
        super().__init__()
        
        # Encoder that maps image → variation embedding
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.LayerNorm(512),
            nn.GELU(),
            nn.Dropout(0.1),
            
            nn.Linear(512, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(0.1),
            
            nn.Linear(256, variation_dim * 2),  # mu and log_var for VAE-style
            nn.Tanh()  # Bounded variation parameters
        )
        
    def forward(self, x):
        params = self.encoder(x)
        mu, log_var = params.chunk(2, dim=-1)
        # Return variation parameters (can be used to reconstruct)
        return mu  # Deterministic for inference


class ManifoldProjection(nn.Module):
    """
    FUNCTION 2: Inter-Class Separator
    
    Projects variation-extracted input onto the digit manifold.
    This learns the "stationary" structure — what makes each digit invariant.
    
    Think of it as: "Given this variation-extracted input, where am I in digit-space?"
    """
    
    def __init__(self, variation_dim=64, manifold_dim=128, num_classes=10):
        super().__init__()
        
        # Project to manifold with class-specific centers
        self.manifold = nn.Sequential(
            nn.Linear(variation_dim, manifold_dim),
            nn.LayerNorm(manifold_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            
            nn.Linear(manifold_dim, manifold_dim),
            nn.LayerNorm(manifold_dim),
            nn.GELU(),
            nn.Dropout(0.1),
        )
        
        # Class centroids (learnable)
        self.class_centroids = nn.Parameter(torch.randn(num_classes, manifold_dim))
        
    def forward(self, variation_embedding):
        manifold_state = self.manifold(variation_embedding)
        
        # Compute distance to each class centroid
        distances = torch.cdist(manifold_state, self.class_centroids)
        
        # Soft assignment (like attention)
        attention = torch.softmax(-distances, dim=-1)
        
        return manifold_state, attention, distances


class EdgeClassifier(nn.Module):
    """
    FUNCTION 3: Cross-Class Edge Handler
    
    Handles the confusion zone between similar digits.
    Uses the manifold topology to make decisions at boundaries.
    """
    
    def __init__(self, manifold_dim=128, num_classes=10):
        super().__init__()
        
        # Edge detection in manifold space
        self.edge_detector = nn.Sequential(
            nn.Linear(manifold_dim * 3, 64),  # manifold_state + top_2_centroids
            nn.LayerNorm(64),
            nn.GELU(),
            nn.Linear(64, 32),
            nn.GELU(),
            nn.Linear(32, num_classes)
        )
        
    def forward(self, manifold_state, top_centroids, class_logits):
        """
        Args:
            manifold_state: Current position in manifold
            top_centroids: Top 2 class centroids (for edge detection)
            class_logits: Initial class prediction
        """
        # Create edge feature: [current_pos, closest_centroid, second_closest_centroid]
        edge_features = torch.cat([
            manifold_state,
            top_centroids[:, 0],
            top_centroids[:, 1]
        ], dim=-1)
        
        # Edge correction logits
        edge_correction = self.edge_detector(edge_features)
        
        # Combine with initial prediction
        final_logits = class_logits + 0.3 * edge_correction
        
        return final_logits


# =============================================================================
# FUNCTION MLP (Complete Architecture)
# =============================================================================

class FunctionMLP(nn.Module):
    """
    The complete Function-Oriented MLP.
    
    Instead of learning static weights for 1M samples,
    it learns 3 functions that can handle infinite variations.
    
    CCT Alignment:
    - VariationExtractor: "Probability" — extracts what changed
    - ManifoldProjection: "Stationary" — finds the invariant core
    - EdgeClassifier: "Cross-boundary" — handles transition zones
    """
    
    def __init__(self, config: FunctionMLPConfig):
        super().__init__()
        
        self.config = config
        
        # Function 1: Extract variation parameters
        self.variation_extractor = VariationExtractor(
            input_dim=config.input_size,
            variation_dim=config.variation_embedding_dim
        )
        
        # Function 2: Project to manifold (class separation)
        self.manifold_projection = ManifoldProjection(
            variation_dim=config.variation_embedding_dim,
            manifold_dim=config.manifold_dim,
            num_classes=config.output_classes
        )
        
        # Function 3: Handle edges
        self.edge_classifier = EdgeClassifier(
            manifold_dim=config.manifold_dim,
            num_classes=config.output_classes
        )
        
        # Direct classification path (fallback)
        self.direct_classifier = nn.Sequential(
            nn.Linear(config.input_size, 256),
            nn.GELU(),
            nn.Linear(256, config.output_classes)
        )
        
    def forward(self, x):
        # Flatten if needed
        if len(x.shape) > 2:
            x = x.view(x.size(0), -1)
            
        # === FUNCTION 1: Variation Extraction ===
        variation_embedding = self.variation_extractor(x)
        
        # === FUNCTION 2: Manifold Projection ===
        manifold_state, attention, distances = self.manifold_projection(variation_embedding)
        
        # Initial class prediction from manifold distance
        initial_logits = -distances  # Closer = higher score
        
        # === FUNCTION 3: Edge Handling ===
        # Get top 2 centroids
        top_centroids = self.class_centroids[torch.argsort(distances, dim=-1)[:, :2]]
        final_logits = self.edge_classifier(manifold_state, top_centroids, initial_logits)
        
        # Fallback: direct classification blend
        direct_logits = self.direct_classifier(x)
        final_logits = 0.7 * final_logits + 0.3 * direct_logits
        
        return final_logits
    
    @property
    def class_centroids(self):
        return self.manifold_projection.class_centroids


# =============================================================================
# SMALLER ARCHITECTURE VARIANTS
# =============================================================================

class MiniFunctionMLP(nn.Module):
    """
    Ultra-minimal version for testing the hypothesis.
    If the function hypothesis is correct, we need far fewer parameters.
    """
    
    def __init__(self, input_dim=784, variation_dim=16, manifold_dim=32, num_classes=10):
        super().__init__()
        
        # Shared function: variation + manifold in one
        self.shared = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.GELU(),
            nn.Linear(128, variation_dim),
            nn.Tanh()
        )
        
        # Class centroids (10 × 32 = 320 params)
        self.centroids = nn.Parameter(torch.randn(num_classes, manifold_dim))
        
        # Variation to manifold projection
        self.projection = nn.Linear(variation_dim, manifold_dim)
        
        # Tiny edge handler
        self.edge = nn.Linear(manifold_dim * 3, num_classes)
        
        # Direct path
        self.direct = nn.Linear(input_dim, num_classes)
        
    def forward(self, x):
        if len(x.shape) > 2:
            x = x.view(x.size(0), -1)
            
        # Variation embedding
        v = self.shared(x)
        
        # Manifold position
        m = self.projection(v)
        
        # Class distances
        dist = torch.cdist(m.unsqueeze(1), self.centroids.unsqueeze(0)).squeeze(1)
        logits = -dist
        
        # Edge handling
        top2_idx = torch.argsort(dist, dim=-1)[:, :2]
        top2 = self.centroids[top2_idx]
        
        edge_feat = torch.cat([m, top2[:, 0], top2[:, 1]], dim=-1)
        edge_correction = self.edge(edge_feat)
        
        # Combine
        direct = self.direct(x)
        final = 0.6 * (logits + 0.2 * edge_correction) + 0.4 * direct
        
        return final


class TinyFunctionMLP(nn.Module):
    """
    Minimal possible function model.
    Just variation extraction + centroid distance.
    Total params: ~50K (vs ~400K for standard MLP)
    """
    
    def __init__(self, num_classes=10):
        super().__init__()
        
        self.encoder = nn.Sequential(
            nn.Linear(784, 64),
            nn.GELU(),
            nn.Linear(64, 32)
        )
        
        self.centroids = nn.Parameter(torch.randn(num_classes, 32))
        
        self.proj = nn.Linear(32, 32)
        
        self.direct = nn.Sequential(
            nn.Linear(784, 64),
            nn.GELU(),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, x):
        if len(x.shape) > 2:
            x = x.view(x.size(0), -1)
            
        # Variation embedding
        v = self.encoder(x)
        
        # Manifold projection
        m = torch.tanh(self.proj(v))
        
        # Centroid distance
        dist = torch.cdist(m.unsqueeze(1), self.centroids.unsqueeze(0)).squeeze(1)
        
        # Direct path
        d = self.direct(x)
        
        return 0.5 * (-dist) + 0.5 * d


# =============================================================================
# SYNTHETIC DATA GENERATOR
# =============================================================================

class SyntheticMNISTGenerator:
    """
    Generate supervised synthetic data for training.
    Each sample is a "question" that probes a variation axis.
    """
    
    def __init__(self, device='cpu'):
        self.device = device
        
        # Load real MNIST
        self.mnist = datasets.MNIST('./data', train=True, download=True, 
                                     transform=transforms.ToTensor())
        
    def generate_batch(self, n_samples, transform_strength=(0.2, 1.5)):
        """
        Generate synthetic samples with random transformations.
        
        Each generated sample tests:
        - "Given this variation of digit X, can the model handle it?"
        """
        images = []
        labels = []
        
        for _ in range(n_samples):
            idx = np.random.randint(len(self.mnist))
            img, label = self.mnist[idx]
            
            img = img.squeeze().numpy()
            
            # Random transformations
            strength = np.random.uniform(*transform_strength)
            
            # Rotation
            angle = np.random.uniform(-30 * strength, 30 * strength)
            img = ndimage.rotate(img, angle, reshape=False, order=1, mode='constant', cval=0)
            
            # Noise
            if np.random.random() < 0.5 * strength:
                noise = np.random.normal(0, 0.1 * strength, img.shape)
                img = np.clip(img + noise, 0, 1)
            
            # Scale
            if np.random.random() < 0.5:
                scale = np.random.uniform(1 - 0.15 * strength, 1 + 0.15 * strength)
                img = ndimage.zoom(img, scale)
                if img.shape[0] > 28:
                    start = (img.shape[0] - 28) // 2
                    img = img[start:start+28, start:start+28]
                elif img.shape[0] < 28:
                    pad = (28 - img.shape[0]) // 2
                    padded = np.zeros((28, 28))
                    padded[pad:pad+img.shape[0], pad:pad+img.shape[1]] = img
                    img = padded
            
            # Translation
            dx = np.random.uniform(-3 * strength, 3 * strength)
            dy = np.random.uniform(-3 * strength, 3 * strength)
            img = ndimage.shift(img, (dy, dx), mode='constant', cval=0)
            
            images.append(img.flatten())
            labels.append(label)
            
        return torch.FloatTensor(np.array(images)), torch.LongTensor(np.array(labels))


# =============================================================================
# TRAINING LOOP
# =============================================================================

def train_function_mlp(model, train_loader, val_loader, epochs, lr, device):
    """Train the Function MLP and track parameters vs accuracy."""
    
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    history = {
        'epoch': [],
        'train_acc': [],
        'val_acc': [],
        'params': sum(p.numel() for p in model.parameters())
    }
    
    for epoch in range(epochs):
        # Train
        model.train()
        correct = 0
        total = 0
        
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            for k in range(2):         
                optimizer.zero_grad()
                out = model(x)
                loss = criterion(out, y)
                loss.backward()
                optimizer.step()
                if k==0:
                    _, pred = out.max(1)
                    correct += pred.eq(y).sum().item()
                    total += y.size(0)
            
        train_acc = 100. * correct / total
        
        # Validate
        model.eval()
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for x, y in val_loader:
                x, y = x.to(device), y.to(device)
                out = model(x)
                _, pred = out.max(1)
                val_correct += pred.eq(y).sum().item()
                val_total += y.size(0)
                
        val_acc = 100. * val_correct / val_total
        
        scheduler.step()
        
        history['epoch'].append(epoch)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)
        
        print(f"Epoch {epoch+1:2d}: Train={train_acc:.2f}%, Val={val_acc:.2f}%")
        
    return history


# =============================================================================
# COMPARISON EXPERIMENT
# =============================================================================

def run_function_comparison(n_train=500000, n_val=10000):
    """
    Compare:
    1. Standard MLP (400K params)
    2. Function MLP (150K params)
    3. Mini Function MLP (50K params)
    4. Tiny Function MLP (10K params)
    
    Hypothesis: Smaller function models should match/better standard MLP
    when trained on synthetic data that explores the variation manifold.
    """
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Device: {device}")
    
    # Load test set
    mnist_test = datasets.MNIST('./data', train=False, download=True,
                                transform=transforms.ToTensor())
    test_x = mnist_test.data.view(-1, 784).float() / 255.0
    test_y = mnist_test.targets
    test_loader = DataLoader(TensorDataset(test_x, test_y), batch_size=512, shuffle=False)
    
    # Subset for validation
    val_loader = DataLoader(
        TensorDataset(test_x[:n_val], test_y[:n_val]),
        batch_size=512, shuffle=False
    )
    
    # Load original MNIST training data
    print("\nLoading original MNIST training data...")
    mnist_train = datasets.MNIST('./data', train=True, download=True,
                                 transform=transforms.ToTensor())
    train_x = mnist_train.data.view(-1, 784).float() / 255.0
    train_y = mnist_train.targets
    
    print(f"  Training samples: {len(train_x):,}")
    print(f"  Test samples: {len(test_x):,}")
    
    train_loader = DataLoader(
        TensorDataset(train_x, train_y),
        batch_size=512, shuffle=True
    )
    
    results = {}
    
    # =========================================================================
    # Function MLP (150K params)
    # =========================================================================
    print("\n" + "="*50)
    print("Function MLP (150K params)")
    print("="*50)
    
    func_mlp = FunctionMLP(FunctionMLPConfig(
        variation_embedding_dim=64*2,
        manifold_dim=128*2
    ))
    
    n_params = sum(p.numel() for p in func_mlp.parameters())
    print(f"Parameters: {n_params:,}")
    
    history = train_function_mlp(func_mlp, train_loader, val_loader,
                                  epochs=100, lr=0.01, device=device)
    
    # Final evaluation on full test set
    print("\nEvaluating on full test set...")
    func_mlp.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for x, y in test_loader:
            x, y = x.to(device), y.to(device)
            out = func_mlp(x)
            pred = out.argmax(dim=1)
            correct += (pred == y).sum().item()
            total += y.size(0)
    
    test_acc = 100 * correct / total
    print(f"Final Test Accuracy: {test_acc:.2f}%")
    
    results['function_mlp'] = {**history, 'params': n_params, 'test_acc': test_acc}
    
    # =========================================================================
    # PLOT RESULTS
    # =========================================================================
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # Accuracy comparison
    ax = axes[0]
    for name, res in results.items():
        ax.plot(res['val_acc'], label=f"{name} ({res['params']:,} params)", marker='o')
    
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Validation Accuracy (%)')
    ax.set_title('Function Models Performance')
    ax.legend()
    ax.grid(True)
    
    # Parameters vs Final Accuracy
    ax = axes[1]
    names = list(results.keys())
    params = [results[n]['params'] for n in names]
    final_acc = [results[n]['test_acc'] for n in names]
    
    ax.scatter(params, final_acc, s=150)
    for i, name in enumerate(names):
        ax.annotate(name, (params[i], final_acc[i]), 
                   textcoords="offset points", xytext=(10,10))
    
    ax.set_xscale('log')
    ax.set_xlabel('Number of Parameters')
    ax.set_ylabel('Final Test Accuracy (%)')
    ax.set_title('Parameters vs Performance')
    ax.grid(True)
    
    plt.tight_layout()
    plt.savefig('./function_mlp_comparison.png', dpi=150)
    print("\n✓ Saved comparison plot")
    
    # Summary
    print("\n" + "="*50)
    print("SUMMARY")
    print("="*50)
    for name, res in results.items():
        print(f"{name:25s}: {res['params']:6,} params, "
              f"Best Val Acc = {max(res['val_acc']):.2f}%, "
              f"Final Test Acc = {res['test_acc']:.2f}%")
    
    return results


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

if __name__ == '__main__':
    results = run_function_comparison(n_train=5000, n_val=10000)