"""
Structured-Factorial Server Model for MNIST
============================================
Server structures initial layers for client models using constrained
permutation topologies derived from the Structured-Factorial theory.

Architecture:
- StructuredFactorialServer: Computes and serves constrained layer configurations
- ClientModel: Uses structured layers from server for classification
- MNIST training/testing pipeline with full evaluation
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
import numpy as np
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional
from enum import Enum
import math
from collections import defaultdict
import warnings

warnings.filterwarnings('ignore')

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

@dataclass
class SFConfig:
    """Configuration for Structured-Factorial Server"""
    # MNIST parameters
    input_size: int = 784
    num_classes: int = 10
    
    # Structured-Factorial parameters
    n_layers: int = 4                    # Number of constrained layers
    constraint_topology: str = "dag"      # Constraint DAG type
    gamma: float = 1.5                   # Confinement exponent
    
    # Server layer configuration
    hidden_dims: List[int] = field(default_factory=lambda: [256, 128, 64])
    server_embedding_dim: int = 32       # Structured embedding dimension
    
    # Training parameters
    batch_size: int = 128
    epochs: int = 20
    learning_rate: float = 0.001
    test_split: float = 0.15             # Portion for testing
    
    # Device
    device: str = "cuda" if torch.cuda.is_available() else "cpu"


# ============================================================================
# STRUCTURED-FACTORIAL SERVER
# ============================================================================

class ConstraintDAG:
    """
    Represents the security topology as a DAG of layer constraints.
    Based on Definition 1.2.1 - partial order constraints on layer configurations.
    """
    
    def __init__(self, n_nodes: int):
        self.n_nodes = n_nodes
        self.edges = set()  # (i, j) means i -> j (i must precede j)
        self.nodes = list(range(n_nodes))
        
    def add_constraint(self, before: int, after: int):
        """Add ordering constraint: before must come before after"""
        if before < self.n_nodes and after < self.n_nodes:
            self.edges.add((before, after))
    
    def get_predecessors(self, node: int) -> List[int]:
        """Get all nodes that must come before this node"""
        return [i for i in self.nodes if (i, node) in self.edges]
    
    def get_successors(self, node: int) -> List[int]:
        """Get all nodes that come after this node"""
        return [i for i in self.nodes if (node, i) in self.edges]
    
    def is_valid_path(self, path: List[int]) -> bool:
        """Check if a path respects all constraints"""
        for i, node in enumerate(path):
            predecessors = self.get_predecessors(node)
            for pred in predecessors:
                if pred not in path[:i]:  # Predecessor not before this node
                    return False
        return True


class StructuredFactorial:
    """
    Computes F_S(n) - the number of valid linear extensions of a poset.
    Chapter 1, Definition 1.2.1: Counts valid permutations under constraints.
    """
    
    def __init__(self, constraint_dag: ConstraintDAG):
        self.dag = constraint_dag
        self.n = constraint_dag.n_nodes
        self._cache = {}
        
    def count_linear_extensions(self) -> int:
        """
        Count the number of linear extensions of the poset.
        This is the Structured-Factorial F_S(n).
        """
        cache_key = tuple(sorted(self.dag.edges))
        if cache_key in self._cache:
            return self._cache[cache_key]
        
        # Base case: single element
        if self.n == 0:
            return 1
        if self.n == 1:
            return 1
        
        # Find elements with no predecessors (source nodes)
        sources = self._find_sources()
        
        count = 0
        for source in sources:
            # Temporarily remove source and recurse
            new_dag = self._remove_node(source)
            sf = StructuredFactorial(new_dag)
            count += sf.count_linear_extensions()
        
        self._cache[cache_key] = count
        return count
    
    def _find_sources(self) -> List[int]:
        """Find all nodes with no incoming edges (can be placed next)"""
        sources = []
        for node in self.dag.nodes:
            if len(self.dag.get_predecessors(node)) == 0:
                sources.append(node)
        return sources
    
    def _remove_node(self, node: int) -> ConstraintDAG:
        """Create a new DAG with the given node removed"""
        new_dag = ConstraintDAG(self.n - 1)
        node_map = {}  # old_node -> new_node
        
        new_idx = 0
        for old_node in self.dag.nodes:
            if old_node != node:
                node_map[old_node] = new_idx
                new_idx += 1
        
        for (u, v) in self.dag.edges:
            if u != node and v != node:
                new_dag.edges.add((node_map[u], node_map[v]))
        
        return new_dag
    
    def security_confinement_ratio(self) -> float:
        """
        R = F_S(n) / n!
        Ratio of valid paths to all possible permutations.
        """
        valid_count = self.count_linear_extensions()
        total_permutations = math.factorial(self.n)
        return valid_count / total_permutations


class StructuredFactorialServer:
    """
    Server that computes structured layer configurations and provides
    them to client models. Implements the Security Threshold Equation
    from Chapter 2.
    
    The server:
    1. Maintains a constraint DAG on layer orderings
    2. Computes the Structured-Factorial for path validation
    3. Generates structured layer configurations (structured embeddings)
    4. Computes security damping coefficient lambda_sec
    """
    
    def __init__(self, config: SFConfig):
        self.config = config
        self.device = torch.device(config.device)
        
        # Build constraint DAG for layer topology
        # Constraint: layer i must precede layer i+1 (chain topology)
        self.constraint_dag = ConstraintDAG(config.n_layers)
        for i in range(config.n_layers - 1):
            self.constraint_dag.add_constraint(i, i + 1)
        
        # Compute Structured-Factorial
        self.sf = StructuredFactorial(self.constraint_dag)
        self.F_S = self.sf.count_linear_extensions()
        self.n_factorial = math.factorial(config.n_layers)
        self.confinement_ratio = self.sf.security_confinement_ratio()
        
        # Compute security damping
        self.lambda_base = 1.0
        self.kappa = 1.0
        self.gamma = config.gamma
        
        # Structured layer templates (computed by server)
        self.layer_configs = self._compute_layer_configs()
        
        # Structured embeddings cache
        self.structured_embeddings = None
        
        print(f"Structured-Factorial Server initialized:")
        print(f"  - F_S({config.n_layers}) = {self.F_S}")
        print(f"  - {config.n_layers}! = {self.n_factorial}")
        print(f"  - Confinement Ratio R = {self.confinement_ratio:.6f}")
        print(f"  - Valid paths: {self._enumerate_valid_paths()}")
    
    def _enumerate_valid_paths(self) -> List[List[int]]:
        """Enumerate all valid paths respecting constraints"""
        paths = []
        
        def backtrack(path: List[int], remaining: set):
            if not remaining:
                paths.append(path.copy())
                return
            
            for node in remaining:
                predecessors = set(self.constraint_dag.get_predecessors(node))
                if predecessors.issubset(set(path)):
                    path.append(node)
                    backtrack(path, remaining - {node})
                    path.pop()
        
        backtrack([], set(self.constraint_dag.nodes))
        return paths
    
    def _compute_layer_configs(self) -> Dict[int, dict]:
        """Compute structured configurations for each layer"""
        configs = {}
        
        for layer_idx in range(self.config.n_layers):
            predecessors = self.constraint_dag.get_predecessors(layer_idx)
            
            configs[layer_idx] = {
                'dim': self.config.hidden_dims[layer_idx] if layer_idx < len(self.config.hidden_dims) 
                       else self.config.hidden_dims[-1],
                'predecessors': predecessors,
                'successors': self.constraint_dag.get_successors(layer_idx),
                'order_index': layer_idx,
                'structured_weight_scale': 1.0 / (1 + len(predecessors)),
            }
        
        return configs
    
    def compute_lambda_sec(self, H_AS: float = 0.0, curl_mag: float = 0.0) -> float:
        """
        Security Threshold Equation (Chapter 2, Theorem 2.5.1):
        λ_sec = λ_base + κ * H_AS * (n! / F_S)^γ + η * ||∇×V_threat||
        """
        kappa = self.kappa
        gamma = self.gamma
        eta = 0.5
        
        # Confinement amplification term
        if self.F_S > 0:
            confinement_ratio = self.n_factorial / self.F_S
            amplification = H_AS * (confinement_ratio ** gamma)
        else:
            amplification = float('inf')
        
        lambda_sec = self.lambda_base + kappa * amplification + eta * curl_mag
        return lambda_sec
    
    def generate_structured_layer(self, layer_idx: int, input_dim: int, 
                                  output_dim: int) -> nn.Module:
        """
        Generate a structured layer based on constraint topology.
        Layers with more constraints get higher damping (lower effective capacity).
        """
        config = self.layer_configs[layer_idx]
        weight_scale = config['structured_weight_scale']
        
        layer = nn.Linear(input_dim, output_dim)
        
        # Scale weights based on constraint topology
        # More constrained layers get smaller initialization (higher damping)
        with torch.no_grad():
            nn.init.xavier_uniform_(layer.weight)
            layer.weight.data *= weight_scale
            nn.init.zeros_(layer.bias)
        
        return layer
    
    def get_structured_embedding(self, batch_size: int, 
                                 base_features: torch.Tensor) -> torch.Tensor:
        """
        Generate structured embeddings from base features.
        Applies topological filtering based on constraint DAG.
        """
        n_layers = self.config.n_layers
        emb_dim = self.config.server_embedding_dim
        
        # Project to structured space
        proj = nn.Linear(base_features.shape[-1], emb_dim * n_layers).to(self.device)
        embeddings = proj(base_features)
        
        # Reshape to (batch, n_layers, emb_dim)
        embeddings = embeddings.view(-1, n_layers, emb_dim)
        
        # Apply layer-wise damping based on constraints
        for i in range(n_layers):
            damping = 1.0 / (1 + len(self.layer_configs[i]['predecessors']))
            embeddings[:, i, :] *= damping
        
        return embeddings
    
    def forward_structured(self, x: torch.Tensor, 
                          H_AS: float = 0.0) -> Tuple[torch.Tensor, float]:
        """
        Process input through structured topology, returning
        processed features and current lambda_sec.
        """
        # Generate structured embedding
        emb = self.get_structured_embedding(x.shape[0], x)
        
        # Apply constraint-based aggregation
        # Features flow through DAG-constrained pathways
        layer_idx = 0
        current = x
        
        # First structured transformation
        if layer_idx < len(self.config.hidden_dims):
            dim = self.config.hidden_dims[layer_idx]
            layer = self.generate_structured_layer(layer_idx, x.shape[-1], dim).to(self.device)
            current = torch.relu(layer(current))
        
        # Apply remaining structured layers
        for layer_idx in range(1, self.config.n_layers):
            if layer_idx < len(self.config.hidden_dims):
                dim = self.config.hidden_dims[layer_idx]
                layer = self.generate_structured_layer(layer_idx, current.shape[-1], dim).to(self.device)
                current = torch.relu(layer(current))
        
        lambda_sec = self.compute_lambda_sec(H_AS)
        return current, lambda_sec


# ============================================================================
# CLIENT MODEL
# ============================================================================

class StructuredClientModel(nn.Module):
    """
    Client model that receives structured initial layers from server.
    Uses the structured layers for MNIST classification.
    """
    
    def __init__(self, config: SFConfig, server: StructuredFactorialServer):
        super().__init__()
        self.config = config
        self.server = server
        self.device = torch.device(config.device)
        
        # Receive structured layers from server
        self.structured_layers = nn.ModuleList()
        self._build_structured_architecture()
        
        # Additional classification head
        self.classifier = nn.Sequential(
            nn.Linear(config.hidden_dims[-1], 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, config.num_classes)
        )
        
        # Structured embedding layer (from server)
        self.structured_proj = nn.Linear(config.input_size, 
                                        config.server_embedding_dim * config.n_layers)
    
    def _build_structured_architecture(self):
        """Build architecture using server's structured layer configs"""
        input_dim = self.config.input_size
        
        for layer_idx in range(self.config.n_layers):
            config = self.server.layer_configs[layer_idx]
            output_dim = config['dim']
            
            # Generate layer from server
            layer = self.server.generate_structured_layer(
                layer_idx, input_dim, output_dim
            )
            self.structured_layers.append(layer)
            
            input_dim = output_dim
    
    def forward(self, x: torch.Tensor, return_lambda: bool = False) -> Tuple:
        """
        Forward pass through structured layers.
        
        Args:
            x: Input tensor (batch, 784)
            return_lambda: Whether to return lambda_sec value
            
        Returns:
            logits: (batch, 10)
            lambda_sec: Security damping coefficient (if return_lambda)
        """
        batch_size = x.shape[0]
        
        # Apply structured projection
        h = self.structured_proj(x)
        h = h.view(batch_size, self.config.n_layers, -1)
        
        # Apply constraint-weighted transformation
        for i, layer in enumerate(self.structured_layers):
            # Get constraint info from server
            predecessors = self.server.layer_configs[i]['predecessors']
            damping = 1.0 / (1 + len(predecessors))
            
            # Apply layer with damping
            h_i = layer(x if i == 0 else h_prev)
            h_i = h_i * damping
            h_i = torch.relu(h_i)
            h_prev = h_i
        
        # Classification
        logits = self.classifier(h_prev)
        
        if return_lambda:
            lambda_sec = self.server.compute_lambda_sec()
            return logits, lambda_sec
        
        return logits


