"""
Structured-Factorial Federated Learning
=========================================
Server (Teacher): Full dataset, computes structured embeddings
Clients (Students): Partial datasets, receive structured guidance

Training: Server governs client learning paths (work-based)
Inference: Clients verify independently (simple)
"""

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


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

class StructuredFactorialServer:
    """
    Server = Teacher with comprehensive knowledge
    - Full training set
    - Computes structured embeddings (topologically constrained)
    - Governs multiple client models
    - High λ_sec during orchestration (work)
    """
    
    def __init__(self, n_layers: int, embedding_dim: int, gamma: float = 1.5):
        self.n_layers = n_layers
        self.embedding_dim = embedding_dim
        self.gamma = gamma
        
        # Structured-Factorial computation
        self.F_S = 1  # Chain topology
        self.n_factorial = math.factorial(n_layers)
        self.confinement_ratio = self.F_S / self.n_factorial
        
        # Work parameters (high for server orchestration)
        self.lambda_base = 10.0   # High work for server
        self.kappa = 2.0
        
        # Damping factors per layer (constraint depth)
        self.damping_factors = [1.0 / (1 + i * 0.5) for i in range(n_layers)]
        
        # Client registry
        self.registered_clients = {}
        
        print(f"Server initialized: F_S={self.F_S}, R={self.confinement_ratio:.6f}")
        print(f"λ_base (orchestration): {self.lambda_base}")
    
    def compute_lambda(self, H_AS: float = 0.0) -> float:
        """Server works hard to compute structured guidance"""
        confinement = (self.n_factorial / self.F_S) ** self.gamma
        return self.lambda_base + self.kappa * H_AS * confinement
    
    def compute_work(self) -> float:
        """Work to structure knowledge for clients"""
        return math.log(self.n_factorial) - math.log(self.F_S)
    
    def register_client(self, client_id: int, client_data_size: int) -> Dict:
        """Register a client and compute its structured configuration"""
        self.registered_clients[client_id] = {
            'data_size': client_data_size,
            'structured_config': self._compute_client_config(client_data_size)
        }
        return self.registered_clients[client_id]
    
    def _compute_client_config(self, client_data_size: int) -> Dict:
        """
        Compute structured configuration for client based on its data size.
        Smaller data = more structured guidance needed.
        """
        # Data scarcity factor (more scarcity = more structure needed)
        scarcity_factor = 1.0 / (1 + math.log(1 + client_data_size))
        
        # Structured embedding based on constraint topology
        config = {
            'embedding_scale': scarcity_factor,
            'damping_weights': [d * scarcity_factor for d in self.damping_factors],
            'structured_dim': int(self.embedding_dim * (1 - scarcity_factor * 0.5)),
            'work_allocation': self.compute_work() * scarcity_factor
        }
        return config
    
    def generate_structured_embedding(self, features: torch.Tensor, 
                                      client_id: int) -> torch.Tensor:
        """
        Generate structured embedding for client.
        Server's work: project features through topological constraints.
        """
        if client_id not in self.registered_clients:
            client_id = 0
        
        config = self.registered_clients[client_id]['structured_config']
        lambda_sec = self.compute_lambda()
        
        batch_size, feat_dim = features.shape
        
        # Project to structured space (server's work)
        proj_weight = torch.randn(feat_dim, config['structured_dim']) * config['embedding_scale']
        proj_weight = proj_weight.to(features.device)
        
        embedded = features @ proj_weight
        
        # Apply layer-wise damping (topological filtering)
        for layer_idx in range(self.n_layers):
            if layer_idx < embedded.shape[-1] // config['structured_dim']:
                damping = config['damping_weights'][layer_idx]
                start_idx = layer_idx * config['structured_dim']
                end_idx = start_idx + config['structured_dim']
                embedded[:, start_idx:end_idx] *= damping
        
        # Apply λ_sec as global damping (server's work cost)
        embedded = embedded / (1 + lambda_sec * 0.1)
        
        return embedded
    
    def distill_guidance(self, features: torch.Tensor, 
                        labels: torch.Tensor, client_id: int) -> Dict[str, torch.Tensor]:
        """
        Server distills knowledge into structured guidance for client.
        This is the hard work: computing what client needs to learn.
        """
        # Generate structured embeddings
        embedded = self.generate_structured_embedding(features, client_id)
        
        # Compute guidance signal (teacher's knowledge)
        config = self.registered_clients[client_id]['structured_config']
        
        guidance = {
            'structured_embedding': embedded,
            'damping_weights': torch.tensor(config['damping_weights']).to(features.device),
            'lambda_sec': self.compute_lambda(),
            'work_done': config['work_allocation']
        }
        
        return guidance


