"""
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()
        
        # Precompute index list per digit once (reused by every generate_batch call)
        labels = self.real_mnist[1]
        self.digit_indices = {d: np.where(labels == d)[0] for d in range(10)}
        
        # Per-digit statistics are expensive and unused by generation, so compute lazily.
        self._digit_statistics = None
    
    @property
    def digit_statistics(self) -> Dict[int, Dict]:
        """Computed on first access (Stationary component)."""
        if self._digit_statistics is None:
            self._digit_statistics = self._compute_digit_statistics()
        return self._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
        distribution = np.asarray(distribution, dtype=np.float64)
        distribution = distribution / distribution.sum()

        # Sample all labels, base images, and strengths up front (vectorized).
        labels = np.random.choice(10, size=n_samples, p=distribution).astype(np.int64)
        base_idx = np.empty(n_samples, dtype=np.int64)
        for d in range(10):
            mask = labels == d
            count = int(mask.sum())
            if count:
                base_idx[mask] = np.random.choice(self.digit_indices[d], size=count)
        strengths = np.random.uniform(strength_range[0], strength_range[1], size=n_samples)

        # Preallocate as float32 (half the memory of float64, and what torch wants).
        synthetic_images = np.empty((n_samples, 28, 28), dtype=np.float32)
        for i in range(n_samples):
            base_image = self.real_mnist[0][base_idx[i]]
            synthetic, _ = self.generate_single(base_image, int(labels[i]), float(strengths[i]))
            synthetic_images[i] = synthetic

        return synthetic_images, 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])
        
        n_total = int(self.config.n_synthetic)
        chunk_size = 100_000
        
        img_path = os.path.join(output_dir, 'synthetic_images.npy')
        lbl_path = os.path.join(output_dir, 'synthetic_labels.npy')
        
        # Disk-backed output. Each chunk is written straight to disk, so peak RAM
        # is one chunk (~78 MB at uint8), not the full dataset. Stored as uint8
        # (0-255): 8x smaller on disk than float64, ~negligible quantization loss.
        images_mm = np.lib.format.open_memmap(
            img_path, mode='w+', dtype=np.uint8, shape=(n_total, 28, 28))
        labels_mm = np.lib.format.open_memmap(
            lbl_path, mode='w+', dtype=np.uint8, shape=(n_total,))
        
        # Running stats (float64 accumulators) so we never need the whole array.
        pixel_sum = 0.0
        pixel_sqsum = 0.0
        class_counts = np.zeros(10, dtype=np.int64)
        
        print("Generating synthetic samples...")
        print("-" * 40)
        
        written = 0
        pbar = tqdm(total=n_total, desc="Generating")
        while written < n_total:
            cur = min(chunk_size, n_total - written)
            images, labels = self.generate_batch(
                n_samples=cur,
                distribution=real_dist.tolist(),
                strength_range=(0.3, 1.5)
            )
            
            # Quantize [0,1] float -> uint8 [0,255] and write to disk immediately.
            images_u8 = np.clip(images * 255.0 + 0.5, 0, 255).astype(np.uint8)
            images_mm[written:written + cur] = images_u8
            labels_mm[written:written + cur] = labels.astype(np.uint8)
            
            # Accumulate stats without holding the data (einsum avoids a big temp).
            pixel_sum += float(images.sum(dtype=np.float64))
            pixel_sqsum += float(np.einsum('ijk,ijk->', images, images, dtype=np.float64))
            class_counts += np.bincount(labels, minlength=10)
            
            written += cur
            pbar.update(cur)
            del images, labels, images_u8  # free chunk before next iteration
        pbar.close()
        
        images_mm.flush()
        labels_mm.flush()
        del images_mm, labels_mm  # close the memmaps
        
        print(f"\nGenerated {written:,} synthetic samples")
        print(f"Saved to {img_path} (uint8) and {lbl_path} (uint8)")
        
        # Finalize statistics from the running accumulators.
        n_pixels = written * 28 * 28
        pixel_mean = pixel_sum / n_pixels
        pixel_var = max(pixel_sqsum / n_pixels - pixel_mean ** 2, 0.0)
        stats = {
            'n_samples': written,
            'original_n_samples': len(self.real_mnist[0]),
            'expansion_ratio': written / len(self.real_mnist[0]),
            'pixel_mean': float(pixel_mean),
            'pixel_std': float(np.sqrt(pixel_var)),
            'dtype': 'uint8',
            'note': 'images are uint8 in [0,255]; divide by 255.0 to restore [0,1] floats',
            'class_distribution': {str(d): int(class_counts[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
        )