# ============================================================================
# TRAINING AND TESTING PIPELINE
# ============================================================================

class MNISTPipeline:
    """Complete training and testing pipeline for MNIST"""
    
    def __init__(self, config: SFConfig):
        self.config = config
        self.device = torch.device(config.device)
        
        # Initialize server
        self.server = StructuredFactorialServer(config)
        
        # Initialize client model
        self.model = StructuredClientModel(config, self.server).to(self.device)
        
        # Optimizer
        self.optimizer = optim.Adam(self.model.parameters(), lr=config.learning_rate)
        
        # Loss
        self.criterion = nn.CrossEntropyLoss()
        
        # Data
        self.train_loader = None
        self.test_loader = None
        self._prepare_data()
        
        # Tracking
        self.history = {
            'train_loss': [], 'train_acc': [],
            'test_loss': [], 'test_acc': [],
            'lambda_sec': []
        }
    
    def _prepare_data(self):
        """Load and prepare MNIST data"""
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        
        full_dataset = datasets.MNIST(
            root='../../data', train=True, download=True, transform=transform
        )
        
        # Split into train/test
        total_size = len(full_dataset)
        test_size = int(total_size * self.config.test_split)
        train_size = total_size - test_size
        
        # Stratified split
        indices = np.random.permutation(total_size)
        train_indices = indices[:train_size]
        test_indices = indices[train_size:]
        
        train_dataset = Subset(full_dataset, train_indices)
        test_dataset = Subset(full_dataset, test_indices)
        
        self.train_loader = DataLoader(
            train_dataset, batch_size=self.config.batch_size, shuffle=True
        )
        self.test_loader = DataLoader(
            test_dataset, batch_size=self.config.batch_size, shuffle=False
        )
        
        print(f"Data prepared: Train={train_size}, Test={test_size}")
    
    def train_epoch(self, epoch: int) -> Tuple[float, float]:
        """Train for one epoch"""
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            # Flatten data
            data = data.view(data.size(0), -1)
            
            self.optimizer.zero_grad()
            
            # Forward pass (with lambda tracking)
            output, lambda_sec = self.model(data, return_lambda=True)
            
            # Loss with lambda regularization
            base_loss = self.criterion(output, target)
            lambda_reg = 0.001 * lambda_sec  # Small regularization
            loss = base_loss + lambda_reg
            
            loss.backward()
            self.optimizer.step()
            
            total_loss += loss.item()
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
        
        avg_loss = total_loss / len(self.train_loader)
        accuracy = 100. * correct / total
        
        return avg_loss, accuracy
    
    def test(self) -> Tuple[float, float]:
        """Evaluate on test set"""
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for data, target in self.test_loader:
                data, target = data.to(self.device), target.to(self.device)
                data = data.view(data.size(0), -1)
                
                output, _ = self.model(data, return_lambda=True)
                loss = self.criterion(output, target)
                
                total_loss += loss.item()
                pred = output.argmax(dim=1)
                correct += (pred == target).sum().item()
                total += target.size(0)
        
        avg_loss = total_loss / len(self.test_loader)
        accuracy = 100. * correct / total
        
        return avg_loss, accuracy
    
    def train(self) -> dict:
        """Full training loop"""
        print("\n" + "="*60)
        print("STRUCTURED-FACTORIAL CLIENT MODEL TRAINING")
        print("="*60)
        print(f"Device: {self.device}")
        print(f"Model parameters: {sum(p.numel() for p in self.model.parameters()):,}")
        print(f"Epochs: {self.config.epochs}")
        print("="*60 + "\n")
        
        best_test_acc = 0.0
        
        for epoch in range(1, self.config.epochs + 1):
            train_loss, train_acc = self.train_epoch(epoch)
            test_loss, test_acc = self.test()
            
            # Get current lambda_sec
            dummy_input = torch.zeros(1, self.config.input_size).to(self.device)
            _, lambda_sec = self.model(dummy_input, return_lambda=True)
            
            # Store history
            self.history['train_loss'].append(train_loss)
            self.history['train_acc'].append(train_acc)
            self.history['test_loss'].append(test_loss)
            self.history['test_acc'].append(test_acc)
            self.history['lambda_sec'].append(lambda_sec)
            
            # Print progress
            print(f"Epoch {epoch:2d}/{self.config.epochs} | "
                  f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | "
                  f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}% | "
                  f"λ_sec: {lambda_sec:.4f}")
            
            if test_acc > best_test_acc:
                best_test_acc = test_acc
        
        print("\n" + "="*60)
        print("TRAINING COMPLETE")
        print("="*60)
        print(f"Best Test Accuracy: {best_test_acc:.2f}%")
        print(f"Final λ_sec: {self.history['lambda_sec'][-1]:.4f}")
        
        return self.history
    
    def evaluate_attack_scenarios(self):
        """Evaluate model under different 'attack surface entropy' scenarios"""
        print("\n" + "="*60)
        print("ATTACK SURFACE ENTROPY ANALYSIS")
        print("="*60)
        
        H_AS_values = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9]
        
        print(f"\n{'H_AS':<10} {'λ_sec':<12} {'Confinement':<15} {'Test Acc':<12}")
        print("-" * 50)
        
        for H_AS in H_AS_values:
            # Compute lambda_sec for this attack level
            lambda_sec = self.server.compute_lambda_sec(H_AS)
            confinement = self.server.confinement_ratio * math.exp(-self.server.gamma * H_AS)
            
            # Evaluate (using stored model, lambda affects only server-side)
            self.model.eval()
            correct = 0
            total = 0
            
            with torch.no_grad():
                for data, target in self.test_loader:
                    data, target = data.to(self.device), target.to(self.device)
                    data = data.view(data.size(0), -1)
                    output = self.model(data)
                    pred = output.argmax(dim=1)
                    correct += (pred == target).sum().item()
                    total += target.size(0)
            
            accuracy = 100. * correct / total
            print(f"{H_AS:<10.1f} {lambda_sec:<12.4f} {confinement:<15.6f} {accuracy:<12.2f}")
    
    def print_model_summary(self):
        """Print detailed model architecture summary"""
        print("\n" + "="*60)
        print("MODEL ARCHITECTURE SUMMARY")
        print("="*60)
        
        print("\n[Structured-Factorial Server]")
        print(f"  Constraint DAG: {self.config.n_layers}-node chain")
        print(f"  F_S({self.config.n_layers}) = {self.server.F_S}")
        print(f"  {self.config.n_layers}! = {self.server.n_factorial}")
        print(f"  Confinement Ratio R = {self.server.confinement_ratio:.6f}")
        print(f"  γ (confinement exponent) = {self.server.gamma}")
        
        print("\n[Layer Configurations]")
        for layer_idx, config in self.server.layer_configs.items():
            print(f"  Layer {layer_idx}: dim={config['dim']}, "
                  f"preds={config['predecessors']}, scale={config['structured_weight_scale']:.4f}")
        
        print("\n[Client Model Layers]")
        for i, layer in enumerate(self.model.structured_layers):
            print(f"  Structured Layer {i}: {layer.in_features} -> {layer.out_features}")
        
        print(f"\n  Classifier: {self.config.hidden_dims[-1]} -> 64 -> {self.config.num_classes}")
        print(f"\n  Total Parameters: {sum(p.numel() for p in self.model.parameters()):,}")


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