# ============================================================================
# CLIENT MODEL (Student)
# ============================================================================

class StructuredClient(nn.Module):
    """
    Client = Student with partial knowledge
    - Small dataset (student book)
    - Receives structured guidance from server (teacher's help)
    - Learns faster with fewer samples due to structured input
    """
    
    def __init__(self, client_id: int, config: 'SFConfig', 
                 server: StructuredFactorialServer):
        super().__init__()
        self.client_id = client_id
        self.config = config
        self.server = server
        
        # Receive structured config from server
        self.server_config = server.register_client(client_id, config.client_data_size)
        
        # Structured embedding dimension (from server)
        self.structured_dim = self.server_config['structured_config']['structured_dim']
        
        # Input projection (receives structured embeddings)
        self.input_proj = nn.Linear(config.input_size, self.structured_dim)
        
        # Structured layers (initialized with server's damping)
        self.structured_layers = nn.ModuleList([
            nn.Linear(self.structured_dim, dim) 
            for dim in config.hidden_dims
        ])
        
        # Classifier (receives structured features)
        self.classifier = nn.Linear(config.hidden_dims[-1], config.num_classes)
        
        # Receive damping weights from server
        self.damping_weights = self.server_config['structured_config']['damping_weights']
        
        self._init_structured()
    
    def _init_structured(self):
        """Initialize with server's structured configuration"""
        for i, layer in enumerate(self.structured_layers):
            damping = self.damping_weights[i] if i < len(self.damping_weights) else 1.0
            nn.init.xavier_uniform_(layer.weight)
            layer.weight.data *= damping
            nn.init.zeros_(layer.bias)
    
    def forward(self, x: torch.Tensor, 
                use_server_guidance: bool = True) -> Tuple[torch.Tensor, Optional[Dict]]:
        """
        Forward pass.
        - Training: Uses server's structured guidance (fast learning)
        - Inference: Independent forward (simple verification)
        """
        # Project input to structured space
        h = self.input_proj(x)
        h = torch.relu(h)
        
        # Apply structured layers with damping
        for i, layer in enumerate(self.structured_layers):
            h = layer(h)
            h = torch.relu(h)
            
            # Apply server's damping (if using guidance)
            if use_server_guidance:
                damping = self.damping_weights[i] if i < len(self.damping_weights) else 1.0
                h = h * damping
        
        logits = self.classifier(h)
        return logits, None
    
    def forward_with_guidance(self, x: torch.Tensor, 
                              guidance: Dict) -> Tuple[torch.Tensor, float]:
        """
        Forward with server's distilled guidance.
        This is the student's learning phase with teacher's help.
        """
        # Use server's structured embedding
        if 'structured_embedding' in guidance:
            h = guidance['structured_embedding']
        else:
            h = self.input_proj(x)
            h = torch.relu(h)
        
        # Apply structured layers with server's damping
        for i, layer in enumerate(self.structured_layers):
            h = layer(h)
            h = torch.relu(h)
            
            damping = guidance['damping_weights'][i].item() if i < len(guidance['damping_weights']) else 1.0
            h = h * damping
        
        logits = self.classifier(h)
        work = guidance.get('work_done', 0.0)
        
        return logits, work


# ============================================================================
# FEDERATED ORCHESTRATION
# ============================================================================

@dataclass
class SFConfig:
    input_size: int = 784
    num_classes: int = 10
    hidden_dims: List[int] = field(default_factory=lambda: [256, 128])
    n_layers: int = 2
    embedding_dim: int = 64
    gamma: float = 1.5
    n_clients: int = 5
    client_data_fraction: float = 0.1  # Each client has 10% of data
    batch_size: int = 64
    server_epochs: int = 10
    client_epochs: int = 3
    learning_rate: float = 0.001


