"""
Flux ResNet-18 for CIFAR-10
=============================
ResNet-18 architecture with Flux Algebra tracking.
Trains on CIFAR-10 with uncertainty quantification and entropy collapse.

Based on: The Algebra of Flux (Conditional Collapse Theory)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from typing import Dict, List, Tuple, Optional
import time
from tqdm import tqdm
import json

from flux_tensor import FluxTensor


# ==================== FLUX BATCH NORM ====================

class FluxBatchNorm2d(nn.Module):
    """
    Batch Normalization with Flux tracking.
    Tracks uncertainty in running statistics.
    """

    def __init__(self, num_features: int, momentum: float = 0.1):
        super().__init__()
        self.bn = nn.BatchNorm2d(num_features, momentum=momentum)

        # Flux state for gamma (weight) and beta (bias)
        self.gamma_flux = None
        self.beta_flux = None

    def initialize_flux(self):
        """Initialize flux tensors after parameters are created."""
        if self.gamma_flux is None:
            self.gamma_flux = FluxTensor(
                self.bn.weight.data.clone(),
                s=torch.ones_like(self.bn.weight.data) * 0.1,
                t=torch.zeros_like(self.bn.weight.data)
            )
            self.beta_flux = FluxTensor(
                self.bn.bias.data.clone(),
                s=torch.ones_like(self.bn.bias.data) * 0.1,
                t=torch.zeros_like(self.bn.bias.data)
            )

    def to(self, *args, **kwargs):
        """Override to() to move flux tensors to target device."""
        super().to(*args, **kwargs)
        if self.gamma_flux is not None:
            self.gamma_flux = self.gamma_flux.to(self.bn.weight.device)
        if self.beta_flux is not None:
            self.beta_flux = self.beta_flux.to(self.bn.weight.device)
        return self

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self.initialize_flux()
        # Ensure flux tensors are on the same device as input
        if self.gamma_flux.v.device != x.device:
            self.gamma_flux = self.gamma_flux.to(x.device)
            self.beta_flux = self.beta_flux.to(x.device)
        # Update BN parameters with current flux values
        self.bn.weight.data = self.gamma_flux.v
        self.bn.bias.data = self.beta_flux.v
        return self.bn(x)
    
    def flux_update(self, loss: torch.Tensor, lr: float, work: float):
        self.initialize_flux()
        with torch.no_grad():
            # Gamma update
            if self.bn.weight.grad is not None:
                delta = -lr * self.bn.weight.grad
                self.gamma_flux = FluxTensor(
                    self.bn.weight + delta,
                    self.gamma_flux.s + torch.abs(delta),
                    self.gamma_flux.t + delta
                ).collapse(work)
                self.bn.weight.data = self.gamma_flux.v
            
            # Beta update
            if self.bn.bias.grad is not None:
                delta = -lr * self.bn.bias.grad
                self.beta_flux = FluxTensor(
                    self.bn.bias + delta,
                    self.beta_flux.s + torch.abs(delta),
                    self.beta_flux.t + delta
                ).collapse(work)
                self.bn.bias.data = self.beta_flux.v
    
    def entropy_stats(self) -> dict:
        self.initialize_flux()
        return {
            'gamma_entropy': self.gamma_flux.s.mean().item(),
            'beta_entropy': self.beta_flux.s.mean().item()
        }


# ==================== FLUX CONV2D ====================

class FluxConv2d(nn.Module):
    """
    Convolutional layer with Flux tracking.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        padding: int = 0,
        stride: int = 1,
        bias: bool = False
    ):
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels, out_channels, kernel_size,
            padding=padding, stride=stride, bias=bias
        )
        self.weight_flux = None
        self.bias_flux = None

    def initialize_flux(self, init_entropy: float = 0.1):
        """Initialize flux tensors."""
        if self.weight_flux is None:
            self.weight_flux = FluxTensor(
                self.conv.weight.data.clone(),
                s=torch.ones_like(self.conv.weight.data) * init_entropy,
                t=torch.zeros_like(self.conv.weight.data)
            )
            if self.conv.bias is not None:
                self.bias_flux = FluxTensor(
                    self.conv.bias.data.clone(),
                    s=torch.ones_like(self.conv.bias.data) * init_entropy,
                    t=torch.zeros_like(self.conv.bias.data)
                )

    def to(self, *args, **kwargs):
        """Override to() to move flux tensors to target device."""
        super().to(*args, **kwargs)
        if self.weight_flux is not None:
            self.weight_flux = self.weight_flux.to(self.conv.weight.device)
        if self.bias_flux is not None:
            self.bias_flux = self.bias_flux.to(self.conv.weight.device)
        return self

    def _sync_flux_to_params(self):
        """Copy flux values back to underlying parameters."""
        if self.weight_flux is not None:
            self.conv.weight.data = self.weight_flux.v
        if self.bias_flux is not None:
            self.conv.bias.data = self.bias_flux.v

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self.initialize_flux()
        # Ensure flux tensors are on the same device as input
        if self.weight_flux.v.device != x.device:
            self.weight_flux = self.weight_flux.to(x.device)
            if self.bias_flux is not None:
                self.bias_flux = self.bias_flux.to(x.device)
        self._sync_flux_to_params()
        return self.conv(x)
    
    def flux_update(self, loss: torch.Tensor, lr: float, work: float):
        self.initialize_flux()
        with torch.no_grad():
            # Weight update
            if self.conv.weight.grad is not None:
                delta = -lr * self.conv.weight.grad
                self.weight_flux = FluxTensor(
                    self.conv.weight + delta,
                    self.weight_flux.s + torch.abs(delta),
                    self.weight_flux.t + delta
                ).collapse(work)
                self.conv.weight.data = self.weight_flux.v
            
            # Bias update
            if self.conv.bias is not None and self.conv.bias.grad is not None:
                delta = -lr * self.conv.bias.grad
                self.bias_flux = FluxTensor(
                    self.conv.bias + delta,
                    self.bias_flux.s + torch.abs(delta),
                    self.bias_flux.t + delta
                ).collapse(work)
                self.conv.bias.data = self.bias_flux.v
    
    def collapse(self, work: float):
        self.initialize_flux()
        self.weight_flux = self.weight_flux.collapse(work)
        if self.conv.bias is not None:
            self.bias_flux = self.bias_flux.collapse(work)
    
    def entropy_stats(self) -> dict:
        self.initialize_flux()
        stats = {'weight_entropy': self.weight_flux.s.mean().item()}
        if self.conv.bias is not None:
            stats['bias_entropy'] = self.bias_flux.s.mean().item()
        return stats