def main():
    """Main execution function"""
    print("\n" + "="*70)
    print("STRUCTURED-FACTORIAL SERVER MODEL")
    print("Mathematical Framework: Topologically-Constrained Probabilities")
    print("="*70 + "\n")
    
    # Configuration
    config = SFConfig(
        n_layers=4,
        hidden_dims=[256, 128, 64, 32],
        gamma=1.5,
        batch_size=128,
        epochs=20,
        learning_rate=0.001
    )
    
    # Initialize and run pipeline
    pipeline = MNISTPipeline(config)
    
    # Print model summary
    pipeline.print_model_summary()
    
    # Train model
    history = pipeline.train()
    
    # Attack surface entropy analysis
    pipeline.evaluate_attack_scenarios()
    
    # Final summary
    print("\n" + "="*70)
    print("FINAL RESULTS")
    print("="*70)
    print(f"Final Training Accuracy: {history['train_acc'][-1]:.2f}%")
    print(f"Final Test Accuracy: {history['test_acc'][-1]:.2f}%")
    print(f"Final λ_sec: {history['lambda_sec'][-1]:.4f}")
    print(f"Confinement Ratio: {config.n_layers}! / F_S = {math.factorial(config.n_layers) / pipeline.server.F_S:.0f}")
    
    return pipeline, history


if __name__ == "__main__":
    pipeline, history = main()