"""
SUPERVISED SYNTHETIC MNIST EXPANSION
=====================================
Framework: ODE-CCT (Conditional Collapse Theory)

Stationary: Real MNIST (60K samples) - the fixed truth
Probability: Transformations applied - the expansion manifold
Work: Generation compute - energy spent to create new samples
Threshold: Intelligence threshold lowered by better manifold coverage

The AI "pays with work" to expand the stationary truth into
a larger probability space, reducing classification entropy.
"""

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

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

@dataclass
class Config:
    """Configuration aligned with CCT framework."""
    # Base data
    base_mnist_path: str = './data'
    n_synthetic: int = 1_000_000  # Target: 1M synthetic samples
    n_original: int = 60_000      # Original MNIST training set
    
    # Generation parameters
    transform_probability: float = 0.8  # How often to apply transform
    max_transform_strength: float = 1.5  # Maximum variance injection
    
    # Training parameters
    batch_size: int = 256
    epochs: int = 20
    learning_rate: float = 0.001
    
    # Device
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'

config = Config()


# =============================================================================
# TRANSFORMATION PIPELINE (Probability Component)
# =============================================================================

class SupervisedTransformPipeline:
    """
    Supervised transformations applied to real MNIST.
    
    Each transform is a "probability" injection that:
    - Preserves the digit label (supervised)
    - Adds variance to the manifold (expansion)
    - Requires "work" (compute) to execute
    
    CCT Alignment:
    - Stationary: The digit identity (unchanged)
    - Probability: The transformation parameters (varied)
    """
    
    def __init__(self, seed=42):
        np.random.seed(seed)
        torch.manual_seed(seed)
        self.rng = np.random.default_rng(seed)
        
        # Learned statistics from real MNIST (Stationary)
        self.pixel_mean = 0.1307
        self.pixel_std = 0.3081
        
    def random_rotation(self, image: np.ndarray, max_angle: float = 30) -> np.ndarray:
        """Rotate image with random angle. Preserves digit identity."""
        angle = self.rng.uniform(-max_angle, max_angle)
        rotated = ndimage.rotate(image, angle, reshape=False, order=1, mode='constant', cval=0)
        return np.clip(rotated, 0, 1)
    
    def random_scale(self, image: np.ndarray, scale_range: Tuple[float, float] = (0.8, 1.2)) -> np.ndarray:
        """Scale image. Preserves digit identity."""
        scale = self.rng.uniform(*scale_range)
        
        if scale == 1.0:
            return image
            
        h, w = image.shape
        scaled = ndimage.zoom(image, scale)
        
        # Crop or pad to original size
        if scale > 1:
            # Crop center
            start_h = (scaled.shape[0] - h) // 2
            start_w = (scaled.shape[1] - w) // 2
            result = scaled[start_h:start_h+h, start_w:start_w+w]
        else:
            # Pad with zeros
            pad_h = (h - scaled.shape[0]) // 2
            pad_w = (w - scaled.shape[1]) // 2
            result = np.zeros((h, w))
            result[pad_h:pad_h+scaled.shape[0], pad_w:pad_w+scaled.shape[1]] = scaled
            
        return np.clip(result, 0, 1)
    
    def random_translation(self, image: np.ndarray, max_shift: float = 4) -> np.ndarray:
        """Translate image within canvas. Preserves digit identity."""
        dx = self.rng.uniform(-max_shift, max_shift)
        dy = self.rng.uniform(-max_shift, max_shift)
        shifted = ndimage.shift(image, (dy, dx), mode='constant', cval=0)
        return shifted
    
    def random_shear(self, image: np.ndarray, max_shear: float = 0.3) -> np.ndarray:
        """Apply affine shear. Preserves digit identity."""
        shear = self.rng.uniform(-max_shear, max_shear)
        
        from scipy.ndimage import affine_transform
        matrix = np.array([[1, shear], [0, 1]])
        sheared = affine_transform(image, matrix, mode='constant', cval=0)
        
        return np.clip(sheared, 0, 1)
    
    def gaussian_noise(self, image: np.ndarray, sigma: float = 0.1) -> np.ndarray:
        """Add Gaussian noise. Slight identity perturbation."""
        noise = self.rng.normal(0, sigma, image.shape)
        noisy = image + noise
        return np.clip(noisy, 0, 1)
    
    def random_erasing(self, image: np.ndarray, p: float = 0.2, 
                       scale_range: Tuple[float, float] = (0.02, 0.1)) -> np.ndarray:
        """Random erasing augmentation. Preserves digit context."""
        if self.rng.random() > p:
            return image
            
        h, w = image.shape
        area = h * w
        
        target_area = self.rng.uniform(*scale_range) * area
        aspect_ratio = self.rng.uniform(0.3, 3)
        
        erase_h = int(np.sqrt(target_area * aspect_ratio))
        erase_w = int(np.sqrt(target_area / aspect_ratio))
        
        if erase_h >= h or erase_w >= w:
            return image
            
        top = self.rng.integers(0, h - erase_h)
        left = self.rng.integers(0, w - erase_w)
        
        result = image.copy()
        result[top:top+erase_h, left:left+erase_w] = 0
        
        return result
    
    def elastic_distortion(self, image: np.ndarray, alpha: float = 5, 
                           sigma: float = 1) -> np.ndarray:
        """Elastic deformation. Preserves local structure."""
        h, w = image.shape
        
        # Random displacement field
        dx = ndimage.gaussian_filter(self.rng.random((h, w)) * 2 - 1, sigma) * alpha
        dy = ndimage.gaussian_filter(self.rng.random((h, w)) * 2 - 1, sigma) * alpha
        
        # Apply displacement
        x, y = np.meshgrid(np.arange(w), np.arange(h))
        indices = (np.clip(y + dy, 0, h-1).astype(int),
                   np.clip(x + dx, 0, w-1).astype(int))
        
        distorted = image[indices]
        return np.clip(distorted, 0, 1)
    
    def morphological_thickness(self, image: np.ndarray, 
                                 thickness_range: Tuple[float, float] = (0.7, 1.3)) -> np.ndarray:
        """Vary stroke thickness using morphology."""
        binary = image > 0.5
        factor = self.rng.uniform(*thickness_range)
        
        iterations = max(1, int(abs(factor - 1) * 3))
        
        from scipy.ndimage import binary_dilation, binary_erosion
        struct = ndimage.generate_binary_structure(2, 1)
        
        if factor > 1:
            result = ndimage.binary_dilation(binary, struct, iterations=iterations)
        else:
            result = ndimage.binary_erosion(binary, struct, iterations=iterations)
            
        return result.astype(float)
    
    def brightness_adjustment(self, image: np.ndarray, 
                               factor_range: Tuple[float, float] = (0.7, 1.3)) -> np.ndarray:
        """Adjust brightness. Preserves structure."""
        factor = self.rng.uniform(*factor_range)
        adjusted = image * factor
        return np.clip(adjusted, 0, 1)
    
    def contrast_adjustment(self, image: np.ndarray, 
                            range_tuple: Tuple[float, float] = (0.8, 1.2)) -> np.ndarray:
        """Adjust contrast. Preserves structure."""
        factor = self.rng.uniform(*range_tuple)
        mean = np.mean(image)
        contrasted = (image - mean) * factor + mean
        return np.clip(contrasted, 0, 1)
    
    def apply_supervised_transforms(self, image: np.ndarray, 
                                     strength: float = 1.0,
                                     n_transforms: int = 3) -> np.ndarray:
        """
        Apply multiple transforms based on strength.
        
        Args:
            image: Input image (28x28, normalized [0,1])
            strength: Transform intensity (0 = no transform, 1 = full)
            n_transforms: Number of transforms to apply
            
        Returns:
            Transformed image with same label
        """
        result = image.copy()
        
        # Select random transforms
        transform_pool = [
            ('rotation', self.random_rotation, {'max_angle': 30 * strength}),
            ('scale', self.random_scale, {'scale_range': (1 - 0.2*strength, 1 + 0.2*strength)}),
            ('translation', self.random_translation, {'max_shift': 4 * strength}),
            ('shear', self.random_shear, {'max_shear': 0.3 * strength}),
            ('noise', self.gaussian_noise, {'sigma': 0.1 * strength}),
            ('erasing', self.random_erasing, {'p': 0.2 * strength}),
            ('elastic', self.elastic_distortion, {'alpha': 5 * strength}),
            ('thickness', self.morphological_thickness, {'thickness_range': (1 - 0.3*strength, 1 + 0.3*strength)}),
            ('brightness', self.brightness_adjustment, {'factor_range': (1 - 0.3*strength, 1 + 0.3*strength)}),
            ('contrast', self.contrast_adjustment, {'range_tuple': (1 - 0.2*strength, 1 + 0.2*strength)}),
        ]
        
        # Randomly select n_transforms
        selected = self.rng.choice(len(transform_pool), size=min(n_transforms, len(transform_pool)), replace=False)
        
        for idx in selected:
            name, func, kwargs = transform_pool[idx]
            try:
                result = func(result, **kwargs)
            except Exception as e:
                # If transform fails, skip it (don't break the pipeline)
                pass
                
        return result