# ==================== FLUX BASIC BLOCK ====================

class FluxBasicBlock(nn.Module):
    """
    ResNet Basic Block (for ResNet-18/34) with Flux tracking.
    
    Architecture:
        x → Conv → BN → ReLU → Conv → BN → +x → ReLU
    """
    
    expansion = 1
    
    def __init__(
        self, 
        in_channels: int, 
        out_channels: int, 
        stride: int = 1,
        downsample: Optional[nn.Module] = None
    ):
        super().__init__()
        self.downsample = downsample
        
        # Conv layers with flux tracking
        self.conv1 = FluxConv2d(in_channels, out_channels, 3, padding=1, stride=stride)
        self.bn1 = FluxBatchNorm2d(out_channels)
        self.conv2 = FluxConv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = FluxBatchNorm2d(out_channels)
        
        self.relu = nn.ReLU(inplace=True)
        self.stride = stride
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        
        out += identity
        out = self.relu(out)
        
        return out
    
    def flux_update(self, loss: torch.Tensor, lr: float, work: float):
        """Update all flux states in this block."""
        self.conv1.flux_update(loss, lr, work)
        self.bn1.flux_update(loss, lr, work)
        self.conv2.flux_update(loss, lr, work)
        self.bn2.flux_update(loss, lr, work)
        
        if self.downsample is not None and isinstance(self.downsample, nn.Sequential):
            if isinstance(self.downsample[0], FluxConv2d):
                self.downsample[0].flux_update(loss, lr, work)
            if isinstance(self.downsample[1], FluxBatchNorm2d):
                self.downsample[1].flux_update(loss, lr, work)
    
    def collapse(self, work: float):
        self.conv1.collapse(work)
        self.conv2.collapse(work)
        if self.downsample is not None and isinstance(self.downsample, nn.Sequential):
            if isinstance(self.downsample[0], FluxConv2d):
                self.downsample[0].collapse(work)
    
    def get_all_entropy(self) -> List[float]:
        """Collect all entropy values in this block."""
        entropy_list = []
        for layer in [self.conv1, self.bn1, self.conv2, self.bn2]:
            stats = layer.entropy_stats()
            entropy_list.extend([v for v in stats.values()])
        return entropy_list