class FederatedPipeline:
    """
    Federated learning with structured-factorial guidance.
    
    Server (Teacher): Orchestrates, computes structured embeddings
    Clients (Students): Learn with guidance, fewer samples needed
    """
    
    def __init__(self, config: SFConfig):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Initialize server
        self.server = StructuredFactorialServer(
            n_layers=config.n_layers,
            embedding_dim=config.embedding_dim,
            gamma=config.gamma
        )
        
        # Initialize clients (students)
        self.clients = []
        self.client_data_loaders = []
        
        self._prepare_clients()
        
        # Global model (aggregated knowledge)
        self.global_model = None
        
        self.history = {
            'round': [], 'server_work': [],
            'client_train_acc': [], 'client_test_acc': []
        }
    
    def _prepare_clients(self):
        """Prepare client datasets (each has partial data)"""
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        
        # Full dataset
        full_dataset = datasets.MNIST(root='../../data', train=True, download=True, transform=transform)
        n_total = len(full_dataset)
        
        # Split into clients (each gets fraction of data)
        n_per_client = int(n_total * self.config.client_data_fraction / self.config.n_clients)
        
        indices = np.random.permutation(n_total)
        
        for i in range(self.config.n_clients):
            # Client gets subset
            start_idx = i * n_per_client
            end_idx = start_idx + n_per_client
            client_indices = indices[start_idx:end_idx]
            
            client_dataset = Subset(full_dataset, client_indices)
            client_loader = DataLoader(
                client_dataset, 
                batch_size=self.config.batch_size, 
                shuffle=True
            )
            self.client_data_loaders.append(client_loader)
            
            # Create client config
            client_config = SFConfig(
                input_size=self.config.input_size,
                num_classes=self.config.num_classes,
                hidden_dims=self.config.hidden_dims,
                n_layers=self.config.n_layers,
                embedding_dim=self.config.embedding_dim,
                gamma=self.config.gamma,
                client_data_size=len(client_dataset)
            )
            
            # Create client (student)
            client = StructuredClient(i, client_config, self.server).to(self.device)
            self.clients.append(client)
            
            print(f"Client {i}: {len(client_dataset)} samples (student's book)")
        
        print(f"\nServer has access to all {n_total} samples (teacher's book)")
        print(f"Each client has ~{100*self.config.client_data_fraction:.0f}% of data")
    
    def _prepare_test_loader(self):
        """Test set (shared by all)"""
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
        return DataLoader(test_dataset, batch_size=self.config.batch_size, shuffle=False)
    
    def server_orchestrates(self, round_num: int) -> Dict:
        """
        Server orchestrates the learning round.
        Work: compute structured guidance for all clients.
        """
        print(f"\n{'='*60}")
        print(f"SERVER ORCHESTRATION ROUND {round_num} (Teacher's Work)")
        print(f"{'='*60}")
        
        total_work = 0.0
        
        # Server samples from its full dataset (teacher's knowledge)
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        full_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
        
        # Server works on samples that clients don't have
        # (teaches what students haven't seen)
        sample_indices = np.random.choice(len(full_dataset), 
                                         size=min(500, len(full_dataset)), 
                                         replace=False)
        
        for idx in sample_indices[:50]:  # Sample some for guidance computation
            data, label = full_dataset[idx]
            data = data.view(-1, self.config.input_size).to(self.device)
            
            # Server computes structured guidance for all clients
            for client_id in range(self.config.n_clients):
                guidance = self.server.distill_guidance(data, label, client_id)
                total_work += guidance['work_done']
        
        lambda_sec = self.server.compute_lambda()
        print(f"Server work done: {total_work:.4f} | λ_sec: {lambda_sec:.2f}")
        
        return {'work': total_work, 'lambda_sec': lambda_sec}
    
    def client_learns(self, client_id: int, round_num: int) -> Tuple[float, float]:
        """
        Client learns with server's structured guidance.
        With guidance, client learns faster with fewer samples.
        """
        client = self.clients[client_id]
        client.train()
        
        total_loss, correct, total = 0, 0, 0
        
        # Server orchestrates this client's learning
        self.server_orchestrates(round_num)
        
        for epoch in range(self.config.client_epochs):
            for batch_idx, (data, target) in enumerate(self.client_data_loaders[client_id]):
                data, target = data.to(self.device), target.to(self.device)
                data = data.view(-1, self.config.input_size)
                
                self.optimizer = optim.Adam(client.parameters(), lr=self.config.learning_rate)
                
                # Get structured guidance from server
                guidance = self.server.distill_guidance(data, target, client_id)
                
                # Client learns with guidance
                output, work = client.forward_with_guidance(data, guidance)
                
                # Loss with work term (client contributes to effort)
                loss = nn.CrossEntropyLoss()(output, target) + 0.001 * work
                
                loss.backward()
                self.optimizer.step()
                
                total_loss += loss.item()
                correct += (output.argmax(1) == target).sum().item()
                total += target.size(0)
        
        return total_loss / len(self.client_data_loaders[client_id]), 100. * correct / total
    
    @torch.no_grad()
    def evaluate_clients(self) -> Tuple[float, float]:
        """Evaluate all clients on test set"""
        test_loader = self._prepare_test_loader()
        
        all_correct, all_total = 0, 0
        
        for client in self.clients:
            client.eval()
            for data, target in test_loader:
                data, target = data.to(self.device), target.to(self.device)
                data = data.view(-1, self.config.input_size)
                
                output, _ = client.forward(data, use_server_guidance=True)
                pred = output.argmax(1)
                all_correct += (pred == target).sum().item()
                all_total += target.size(0)
        
        return 0.0, 100. * all_correct / all_total
    
    def federated_train(self):
        """Full federated learning loop"""
        print("\n" + "="*70)
        print("FEDERATED STRUCTURED-LEARNING")
        print("Server (Teacher): Full dataset | Clients (Students): Partial datasets")
        print("="*70)
        
        for round_num in range(1, self.config.server_epochs + 1):
            # Server orchestrates
            orch_result = self.server_orchestrates(round_num)
            
            # All clients learn with guidance
            round_train_acc = []
            for client_id in range(self.config.n_clients):
                loss, acc = self.client_learns(client_id, round_num)
                round_train_acc.append(acc)
            
            # Evaluate
            _, test_acc = self.evaluate_clients()
            
            # Store history
            self.history['round'].append(round_num)
            self.history['server_work'].append(orch_result['work'])
            self.history['client_train_acc'].append(np.mean(round_train_acc))
            self.history['client_test_acc'].append(test_acc)
            
            print(f"\nRound {round_num}: Client Avg Acc: {np.mean(round_train_acc):.1f}% | "
                  f"Test Acc: {test_acc:.1f}% | Server Work: {orch_result['work']:.2f}")
        
        print("\n" + "="*70)
        print("FEDERATED LEARNING COMPLETE")
        print("="*70)
        self._final_summary()
        
        return self.history
    
    def _final_summary(self):
        """Summary of federated learning"""
        print(f"\nClients learned with ~{100*self.config.client_data_fraction:.0f}% of data each")
        print(f"Server provided structured guidance (teacher's work)")
        print(f"\nFinal Results:")
        print(f"  Client Avg Train Acc: {self.history['client_train_acc'][-1]:.2f}%")
        print(f"  Test Acc: {self.history['client_test_acc'][-1]:.2f}%")
        print(f"  Total Server Work: {sum(self.history['server_work']):.2f}")
    
    def compare_no_guidance(self):
        """Compare client learning WITH vs WITHOUT server guidance"""
        print("\n" + "="*60)
        print("ABLATION: With Guidance vs Without Guidance")
        print("="*60)
        
        # Train one client WITH guidance
        print("\n[Client WITH server guidance]")
        loss, acc = self.client_learns(0, 0)
        print(f"  Train Loss: {loss:.4f}, Acc: {acc:.1f}%")
        
        # Reinitialize client
        client_config = SFConfig(client_data_size=len(list(self.client_data_loaders[0].dataset)))
        client_no_guidance = StructuredClient(99, client_config, self.server).to(self.device)
        
        print("\n[Client WITHOUT server guidance]")
        client_no_guidance.train()
        optimizer = optim.Adam(client_no_guidance.parameters(), lr=self.config.learning_rate)
        
        for epoch in range(self.config.client_epochs):
            total_loss, correct, total = 0, 0, 0
            for data, target in self.client_data_loaders[0]:
                data, target = data.to(self.device), target.to(self.device)
                data = data.view(-1, self.config.input_size)
                
                optimizer.zero_grad()
                output, _ = client_no_guidance.forward(data, use_server_guidance=False)
                loss = nn.CrossEntropyLoss()(output, target)
                loss.backward()
                optimizer.step()
                
                total_loss += loss.item()
                correct += (output.argmax(1) == target).sum().item()
                total += target.size(0)
        
        no_guide_acc = 100. * correct / total
        print(f"  Train Loss: {total_loss/len(self.client_data_loaders[0]):.4f}, Acc: {no_guide_acc:.1f}%")
        
        print(f"\n[Result]")
        print(f"  WITH guidance: ~{acc:.1f}%")
        print(f"  WITHOUT guidance: ~{no_guide_acc:.1f}%")
        print(f"  → Server guidance improves learning with same data!")


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

def main():
    print("\n" + "="*70)
    print("STRUCTURED-FACTORIAL FEDERATED LEARNING")
    print("Server = Teacher (full knowledge) | Clients = Students (partial knowledge)")
    print("="*70)
    
    config = SFConfig(
        n_clients=5,
        client_data_fraction=0.1,  # Each client has 10% of data
        server_epochs=5,
        client_epochs=2,
        hidden_dims=[256, 128],
        n_layers=2
    )
    
    pipeline = FederatedPipeline(config)
    history = pipeline.federated_train()
    
    # Ablation study
    pipeline.compare_no_guidance()
    
    return pipeline, history


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