# =============================================================================
# SUPERVISED SYNTHETIC MNIST GENERATOR
# =============================================================================

class SupervisedSyntheticMNIST:
    """
    Generates supervised synthetic MNIST from real MNIST.
    
    CCT Framework:
    - Stationary: Real MNIST samples (ground truth, labels preserved)
    - Probability: Transformation parameters (variance injected)
    - Work: Compute spent generating N variations
    - Threshold: Classification threshold lowered by expanded manifold
    """
    
    def __init__(self, config: Config):
        self.config = config
        self.transform_pipeline = SupervisedTransformPipeline()
        
        # Load real MNIST
        self.real_mnist = self._load_real_mnist()
        
        # Pre-compute per-digit statistics (Stationary)
        self.digit_statistics = self._compute_digit_statistics()
        
    def _load_real_mnist(self) -> Tuple[np.ndarray, np.ndarray]:
        """Load real MNIST training set from torchvision."""
        print("Loading real MNIST from torchvision...")
        
        transform = transforms.Compose([
            transforms.ToTensor()
        ])
        
        mnist_train = datasets.MNIST(
            root=self.config.base_mnist_path,
            train=True,
            download=True,
            transform=transform
        )
        
        # Extract data
        images = mnist_train.data.numpy().astype(np.float32) / 255.0
        labels = mnist_train.targets.numpy()
        
        print(f"  Loaded {len(images):,} real MNIST samples")
        print(f"  Shape: {images.shape}")
        print(f"  Labels distribution: {np.bincount(labels)}")
        
        return images, labels
    
    def _compute_digit_statistics(self) -> Dict[int, Dict]:
        """Compute statistics per digit class (Stationary component)."""
        stats = {}
        for digit in range(10):
            mask = self.real_mnist[1] == digit
            digit_images = self.real_mnist[0][mask]
            
            stats[digit] = {
                'count': len(digit_images),
                'mean_pixel': np.mean(digit_images),
                'std_pixel': np.std(digit_images),
                'mean_center_x': np.mean([np.where(img > 0.5)[1].mean() if np.any(img > 0.5) else 14 for img in digit_images]),
                'mean_center_y': np.mean([np.where(img > 0.5)[0].mean() if np.any(img > 0.5) else 14 for img in digit_images]),
                'density': np.mean([np.mean(img > 0.5) for img in digit_images]),
            }
            
        return stats
    
    def generate_single(self, base_image: np.ndarray, base_label: int,
                        strength: float = 1.0) -> Tuple[np.ndarray, int]:
        """
        Generate a single synthetic sample from a real sample.
        
        Args:
            base_image: Real MNIST image (28x28)
            base_label: The digit label (preserved)
            strength: Transform intensity (0-1.5)
            
        Returns:
            (synthetic_image, label) - supervised pair
        """
        synthetic = self.transform_pipeline.apply_supervised_transforms(
            base_image, 
            strength=strength,
            n_transforms=int(3 + strength * 2)  # More transforms for higher strength
        )
        
        return synthetic, base_label
    
    def generate_batch(self, n_samples: int, 
                       distribution: Optional[List[float]] = None,
                       strength_range: Tuple[float, float] = (0.3, 1.5)) -> Tuple[np.ndarray, np.ndarray]:
        """
        Generate a batch of synthetic samples.
        
        Args:
            n_samples: Number of samples to generate
            distribution: Class distribution (default: uniform)
            strength_range: Range of transform strengths
            
        Returns:
            (synthetic_images, labels)
        """
        if distribution is None:
            distribution = [0.1] * 10
            
        synthetic_images = []
        synthetic_labels = []
        
        # Create index mapping for sampling
        digit_indices = {d: np.where(self.real_mnist[1] == d)[0] for d in range(10)}
        
        for _ in range(n_samples):
            # Sample class based on distribution
            label = np.random.choice(10, p=distribution)
            
            # Sample random base image from that class
            base_idx = np.random.choice(digit_indices[label])
            base_image = self.real_mnist[0][base_idx]
            
            # Sample transform strength
            strength = np.random.uniform(*strength_range)
            
            # Generate synthetic sample
            synthetic, _ = self.generate_single(base_image, label, strength)
            
            synthetic_images.append(synthetic)
            synthetic_labels.append(label)
            
        return np.array(synthetic_images), np.array(synthetic_labels)
    
    def generate_1m_dataset(self, output_dir: str = './synthetic_mnist_supervised') -> Dict:
        """
        Main generation pipeline: Create 1M supervised synthetic samples.
        
        CCT Alignment:
        - Each generated sample = a "question" probing the manifold
        - Work = compute spent on transformations
        - Threshold = lowered by expanded coverage
        """
        print("=" * 60)
        print("SUPERVISED SYNTHETIC MNIST GENERATION")
        print("=" * 60)
        print(f"Target: {self.config.n_synthetic:,} synthetic samples")
        print(f"Base: {len(self.real_mnist[0]):,} real MNIST samples")
        print()
        
        os.makedirs(output_dir, exist_ok=True)
        
        # Class distribution matching real MNIST
        real_dist = np.bincount(self.real_mnist[1]) / len(self.real_mnist[1])
        
        # Generate in chunks
        chunk_size = 100_000
        n_chunks = self.config.n_synthetic // chunk_size
        
        all_images = []
        all_labels = []
        
        print("Generating synthetic samples...")
        print("-" * 40)
        
        for chunk_idx in tqdm(range(n_chunks), desc="Generating"):
            images, labels = self.generate_batch(
                n_samples=chunk_size,
                distribution=real_dist.tolist(),
                strength_range=(0.3, 1.5)
            )
            all_images.append(images)
            all_labels.append(labels)
            
        # Combine
        synthetic_images = np.concatenate(all_images, axis=0)
        synthetic_labels = np.concatenate(all_labels, axis=0)
        
        print(f"\nGenerated {len(synthetic_images):,} synthetic samples")
        
        # Save
        print("\nSaving dataset...")
        np.save(os.path.join(output_dir, 'synthetic_images.npy'), synthetic_images)
        np.save(os.path.join(output_dir, 'synthetic_labels.npy'), synthetic_labels)
        
        # Compute and save statistics
        stats = {
            'n_samples': len(synthetic_images),
            'original_n_samples': len(self.real_mnist[0]),
            'expansion_ratio': len(synthetic_images) / len(self.real_mnist[0]),
            'pixel_mean': float(np.mean(synthetic_images)),
            'pixel_std': float(np.std(synthetic_images)),
            'class_distribution': {str(d): int(np.sum(synthetic_labels == d)) for d in range(10)},
            'generation_method': 'SupervisedSyntheticMNIST v1.0',
            'framework': 'ODE-CCT Probability Expansion'
        }
        
        with open(os.path.join(output_dir, 'stats.json'), 'w') as f:
            json.dump(stats, f, indent=2)
            
        print(f"✓ Saved to {output_dir}/")
        
        return stats