# ==================== FLUX RESNET ====================

class FluxResNet(nn.Module):
    """
    ResNet with Flux Algebra tracking.
    
    Supports ResNet-18, 34, 50, 101, 152 architectures.
    """
    
    def __init__(
        self, 
        block: nn.Module, 
        layers: List[int], 
        num_classes: int = 10,
        init_entropy: float = 0.1
    ):
        super().__init__()
        self.in_channels = 64
        self.init_entropy = init_entropy
        
        # Initial conv (7x7 for ImageNet, adapted to 3x3 for CIFAR)
        self.conv1 = FluxConv2d(3, 64, 3, padding=1, stride=1)
        self.bn1 = FluxBatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.Identity()  # Remove for CIFAR
        
        # ResNet layers
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        
        # Final classifier
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = FluxConv2d(512 * block.expansion, num_classes, 1)
        
        # Initialize all flux states
        self._initialize_all_flux()
    
    def _initialize_all_flux(self):
        """Initialize flux tensors for all layers."""
        self.conv1.initialize_flux(self.init_entropy)
        self.bn1.initialize_flux()
        self.bn1.gamma_flux = FluxTensor(
            self.bn1.bn.weight.data.clone(),
            s=torch.ones_like(self.bn1.bn.weight.data) * self.init_entropy,
            t=torch.zeros_like(self.bn1.bn.weight.data)
        )
        self.bn1.beta_flux = FluxTensor(
            self.bn1.bn.bias.data.clone(),
            s=torch.ones_like(self.bn1.bn.bias.data) * self.init_entropy,
            t=torch.zeros_like(self.bn1.bn.bias.data)
        )
        self.fc.initialize_flux(self.init_entropy)
    
    def _make_layer(
        self, 
        block: nn.Module, 
        channels: int, 
        blocks: int, 
        stride: int = 1
    ) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.in_channels != channels * block.expansion:
            downsample = nn.Sequential(
                FluxConv2d(self.in_channels, channels * block.expansion, 1, stride=stride),
                FluxBatchNorm2d(channels * block.expansion)
            )
        
        layers = []
        layers.append(block(self.in_channels, channels, stride, downsample))
        self.in_channels = channels * block.expansion
        
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, channels))
        
        return nn.Sequential(*layers)

    def to(self, *args, **kwargs):
        """Override to() to move all flux tensors to target device."""
        super().to(*args, **kwargs)
        self.conv1.to(*args, **kwargs)
        self.bn1.to(*args, **kwargs)
        self.fc.to(*args, **kwargs)
        for layer_group in [self.layer1, self.layer2, self.layer3, self.layer4]:
            for block in layer_group:
                block.conv1.to(*args, **kwargs)
                block.bn1.to(*args, **kwargs)
                block.conv2.to(*args, **kwargs)
                block.bn2.to(*args, **kwargs)
                if block.downsample is not None:
                    for module in block.downsample:
                        if hasattr(module, 'to') and isinstance(module, nn.Module):
                            module.to(*args, **kwargs)
        return self

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.avgpool(x)
        x = self.fc(x)
        x = x.view(x.size(0), -1)
        
        return x
    
    def flux_update_all(self, loss: torch.Tensor, lr: float, work: float):
        """Update flux states for all layers."""
        self.conv1.flux_update(loss, lr, work)
        self.bn1.flux_update(loss, lr, work)
        
        for layer_group in [self.layer1, self.layer2, self.layer3, self.layer4]:
            for block in layer_group:
                block.flux_update(loss, lr, work)
        
        self.fc.flux_update(loss, lr, work)
    
    def collapse_all(self, work: float):
        """Apply collapse to all flux states."""
        self.conv1.collapse(work)
        self.fc.collapse(work)
        
        for layer_group in [self.layer1, self.layer2, self.layer3, self.layer4]:
            for block in layer_group:
                block.collapse(work)
    
    def get_entropy_report(self) -> dict:
        """Get comprehensive entropy report."""
        report = {
            'conv1': self.conv1.entropy_stats(),
            'bn1': self.bn1.entropy_stats(),
            'fc': self.fc.entropy_stats(),
        }
        
        for i, layer_group in enumerate([self.layer1, self.layer2, self.layer3, self.layer4]):
            layer_entropies = []
            for block in layer_group:
                layer_entropies.extend(block.get_all_entropy())
            report[f'layer{i+1}'] = {
                'mean': np.mean(layer_entropies),
                'max': np.max(layer_entropies),
                'min': np.min(layer_entropies)
            }
        
        return report
    
    def total_entropy(self) -> float:
        """Calculate total entropy across all parameters."""
        total = 0.0
        
        # Conv1 and BN1
        for stats in [self.conv1.entropy_stats(), self.bn1.entropy_stats()]:
            total += sum(stats.values())
        
        # FC
        for stats in self.fc.entropy_stats().values():
            total += stats
        
        # All layers
        for layer_group in [self.layer1, self.layer2, self.layer3, self.layer4]:
            for block in layer_group:
                for ent in block.get_all_entropy():
                    total += ent
        
        return total