# =============================================================================
# MLP CLASSIFIER FOR COMPARISON
# =============================================================================

class MLPClassifier(nn.Module):
    """
    Simple MLP for MNIST classification.
    Used to compare real vs synthetic vs combined datasets.
    """
    
    def __init__(self, input_size=784, hidden_sizes=[512, 256, 128], num_classes=10):
        super().__init__()
        
        layers = []
        prev_size = input_size
        
        for hidden_size in hidden_sizes:
            layers.extend([
                nn.Linear(prev_size, hidden_size),
                nn.ReLU(),
                nn.BatchNorm1d(hidden_size),
                nn.Dropout(0.2)
            ])
            prev_size = hidden_size
            
        layers.append(nn.Linear(prev_size, num_classes))
        
        self.network = nn.Sequential(*layers)
        
    def forward(self, x):
        return self.network(x)


def train_model(model, train_loader, val_loader, epochs, lr, device):
    """Train MLP and return training history."""
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, factor=0.5)
    
    history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    
    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0
        train_correct = 0
        train_total = 0
        
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            train_total += batch_y.size(0)
            train_correct += predicted.eq(batch_y).sum().item()
            
        train_acc = 100. * train_correct / train_total
        history['train_loss'].append(train_loss / len(train_loader))
        history['train_acc'].append(train_acc)
        
        # Validation
        model.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                outputs = model(batch_x)
                loss = criterion(outputs, batch_y)
                
                val_loss += loss.item()
                _, predicted = outputs.max(1)
                val_total += batch_y.size(0)
                val_correct += predicted.eq(batch_y).sum().item()
                
        val_acc = 100. * val_correct / val_total
        history['val_loss'].append(val_loss / len(val_loader))
        history['val_acc'].append(val_acc)
        
        scheduler.step(val_loss)
        
        print(f"  Epoch {epoch+1:2d}: Train Acc={train_acc:.2f}%, Val Acc={val_acc:.2f}%")
        
    return history


def run_comparison_experiments(n_synthetic=500_000, n_test=10000):
    """
    Run experiments comparing:
    1. Real MNIST only (60K)
    2. Synthetic only (500K)
    3. Real + Synthetic combined
    
    CCT Hypothesis: Combined should have lowest entropy = best generalization
    """
    print("=" * 60)
    print("COMPARISON EXPERIMENTS: Real vs Synthetic vs Combined")
    print("=" * 60)
    
    device = config.device
    print(f"Device: {device}")
    
    # Load real MNIST
    transform = transforms.Compose([transforms.ToTensor()])
    mnist_train = datasets.MNIST('../data', train=True, download=True, transform=transform)
    mnist_test = datasets.MNIST('../data', train=False, download=True, transform=transform)
    
    # Test set
    test_images = mnist_test.data.view(-1, 784).float() / 255.0
    test_labels = mnist_test.targets
    
    test_loader = DataLoader(TensorDataset(test_images, test_labels), batch_size=256, shuffle=False)
    
    # Subsample test for faster evaluation
    test_subset = DataLoader(TensorDataset(
        test_images[:n_test], test_labels[:n_test]
    ), batch_size=256, shuffle=False)
    
    results = {}
    
    # =========================================================================
    # EXPERIMENT 1: Real MNIST Only
    # =========================================================================
    print("\n" + "-" * 40)
    print("Experiment 1: Real MNIST Only (60K)")
    print("-" * 40)
    
    train_images = mnist_train.data.view(-1, 784).float() / 255.0
    train_labels = mnist_train.targets
    
    train_loader = DataLoader(
        TensorDataset(train_images, train_labels),
        batch_size=config.batch_size, shuffle=True
    )
    
    model1 = MLPClassifier()
    history1 = train_model(model1, train_loader, test_subset, 
                           epochs=config.epochs, lr=config.learning_rate, device=device)
    results['real_only'] = history1
    
    # =========================================================================
    # EXPERIMENT 2: Synthetic Only
    # =========================================================================
    print("\n" + "-" * 40)
    print(f"Experiment 2: Synthetic Only ({n_synthetic:,})")
    print("-" * 40)
    
    generator = SupervisedSyntheticMNIST(config)
    synth_images, synth_labels = generator.generate_batch(n_synthetic)
    
    synth_images_flat = synth_images.reshape(n_synthetic, 784)
    
    synth_loader = DataLoader(
        TensorDataset(torch.FloatTensor(synth_images_flat), torch.LongTensor(synth_labels)),
        batch_size=config.batch_size, shuffle=True
    )
    
    model2 = MLPClassifier()
    history2 = train_model(model2, synth_loader, test_subset,
                           epochs=config.epochs, lr=config.learning_rate, device=device)
    results['synthetic_only'] = history2
    
    # =========================================================================
    # EXPERIMENT 3: Combined (Real + Synthetic)
    # =========================================================================
    print("\n" + "-" * 40)
    print(f"Experiment 3: Real + Synthetic Combined ({60000 + n_synthetic:,})")
    print("-" * 40)
    
    # Combine datasets
    combined_images = torch.cat([
        train_images,
        torch.FloatTensor(synth_images_flat)
    ], dim=0)
    combined_labels = torch.cat([
        train_labels,
        torch.LongTensor(synth_labels)
    ])
    
    combined_loader = DataLoader(
        TensorDataset(combined_images, combined_labels),
        batch_size=config.batch_size, shuffle=True
    )
    
    model3 = MLPClassifier()
    history3 = train_model(model3, combined_loader, test_subset,
                           epochs=config.epochs, lr=config.learning_rate, device=device)
    results['combined'] = history3
    
    # =========================================================================
    # PLOT RESULTS
    # =========================================================================
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    
    # Accuracy plot
    ax = axes[0]
    ax.plot(history1['val_acc'], label='Real Only (60K)', marker='o')
    ax.plot(history2['val_acc'], label=f'Synthetic Only ({n_synthetic//1000}K)', marker='s')
    ax.plot(history3['val_acc'], label='Combined', marker='^')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Validation Accuracy (%)')
    ax.set_title('MNIST Classification: ODE-CCT Expansion Comparison')
    ax.legend()
    ax.grid(True)
    
    # Loss plot
    ax = axes[1]
    ax.plot(history1['val_loss'], label='Real Only', marker='o')
    ax.plot(history2['val_loss'], label='Synthetic Only', marker='s')
    ax.plot(history3['val_loss'], label='Combined', marker='^')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Validation Loss')
    ax.set_title('Loss Convergence')
    ax.legend()
    ax.grid(True)
    
    plt.tight_layout()
    plt.savefig('./comparison_results.png', dpi=150)
    print("\n✓ Saved comparison plot to ./comparison_results.png")
    
    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"Real Only (60K):      Best Val Acc = {max(history1['val_acc']):.2f}%")
    print(f"Synthetic Only ({n_synthetic//1000}K): Best Val Acc = {max(history2['val_acc']):.2f}%")
    print(f"Combined:             Best Val Acc = {max(history3['val_acc']):.2f}%")
    
    return results