def flux_resnet18(num_classes: int = 10, init_entropy: float = 0.1) -> FluxResNet:
    """Create ResNet-18 with Flux tracking."""
    return FluxResNet(FluxBasicBlock, [2, 2, 2, 2], num_classes, init_entropy)


# ==================== FLUX TRAINER ====================

class FluxTrainer:
    """
    Training loop with Flux Algebra integration.
    
    Tracks entropy, applies collapse, and monitors learning dynamics.
    """
    
    def __init__(
        self,
        model: FluxResNet,
        device: torch.device,
        lr: float = 0.1,
        momentum: float = 0.9,
        weight_decay: float = 5e-4,
        work_schedule: Optional[str] = 'linear',
        collapse_every_n: int = 5,
        collapse_work: float = 0.5
    ):
        self.model = model
        self.device = device
        
        # SGD optimizer (standard for ResNet)
        self.optimizer = optim.SGD(
            [p for p in model.parameters() if p.requires_grad],
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay
        )
        
        # Learning rate scheduler (cosine annealing)
        self.scheduler = None
        
        # Flux parameters
        self.work_schedule = work_schedule
        self.collapse_every_n = collapse_every_n
        self.collapse_work = collapse_work
        
        # Tracking
        self.history = {
            'train_loss': [],
            'train_acc': [],
            'val_loss': [],
            'val_acc': [],
            'entropy': [],
            'learning_rates': [],
            'work_budgets': []
        }
        
        self.epoch = 0
    
    def compute_work_budget(self, epoch: int, total_epochs: int) -> float:
        """Compute work budget based on schedule."""
        if self.work_schedule == 'constant':
            return 0.01
        elif self.work_schedule == 'linear':
            # Linear increase
            return 0.01 + 0.04 * (epoch / total_epochs)
        elif self.work_schedule == 'exponential':
            return 0.01 * (1.1 ** epoch)
        elif self.work_schedule == 'cosine':
            # Cosine annealing
            return 0.01 + 0.04 * (1 + np.cos(np.pi * epoch / total_epochs)) / 2
        else:
            return 0.01
    
    def train_epoch(
        self, 
        train_loader: DataLoader, 
        epoch: int, 
        total_epochs: int,
        verbose: bool = True
    ) -> dict:
        """Train for one epoch with flux tracking."""
        self.model.train()
        
        work_budget = self.compute_work_budget(epoch, total_epochs)
        criterion = nn.CrossEntropyLoss()
        
        running_loss = 0.0
        correct = 0
        total = 0
        
        pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{total_epochs} [Train]', 
                    disable=not verbose)
        
        for batch_idx, (inputs, targets) in enumerate(pbar):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            
            # Forward pass
            outputs = self.model(inputs)
            loss = criterion(outputs, targets)
            
            # Backward pass
            self.optimizer.zero_grad()
            loss.backward()
            
            # Flux-aware update
            lr = self.optimizer.param_groups[0]['lr']
            
            # Standard optimizer step
            self.optimizer.step()
            
            # Flux update (manual tracking)
            self.model.flux_update_all(loss.detach(), lr * 0.1, work_budget)
            
            # Periodic collapse
            if (batch_idx + 1) % self.collapse_every_n == 0:
                self.model.collapse_all(self.collapse_work * 0.1)
            
            # Statistics
            running_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            pbar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'acc': f'{100.*correct/total:.2f}%',
                'entropy': f'{self.model.total_entropy():.2f}'
            })
        
        avg_loss = running_loss / total
        accuracy = 100. * correct / total
        
        return {
            'loss': avg_loss,
            'accuracy': accuracy,
            'entropy': self.model.total_entropy(),
            'work_budget': work_budget
        }
    
    @torch.no_grad()
    def evaluate(self, val_loader: DataLoader, verbose: bool = True) -> dict:
        """Evaluate model on validation set."""
        self.model.eval()
        
        criterion = nn.CrossEntropyLoss()
        running_loss = 0.0
        correct = 0
        total = 0
        
        pbar = tqdm(val_loader, desc='Evaluating', disable=not verbose)
        
        for inputs, targets in pbar:
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            
            outputs = self.model(inputs)
            loss = criterion(outputs, targets)
            
            running_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            pbar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'acc': f'{100.*correct/total:.2f}%'
            })
        
        avg_loss = running_loss / total
        accuracy = 100. * correct / total
        
        return {
            'loss': avg_loss,
            'accuracy': accuracy
        }
    
    def train(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        num_epochs: int = 100,
        verbose: bool = True
    ) -> dict:
        """Full training loop."""
        print(f"\n{'='*60}")
        print(f"Flux ResNet Training")
        print(f"{'='*60}")
        print(f"Device: {self.device}")
        print(f"Initial entropy: {self.model.total_entropy():.2f}")
        print(f"Work schedule: {self.work_schedule}")
        print(f"{'='*60}\n")
        
        # LR scheduler: StepLR with milestones
        self.scheduler = optim.lr_scheduler.MultiStepLR(
            self.optimizer,
            milestones=[int(num_epochs * 0.5), int(num_epochs * 0.75)],
            gamma=0.1
        )
        
        best_val_acc = 0.0
        best_model_state = None
        
        for epoch in range(num_epochs):
            self.epoch = epoch
            
            # Train
            train_stats = self.train_epoch(train_loader, epoch, num_epochs, verbose)
            
            # Evaluate
            val_stats = self.evaluate(val_loader, verbose)
            
            # Update LR
            self.scheduler.step()
            current_lr = self.optimizer.param_groups[0]['lr']
            
            # Record history
            self.history['train_loss'].append(train_stats['loss'])
            self.history['train_acc'].append(train_stats['accuracy'])
            self.history['val_loss'].append(val_stats['loss'])
            self.history['val_acc'].append(val_stats['accuracy'])
            self.history['entropy'].append(train_stats['entropy'])
            self.history['learning_rates'].append(current_lr)
            self.history['work_budgets'].append(train_stats['work_budget'])
            
            # Print summary
            if verbose or (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch+1:3d} | "
                      f"Train: {train_stats['accuracy']:.2f}% | "
                      f"Val: {val_stats['accuracy']:.2f}% | "
                      f"Entropy: {train_stats['entropy']:.2f} | "
                      f"LR: {current_lr:.4f}")
            
            # Save best model
            if val_stats['accuracy'] > best_val_acc:
                best_val_acc = val_stats['accuracy']
                best_model_state = {
                    'epoch': epoch,
                    'model_state': self.model.state_dict(),
                    'optimizer_state': self.optimizer.state_dict(),
                    'val_acc': best_val_acc,
                    'train_acc': train_stats['accuracy'],
                    'entropy': train_stats['entropy']
                }
        
        print(f"\n{'='*60}")
        print(f"Training Complete!")
        print(f"Best Validation Accuracy: {best_val_acc:.2f}%")
        print(f"Final Entropy: {self.history['entropy'][-1]:.2f}")
        print(f"{'='*60}\n")
        
        return best_model_state
    
    def plot_history(self, save_path: Optional[str] = None):
        """Plot training history."""
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # Loss
        axes[0, 0].plot(self.history['train_loss'], label='Train', marker='o', markersize=3)
        axes[0, 0].plot(self.history['val_loss'], label='Validation', marker='s', markersize=3)
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Loss')
        axes[0, 0].set_title('Loss over Time')
        axes[0, 0].legend()
        axes[0, 0].grid(True, alpha=0.3)
        
        # Accuracy
        axes[0, 1].plot(self.history['train_acc'], label='Train', marker='o', markersize=3)
        axes[0, 1].plot(self.history['val_acc'], label='Validation', marker='s', markersize=3)
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Accuracy (%)')
        axes[0, 1].set_title('Accuracy over Time')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)
        
        # Entropy
        axes[1, 0].plot(self.history['entropy'], label='Total Entropy', color='red', marker='o', markersize=3)
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('Entropy')
        axes[1, 0].set_title('Flux Entropy over Time')
        axes[1, 0].legend()
        axes[1, 0].grid(True, alpha=0.3)
        
        # Learning Rate
        axes[1, 1].plot(self.history['learning_rates'], label='Learning Rate', color='green', marker='o', markersize=3)
        axes[1, 1].set_xlabel('Epoch')
        axes[1, 1].set_ylabel('Learning Rate')
        axes[1, 1].set_title('Learning Rate Schedule')
        axes[1, 1].legend()
        axes[1, 1].grid(True, alpha=0.3)
        axes[1, 1].set_yscale('log')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"Plots saved to {save_path}")
        
        plt.show()


# ==================== CIFAR-10 DATA LOADER ====================

def get_cifar10_loaders(
    data_dir: str = './data',
    batch_size: int = 128,
    num_workers: int = 2
) -> Tuple[DataLoader, DataLoader]:
    """Create CIFAR-10 train and test DataLoaders."""
    
    # Training transforms with augmentation
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    
    # Test transforms (no augmentation)
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    
    # Datasets
    train_dataset = datasets.CIFAR10(
        root=data_dir,
        train=True,
        download=True,
        transform=train_transform
    )
    
    test_dataset = datasets.CIFAR10(
        root=data_dir,
        train=False,
        download=True,
        transform=test_transform
    )
    
    # DataLoaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, test_loader


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

def main():
    """Main training and testing pipeline."""
    print("\n" + "="*60)
    print("Flux ResNet-18 on CIFAR-10")
    print("Based on: The Algebra of Flux (Conditional Collapse Theory)")
    print("="*60 + "\n")
    
    # Configuration
    config = {
        'batch_size': 128,
        'num_epochs': 50,
        'learning_rate': 0.1,
        'momentum': 0.9,
        'weight_decay': 5e-4,
        'work_schedule': 'linear',
        'collapse_every_n': 10,
        'collapse_work': 0.5,
        'init_entropy': 0.1,
        'num_workers': 2,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu'
    }
    
    print("Configuration:")
    for key, value in config.items():
        print(f"  {key}: {value}")
    print()
    
    device = torch.device(config['device'])
    
    # Data
    print("Loading CIFAR-10...")
    train_loader, test_loader = get_cifar10_loaders(
        batch_size=config['batch_size'],
        num_workers=config['num_workers']
    )
    print(f"Train samples: {len(train_loader.dataset)}")
    print(f"Test samples: {len(test_loader.dataset)}\n")
    
    # Model
    print("Creating Flux ResNet-18...")
    model = flux_resnet18(
        num_classes=10,
        init_entropy=config['init_entropy']
    ).to(device)
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Total parameters: {total_params:,}")
    print(f"Trainable parameters: {trainable_params:,}")
    print(f"Initial entropy: {model.total_entropy():.2f}\n")
    
    # Trainer
    trainer = FluxTrainer(
        model=model,
        device=device,
        lr=config['learning_rate'],
        momentum=config['momentum'],
        weight_decay=config['weight_decay'],
        work_schedule=config['work_schedule'],
        collapse_every_n=config['collapse_every_n'],
        collapse_work=config['collapse_work']
    )
    
    # Train
    print("Starting training...\n")
    start_time = time.time()
    
    best_state = trainer.train(
        train_loader=train_loader,
        val_loader=test_loader,
        num_epochs=config['num_epochs'],
        verbose=True
    )
    
    training_time = time.time() - start_time
    print(f"Training time: {training_time:.2f}s ({training_time/60:.2f}min)\n")
    
    # Load best model for testing
    print("Loading best model for final testing...")
    model.load_state_dict(best_state['model_state'])
    
    # Final evaluation
    print("\n" + "="*60)
    print("FINAL EVALUATION ON TEST SET")
    print("="*60)
    
    model.eval()
    criterion = nn.CrossEntropyLoss()
    
    test_loss = 0.0
    correct = 0
    total = 0
    
    # Per-class accuracy
    class_correct = [0] * 10
    class_total = [0] * 10
    cifar_classes = ['plane', 'car', 'bird', 'cat', 'deer', 
                     'dog', 'frog', 'horse', 'ship', 'truck']
    
    with torch.no_grad():
        for inputs, targets in tqdm(test_loader, desc='Testing'):
            inputs, targets = inputs.to(device), targets.to(device)
            
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            test_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            # Per-class
            for i in range(targets.size(0)):
                label = targets[i]
                class_correct[label] += predicted[i].eq(label).sum().item()
                class_total[label] += 1
    
    avg_test_loss = test_loss / total
    test_accuracy = 100. * correct / total
    
    print(f"\nTest Loss: {avg_test_loss:.4f}")
    print(f"Test Accuracy: {test_accuracy:.2f}%")
    print(f"Correct: {correct}/{total}")
    
    # Per-class accuracy
    print("\nPer-class Accuracy:")
    print("-" * 40)
    for i, class_name in enumerate(cifar_classes):
        if class_total[i] > 0:
            acc = 100. * class_correct[i] / class_total[i]
            print(f"  {class_name:8s}: {acc:6.2f}%")
    
    # Entropy report
    print("\n" + "="*60)
    print("ENTROPY REPORT")
    print("="*60)
    
    entropy_report = model.get_entropy_report()
    print(f"\nTotal Entropy: {model.total_entropy():.2f}")
    print(f"Final Entropy (from training): {trainer.history['entropy'][-1]:.2f}")
    
    print("\nLayer-wise Entropy:")
    for layer_name, stats in entropy_report.items():
        if isinstance(stats, dict) and 'mean' in stats:
            print(f"  {layer_name:12s}: mean={stats['mean']:.4f}, "
                  f"max={stats['max']:.4f}, min={stats['min']:.4f}")
    
    # Summary
    print("\n" + "="*60)
    print("TRAINING SUMMARY")
    print("="*60)
    
    summary = {
        'config': config,
        'results': {
            'test_accuracy': test_accuracy,
            'test_loss': avg_test_loss,
            'best_val_acc': best_state['val_acc'],
            'best_train_acc': best_state['train_acc'],
            'final_entropy': trainer.history['entropy'][-1],
            'training_time': training_time
        },
        'history': {
            'train_acc': trainer.history['train_acc'],
            'val_acc': trainer.history['val_acc'],
            'entropy': trainer.history['entropy']
        }
    }
    
    # Print summary
    print(f"Best Validation Accuracy: {best_state['val_acc']:.2f}%")
    print(f"Final Test Accuracy:      {test_accuracy:.2f}%")
    print(f"Training Time:            {training_time:.2f}s")
    print(f"Entropy Reduction:        "
          f"{config['init_entropy'] * 11000000:.2f} → "
          f"{trainer.history['entropy'][-1]:.2f}")
    
    # Save results
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    
    # Save model
    model_save_path = f"flux_resnet18_cifar10_{timestamp}.pth"
    torch.save({
        'model_state': best_state['model_state'],
        'optimizer_state': best_state['optimizer_state'],
        'config': config,
        'test_accuracy': test_accuracy,
        'entropy': trainer.history['entropy'][-1]
    }, model_save_path)
    print(f"\nModel saved to: {model_save_path}")
    
    # Save history
    history_save_path = f"flux_resnet18_history_{timestamp}.json"
    with open(history_save_path, 'w') as f:
        json.dump({
            'train_loss': trainer.history['train_loss'],
            'train_acc': trainer.history['train_acc'],
            'val_loss': trainer.history['val_loss'],
            'val_acc': trainer.history['val_acc'],
            'entropy': trainer.history['entropy'],
            'learning_rates': trainer.history['learning_rates'],
            'work_budgets': trainer.history['work_budgets']
        }, f, indent=2)
    print(f"History saved to: {history_save_path}")
    
    # Plot
    try:
        plot_save_path = f"flux_resnet18_training_{timestamp}.png"
        trainer.plot_history(save_path=plot_save_path)
    except Exception as e:
        print(f"\nCould not save plot: {e}")
        print("Trying to show plot instead...")
        try:
            trainer.plot_history()
        except:
            print("Plotting not available (no display)")
    
    print("\n" + "="*60)
    print("Flux ResNet-18 CIFAR-10 Complete!")
    print("="*60 + "\n")
    
    return summary


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