# =============================================================================
# VISUALIZATION
# =============================================================================

def visualize_samples(real_images, synthetic_images, labels, n_samples=50):
    """Visualize real vs synthetic samples."""
    
    fig, axes = plt.subplots(4, 25, figsize=(20, 6))
    
    for i in range(25):
        # Real samples (row 0-1)
        idx = np.random.randint(len(real_images))
        axes[0, i].imshow(real_images[idx], cmap='gray')
        axes[0, i].axis('off')
        axes[1, i].imshow(real_images[idx], cmap='gray')
        axes[1, i].axis('off')
        
        # Synthetic samples (row 2-3)
        synth_mask = labels == np.bincount(labels).argmax()  # Use most common label
        synth_idx = np.random.randint(len(synthetic_images))
        axes[2, i].imshow(synthetic_images[synth_idx], cmap='gray')
        axes[2, i].axis('off')
        axes[3, i].imshow(synthetic_images[synth_idx], cmap='gray')
        axes[3, i].axis('off')
    
    axes[0, 0].set_ylabel('Real A', fontsize=10)
    axes[1, 0].set_ylabel('Real B', fontsize=10)
    axes[2, 0].set_ylabel('Synthetic A', fontsize=10)
    axes[3, 0].set_ylabel('Synthetic B', fontsize=10)
    
    plt.suptitle('Real MNIST (top) vs Supervised Synthetic (bottom)', fontsize=14)
    plt.tight_layout()
    plt.savefig('./real_vs_synthetic.png', dpi=150)
    print("✓ Saved comparison visualization to ./real_vs_synthetic.png")


# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='Supervised Synthetic MNIST Generation')
    parser.add_argument('--mode', choices=['generate', 'train', 'compare', 'all'], default='all',
                        help='Mode: generate dataset, train, compare, or all')
    parser.add_argument('--n_synthetic', type=int, default=500000,
                        help='Number of synthetic samples to generate')
    parser.add_argument('--n_epochs', type=int, default=15,
                        help='Number of training epochs')
    
    args = parser.parse_args()
    
    # Update config
    config.n_synthetic = args.n_synthetic
    config.epochs = args.n_epochs
    
    if args.mode in ['generate', 'all']:
        # Generate 1M synthetic samples
        generator = SupervisedSyntheticMNIST(config)
        stats = generator.generate_1m_dataset()
        
        # Visualize comparison
        real_images = generator.real_mnist[0][:100]
        synth_images, synth_labels = generator.generate_batch(100)
        visualize_samples(real_images, synth_images, synth_labels)
        
    if args.mode in ['compare', 'all']:
        # Run comparison experiments
        results = run_comparison_experiments(
            n_synthetic=args.n_synthetic,
            n_test=10000
        )