Privacy-Preserving Machine Learning in Production - Implementing Differential Privacy at Scale

import numpy as np
import torch
import torch.nn as nn
from opacus import PrivacyEngine
from opacus.utils.batch_memory_manager import BatchMemoryManager
from typing import Tuple, List, Dict, Optional
import math
import logging

class PrivacyBudgetManager:
    def __init__(self, total_epsilon: float, total_delta: float, max_epochs: int):
        """
        Manage privacy budget allocation across training
        
        Args:
            total_epsilon: Total privacy budget (smaller = more private)
            total_delta: Failure probability (typically 1e-5)
            max_epochs: Maximum training epochs
        """
        self.total_epsilon = total_epsilon
        self.total_delta = total_delta
        self.max_epochs = max_epochs
        self.spent_epsilon = 0.0
        self.epoch_budgets = []
        
        # Allocate budget across epochs (more early, less later)
        self._allocate_budget()
        
    def _allocate_budget(self):
        """Allocate privacy budget across epochs using diminishing returns"""
        
        # Use harmonic series for budget allocation
        harmonic_weights = [1.0 / (i + 1) for i in range(self.max_epochs)]
        total_weight = sum(harmonic_weights)
        
        for i in range(self.max_epochs):
            epoch_epsilon = (harmonic_weights[i] / total_weight) * self.total_epsilon
            self.epoch_budgets.append(epoch_epsilon)
            
        logging.info(f"Privacy budget allocated: {self.epoch_budgets[:5]}... (first 5 epochs)")
        
    def get_epoch_budget(self, epoch: int) -> float:
        """Get privacy budget for specific epoch"""
        if epoch < len(self.epoch_budgets):
            return self.epoch_budgets[epoch]
        return 0.0
        
    def spend_budget(self, epsilon: float) -> bool:
        """Spend privacy budget and check if within limits"""
        if self.spent_epsilon + epsilon <= self.total_epsilon:
            self.spent_epsilon += epsilon
            return True
        return False
        
    def get_remaining_budget(self) -> float:
        """Get remaining privacy budget"""
        return max(0.0, self.total_epsilon - self.spent_epsilon)
        
    def privacy_accounting(self, noise_multiplier: float, sample_rate: float, 
                          steps: int) -> Tuple[float, float]:
        """
        Calculate privacy loss using Renyi Differential Privacy accounting
        
        Args:
            noise_multiplier: Noise scale relative to clipping norm
            sample_rate: Fraction of data used per step
            steps: Number of training steps
            
        Returns:
            (epsilon, delta) privacy loss
        """
        from opacus.accountants import RDPAccountant
        
        accountant = RDPAccountant()
        
        # Add noise for each step
        for _ in range(steps):
            accountant.step(
                noise_multiplier=noise_multiplier,
                sample_rate=sample_rate
            )
            
        # Get privacy loss
        epsilon = accountant.get_epsilon(delta=self.total_delta)
        
        return epsilon, self.total_delta

class DifferentiallyPrivateTrainer:
    def __init__(self, model: nn.Module, privacy_budget_manager: PrivacyBudgetManager):
        self.model = model
        self.privacy_manager = privacy_budget_manager
        self.privacy_engine = None
        
        # DP-SGD hyperparameters
        self.max_grad_norm = 1.0  # Gradient clipping bound
        self.noise_multiplier = 1.1  # Noise scale
        self.secure_rng = None
        
    def setup_differential_privacy(self, optimizer, data_loader, 
                                 target_epsilon: float, target_delta: float):
        """Setup differential privacy with Opacus"""
        
        # Attach privacy engine
        self.privacy_engine = PrivacyEngine(secure_mode=True)
        
        # Make model, optimizer, and data_loader private
        self.model, optimizer, data_loader = self.privacy_engine.make_private_with_epsilon(
            module=self.model,
            optimizer=optimizer,
            data_loader=data_loader,
            target_epsilon=target_epsilon,
            target_delta=target_delta,
            epochs=self.privacy_manager.max_epochs,
            max_grad_norm=self.max_grad_norm
        )
        
        # Get actual noise multiplier used
        self.noise_multiplier = optimizer.noise_multiplier
        
        logging.info(f"DP-SGD setup: noise_multiplier={self.noise_multiplier:.3f}, "
                    f"max_grad_norm={self.max_grad_norm}")
        
        return self.model, optimizer, data_loader
        
    def private_training_step(self, batch_data, batch_labels, optimizer, criterion):
        """Execute differentially private training step"""
        
        # Clear gradients
        optimizer.zero_grad()
        
        # Forward pass
        outputs = self.model(batch_data)
        loss = criterion(outputs, batch_labels)
        
        # Backward pass with gradient clipping and noise addition
        loss.backward()
        
        # Optimizer step automatically adds calibrated noise
        optimizer.step()
        
        return loss.item(), outputs
        
    def get_privacy_spent(self) -> Tuple[float, float]:
        """Get current privacy expenditure"""
        if self.privacy_engine:
            epsilon = self.privacy_engine.accountant.get_epsilon(
                delta=self.privacy_manager.total_delta
            )
            return epsilon, self.privacy_manager.total_delta
        return 0.0, 0.0

import hashlib
import secrets
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64

class SecureAggregator:
    def __init__(self, num_parties: int, threshold: int):
        """
        Secure aggregation using secret sharing
        
        Args:
            num_parties: Total number of participating parties
            threshold: Minimum parties needed for reconstruction
        """
        self.num_parties = num_parties
        self.threshold = threshold
        self.party_keys = {}
        self.shared_secrets = {}
        
    def generate_party_keys(self) -> Dict[int, bytes]:
        """Generate unique keys for each party"""
        
        for party_id in range(self.num_parties):
            # Generate unique key for each party
            key = Fernet.generate_key()
            self.party_keys[party_id] = key
            
        return self.party_keys
        
    def secret_share_weights(self, weights: torch.Tensor, party_id: int) -> List[torch.Tensor]:
        """
        Share model weights using Shamir's Secret Sharing
        
        Args:
            weights: Model weights to share
            party_id: ID of the party sharing weights
            
        Returns:
            List of shares for each party
        """
        
        # Flatten weights for easier processing
        flat_weights = weights.flatten().numpy()
        
        # Generate random polynomial coefficients
        coefficients = [flat_weights]  # a_0 = secret
        
        for _ in range(self.threshold - 1):
            coeff = np.random.randn(len(flat_weights)).astype(np.float32)
            coefficients.append(coeff)
            
        # Generate shares for each party
        shares = []
        for party in range(self.num_parties):
            x = party + 1  # Avoid x=0
            share = np.zeros_like(flat_weights)
            
            for i, coeff in enumerate(coefficients):
                share += coeff * (x ** i)
                
            shares.append(torch.tensor(share.reshape(weights.shape)))
            
        return shares
        
    def reconstruct_secret(self, shares: List[Tuple[int, torch.Tensor]], 
                          original_shape: torch.Size) -> torch.Tensor:
        """
        Reconstruct secret from shares using Lagrange interpolation
        
        Args:
            shares: List of (party_id, share) tuples
            original_shape: Original tensor shape
            
        Returns:
            Reconstructed secret tensor
        """
        
        if len(shares) < self.threshold:
            raise ValueError(f"Need at least {self.threshold} shares, got {len(shares)}")
            
        # Extract party IDs and share values
        x_values = [party_id + 1 for party_id, _ in shares[:self.threshold]]
        y_values = [share.flatten().numpy() for _, share in shares[:self.threshold]]
        
        # Lagrange interpolation to recover secret (a_0)
        secret = np.zeros_like(y_values[0])
        
        for i in range(self.threshold):
            # Calculate Lagrange basis polynomial
            basis = np.ones_like(y_values[i])
            
            for j in range(self.threshold):
                if i != j:
                    basis *= (0 - x_values[j]) / (x_values[i] - x_values[j])
                    
            secret += basis * y_values[i]
            
        return torch.tensor(secret.reshape(original_shape))

class FederatedLearningCoordinator:
    def __init__(self, global_model: nn.Module, num_clients: int, 
                 aggregation_method: str = "fedavg"):
        self.global_model = global_model
        self.num_clients = num_clients
        self.aggregation_method = aggregation_method
        self.client_models = {}
        self.client_weights = {}
        
        # Privacy parameters
        self.secure_aggregator = SecureAggregator(num_clients, threshold=num_clients//2 + 1)
        self.privacy_budgets = {}
        
    def initialize_clients(self, privacy_budgets: Dict[int, float]):
        """Initialize federated learning clients"""
        
        self.privacy_budgets = privacy_budgets
        
        # Generate secure keys
        party_keys = self.secure_aggregator.generate_party_keys()
        
        for client_id in range(self.num_clients):
            # Create client model (copy of global model)
            self.client_models[client_id] = self._create_client_model()
            
            # Initialize client privacy budget
            self.client_weights[client_id] = 1.0 / self.num_clients  # Equal weights initially
            
        return party_keys
        
    def _create_client_model(self) -> nn.Module:
        """Create a copy of global model for client"""
        client_model = type(self.global_model)()  # Assuming default constructor
        client_model.load_state_dict(self.global_model.state_dict())
        return client_model
        
    def client_update(self, client_id: int, local_data_loader, epochs: int = 1) -> Dict:
        """Perform local client update with differential privacy"""
        
        client_model = self.client_models[client_id]
        client_model.train()
        
        # Setup differential privacy for client
        privacy_manager = PrivacyBudgetManager(
            total_epsilon=self.privacy_budgets[client_id],
            total_delta=1e-5,
            max_epochs=epochs
        )
        
        dp_trainer = DifferentiallyPrivateTrainer(client_model, privacy_manager)
        
        optimizer = torch.optim.SGD(client_model.parameters(), lr=0.01)
        criterion = nn.CrossEntropyLoss()
        
        # Make training differentially private
        client_model, optimizer, local_data_loader = dp_trainer.setup_differential_privacy(
            optimizer, local_data_loader,
            target_epsilon=self.privacy_budgets[client_id],
            target_delta=1e-5
        )
        
        # Local training
        for epoch in range(epochs):
            for batch_idx, (data, target) in enumerate(local_data_loader):
                loss, _ = dp_trainer.private_training_step(data, target, optimizer, criterion)
                
        # Get privacy spent
        epsilon_spent, delta = dp_trainer.get_privacy_spent()
        
        # Extract model updates (difference from global model)
        model_update = {}
        global_params = dict(self.global_model.named_parameters())
        
        for name, param in client_model.named_parameters():
            model_update[name] = param.data - global_params[name].data
            
        return {
            'client_id': client_id,
            'model_update': model_update,
            'privacy_spent': epsilon_spent,
            'num_samples': len(local_data_loader.dataset)
        }
        
    def secure_aggregate(self, client_updates: List[Dict]) -> Dict[str, torch.Tensor]:
        """Securely aggregate client updates"""
        
        if self.aggregation_method == "fedavg":
            return self._federated_averaging(client_updates)
        elif self.aggregation_method == "secure_fedavg":
            return self._secure_federated_averaging(client_updates)
        else:
            raise ValueError(f"Unknown aggregation method: {self.aggregation_method}")
            
    def _federated_averaging(self, client_updates: List[Dict]) -> Dict[str, torch.Tensor]:
        """Standard federated averaging"""
        
        # Calculate sample-weighted average
        total_samples = sum(update['num_samples'] for update in client_updates)
        aggregated_update = {}
        
        # Initialize aggregated parameters
        first_update = client_updates[0]['model_update']
        for param_name in first_update:
            aggregated_update[param_name] = torch.zeros_like(first_update[param_name])
            
        # Weighted aggregation
        for update in client_updates:
            weight = update['num_samples'] / total_samples
            
            for param_name, param_update in update['model_update'].items():
                aggregated_update[param_name] += weight * param_update
                
        return aggregated_update
        
    def _secure_federated_averaging(self, client_updates: List[Dict]) -> Dict[str, torch.Tensor]:
        """Secure federated averaging using secret sharing"""
        
        aggregated_update = {}
        first_update = client_updates[0]['model_update']
        
        for param_name in first_update:
            # Collect shares for this parameter
            shares = []
            
            for update in client_updates:
                client_id = update['client_id']
                param_tensor = update['model_update'][param_name]
                
                # Create shares
                param_shares = self.secure_aggregator.secret_share_weights(param_tensor, client_id)
                shares.append((client_id, param_shares[0]))  # Use first share for simplicity
                
            # Reconstruct aggregated parameter
            original_shape = first_update[param_name].shape
            aggregated_param = self.secure_aggregator.reconstruct_secret(shares, original_shape)
            
            # Average across clients
            aggregated_update[param_name] = aggregated_param / len(client_updates)
            
        return aggregated_update
        
    def update_global_model(self, aggregated_update: Dict[str, torch.Tensor]):
        """Update global model with aggregated updates"""
        
        with torch.no_grad():
            for name, param in self.global_model.named_parameters():
                if name in aggregated_update:
                    param.data += aggregated_update[name]
                    
    def federated_learning_round(self, client_data_loaders: Dict[int, torch.utils.data.DataLoader],
                                local_epochs: int = 1) -> Dict:
        """Execute one round of federated learning"""
        
        client_updates = []
        round_privacy_spent = {}
        
        # Collect updates from all clients
        for client_id, data_loader in client_data_loaders.items():
            update = self.client_update(client_id, data_loader, local_epochs)
            client_updates.append(update)
            round_privacy_spent[client_id] = update['privacy_spent']
            
        # Securely aggregate updates
        aggregated_update = self.secure_aggregate(client_updates)
        
        # Update global model
        self.update_global_model(aggregated_update)
        
        return {
            'num_clients': len(client_updates),
            'privacy_spent': round_privacy_spent,
            'global_model_state': self.global_model.state_dict()
        }

import tenseal as ts
import torch.nn.functional as F
from typing import List, Tuple

class HomomorphicInferenceEngine:
    def __init__(self, model: nn.Module, poly_modulus_degree: int = 8192):
        """
        Setup homomorphic encryption for private inference
        
        Args:
            model: Trained model for inference
            poly_modulus_degree: Security parameter for encryption
        """
        self.model = model
        self.context = None
        self.encrypted_weights = {}
        self.setup_encryption_context(poly_modulus_degree)
        
    def setup_encryption_context(self, poly_modulus_degree: int):
        """Setup TenSEAL encryption context"""
        
        # Create TenSEAL context
        self.context = ts.context(
            ts.SCHEME_TYPE.CKKS,
            poly_modulus_degree=poly_modulus_degree,
            coeff_mod_bit_sizes=[60, 40, 40, 60]
        )
        
        # Set scale for fixed-point arithmetic
        self.context.global_scale = 2**40
        
        # Generate keys
        self.context.generate_galois_keys()
        
        logging.info(f"Homomorphic encryption context created with "
                    f"polynomial degree {poly_modulus_degree}")
        
    def encrypt_model_weights(self):
        """Encrypt model weights for homomorphic computation"""
        
        self.model.eval()
        
        for name, param in self.model.named_parameters():
            # Convert to list for encryption
            param_list = param.detach().flatten().tolist()
            
            # Encrypt weights
            encrypted_param = ts.ckks_vector(self.context, param_list)
            
            self.encrypted_weights[name] = {
                'encrypted_data': encrypted_param,
                'original_shape': param.shape
            }
            
        logging.info(f"Encrypted {len(self.encrypted_weights)} weight tensors")
        
    def encrypt_input(self, input_data: torch.Tensor) -> ts.CKKSVector:
        """Encrypt input data for private inference"""
        
        # Flatten input and convert to list
        input_list = input_data.flatten().tolist()
        
        # Encrypt input
        encrypted_input = ts.ckks_vector(self.context, input_list)
        
        return encrypted_input
        
    def homomorphic_linear_layer(self, encrypted_input: ts.CKKSVector, 
                                layer_name: str, bias: bool = True) -> ts.CKKSVector:
        """
        Perform linear transformation in encrypted space
        
        Args:
            encrypted_input: Encrypted input vector
            layer_name: Name of the linear layer
            bias: Whether to add bias term
            
        Returns:
            Encrypted output of linear layer
        """
        
        # Get encrypted weights
        weight_name = f"{layer_name}.weight"
        encrypted_weights = self.encrypted_weights[weight_name]['encrypted_data']
        
        # Matrix multiplication in encrypted space
        # This is simplified - real implementation needs proper matrix operations
        encrypted_output = encrypted_input.dot(encrypted_weights)
        
        # Add bias if present
        if bias:
            bias_name = f"{layer_name}.bias"
            if bias_name in self.encrypted_weights:
                encrypted_bias = self.encrypted_weights[bias_name]['encrypted_data']
                encrypted_output = encrypted_output + encrypted_bias
                
        return encrypted_output
        
    def homomorphic_activation(self, encrypted_input: ts.CKKSVector, 
                             activation: str = "relu") -> ts.CKKSVector:
        """
        Apply activation function in encrypted space
        
        Note: Non-linear activations are challenging in homomorphic encryption.
        This uses polynomial approximations.
        """
        
        if activation == "relu":
            # ReLU approximation using polynomial
            # ReLU(x) ≈ max(0, x) can be approximated with degree-2 polynomial
            # This is a simplified approximation
            
            # Square the input (degree 2)
            squared = encrypted_input * encrypted_input
            
            # Linear combination for ReLU approximation
            # This needs careful calibration based on input range
            approx_relu = encrypted_input * 0.5 + squared * 0.25
            
            return approx_relu
            
        elif activation == "sigmoid":
            # Sigmoid approximation: σ(x) ≈ 0.5 + 0.25x (linear approximation)
            # For better approximation, use higher degree polynomials
            
            linear_approx = encrypted_input * 0.25
            constant_term = ts.ckks_vector(self.context, [0.5] * len(encrypted_input.decrypt()))
            
            return constant_term + linear_approx
            
        else:
            raise ValueError(f"Activation {activation} not supported in homomorphic encryption")
            
    def private_inference(self, encrypted_input: ts.CKKSVector, 
                         input_shape: Tuple[int, ...]) -> ts.CKKSVector:
        """
        Perform full model inference in encrypted space
        
        Args:
            encrypted_input: Encrypted input data
            input_shape: Original shape of input
            
        Returns:
            Encrypted model output
        """
        
        # This is a simplified example for a basic feedforward network
        # Real implementation needs to handle specific model architecture
        
        current_encrypted = encrypted_input
        
        # Assuming a simple feedforward model
        layer_names = [name for name, _ in self.model.named_modules() 
                      if isinstance(_, (nn.Linear, nn.Conv2d))]
        
        for i, layer_name in enumerate(layer_names):
            # Linear transformation
            current_encrypted = self.homomorphic_linear_layer(
                current_encrypted, layer_name, bias=True
            )
            
            # Apply activation (except for output layer)
            if i < len(layer_names) - 1:
                current_encrypted = self.homomorphic_activation(
                    current_encrypted, "relu"
                )
                
        return current_encrypted
        
    def decrypt_result(self, encrypted_output: ts.CKKSVector, 
                      output_shape: Tuple[int, ...]) -> torch.Tensor:
        """Decrypt final result"""
        
        # Decrypt output
        decrypted_list = encrypted_output.decrypt()
        
        # Convert back to tensor with original shape
        result_tensor = torch.tensor(decrypted_list).reshape(output_shape)
        
        return result_tensor

class PrivateInferenceService:
    def __init__(self, model_path: str):
        """Initialize private inference service"""
        
        # Load model
        self.model = torch.load(model_path)
        self.model.eval()
        
        # Setup homomorphic encryption
        self.he_engine = HomomorphicInferenceEngine(self.model)
        self.he_engine.encrypt_model_weights()
        
        # Performance metrics
        self.inference_times = []
        self.privacy_overhead = []
        
    def private_predict(self, input_data: torch.Tensor) -> torch.Tensor:
        """
        Perform privacy-preserving prediction
        
        Args:
            input_data: Input tensor for prediction
            
        Returns:
            Prediction result (same as normal inference)
        """
        
        start_time = time.time()
        
        # Encrypt input
        encrypted_input = self.he_engine.encrypt_input(input_data)
        
        encryption_time = time.time()
        
        # Perform encrypted inference
        encrypted_output = self.he_engine.private_inference(
            encrypted_input, input_data.shape
        )
        
        inference_time = time.time()
        
        # Decrypt result
        result = self.he_engine.decrypt_result(
            encrypted_output, self.model(input_data).shape
        )
        
        total_time = time.time() - start_time
        
        # Record performance metrics
        self.inference_times.append(total_time)
        
        # Calculate privacy overhead (vs normal inference)
        normal_start = time.time()
        with torch.no_grad():
            normal_result = self.model(input_data)
        normal_time = time.time() - normal_start
        
        overhead = total_time / normal_time
        self.privacy_overhead.append(overhead)
        
        logging.info(f"Private inference completed: "
                    f"total_time={total_time:.3f}s, "
                    f"overhead={overhead:.1f}x, "
                    f"encryption={encryption_time-start_time:.3f}s, "
                    f"inference={inference_time-encryption_time:.3f}s")
        
        return result
        
    def get_performance_stats(self) -> Dict:
        """Get performance statistics"""
        
        if not self.inference_times:
            return {}
            
        return {
            'avg_inference_time': np.mean(self.inference_times),
            'p95_inference_time': np.percentile(self.inference_times, 95),
            'avg_privacy_overhead': np.mean(self.privacy_overhead),
            'p95_privacy_overhead': np.percentile(self.privacy_overhead, 95),
            'total_inferences': len(self.inference_times)
        }

from kubernetes import client, config
import asyncio
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

class PrivacyPreservingMLPlatform:
    def __init__(self):
        self.federated_coordinators = {}
        self.private_inference_services = {}
        self.privacy_auditor = PrivacyAuditor()
        
        # Kubernetes client for managing federated workers
        config.load_incluster_config()
        self.k8s_client = client.AppsV1Api()
        
    def deploy_federated_learning_job(self, job_config: Dict) -> str:
        """Deploy federated learning job on Kubernetes"""
        
        job_name = job_config['job_name']
        num_clients = job_config['num_clients']
        privacy_budget = job_config['privacy_budget']
        
        # Create federated coordinator
        coordinator = FederatedLearningCoordinator(
            global_model=self._load_model(job_config['model_config']),
            num_clients=num_clients
        )
        
        # Initialize privacy budgets
        client_budgets = {i: privacy_budget / num_clients for i in range(num_clients)}
        coordinator.initialize_clients(client_budgets)
        
        self.federated_coordinators[job_name] = coordinator
        
        # Deploy Kubernetes jobs for each client
        self._deploy_federated_clients(job_name, job_config)
        
        return job_name
        
    def _deploy_federated_clients(self, job_name: str, job_config: Dict):
        """Deploy federated learning client pods"""
        
        for client_id in range(job_config['num_clients']):
            deployment_manifest = {
                "apiVersion": "apps/v1",
                "kind": "Deployment",
                "metadata": {
                    "name": f"{job_name}-client-{client_id}",
                    "labels": {"app": "federated-client", "job": job_name}
                },
                "spec": {
                    "replicas": 1,
                    "selector": {"matchLabels": {"client-id": str(client_id)}},
                    "template": {
                        "metadata": {"labels": {"client-id": str(client_id)}},
                        "spec": {
                            "containers": [{
                                "name": "federated-client",
                                "image": "privacy-ml:latest",
                                "env": [
                                    {"name": "CLIENT_ID", "value": str(client_id)},
                                    {"name": "JOB_NAME", "value": job_name},
                                    {"name": "COORDINATOR_URL", "value": f"http://{job_name}-coordinator:8000"},
                                    {"name": "PRIVACY_BUDGET", "value": str(job_config['privacy_budget'] / job_config['num_clients'])}
                                ],
                                "resources": {
                                    "requests": {"memory": "4Gi", "cpu": "2"},
                                    "limits": {"memory": "8Gi", "cpu": "4"}
                                }
                            }]
                        }
                    }
                }
            }
            
            # Deploy client
            self.k8s_client.create_namespaced_deployment(
                namespace="federated-learning",
                body=deployment_manifest
            )

class PrivacyAuditor:
    def __init__(self):
        self.audit_logs = []
        self.privacy_violations = []
        
    def audit_privacy_budget(self, job_name: str, client_id: int, 
                           epsilon_spent: float, total_budget: float) -> bool:
        """Audit privacy budget usage"""
        
        audit_entry = {
            'timestamp': time.time(),
            'job_name': job_name,
            'client_id': client_id,
            'epsilon_spent': epsilon_spent,
            'total_budget': total_budget,
            'budget_remaining': total_budget - epsilon_spent
        }
        
        self.audit_logs.append(audit_entry)
        
        # Check for budget violation
        if epsilon_spent > total_budget:
            violation = {
                'timestamp': time.time(),
                'job_name': job_name,
                'client_id': client_id,
                'violation_type': 'budget_exceeded',
                'severity': 'critical',
                'details': f"Privacy budget exceeded: {epsilon_spent} > {total_budget}"
            }
            
            self.privacy_violations.append(violation)
            
            # Alert administrators
            self._send_privacy_alert(violation)
            
            return False
            
        return True
        
    def generate_privacy_report(self, job_name: str) -> Dict:
        """Generate comprehensive privacy report"""
        
        job_logs = [log for log in self.audit_logs if log['job_name'] == job_name]
        job_violations = [v for v in self.privacy_violations if v['job_name'] == job_name]
        
        if not job_logs:
            return {'error': 'No audit logs found for job'}
            
        # Calculate statistics
        total_budget_used = sum(log['epsilon_spent'] for log in job_logs)
        total_budget_allocated = sum(log['total_budget'] for log in job_logs)
        
        client_budgets = {}
        for log in job_logs:
            client_id = log['client_id']
            if client_id not in client_budgets:
                client_budgets[client_id] = {
                    'spent': 0.0,
                    'allocated': log['total_budget']
                }
            client_budgets[client_id]['spent'] += log['epsilon_spent']
            
        return {
            'job_name': job_name,
            'total_privacy_budget_used': total_budget_used,
            'total_privacy_budget_allocated': total_budget_allocated,
            'budget_utilization_rate': total_budget_used / total_budget_allocated,
            'client_privacy_budgets': client_budgets,
            'privacy_violations': len(job_violations),
            'audit_entries': len(job_logs),
            'compliance_status': 'compliant' if len(job_violations) == 0 else 'violations_detected'
        }
        
    def _send_privacy_alert(self, violation: Dict):
        """Send privacy violation alert"""
        # Implementation would send alerts via Slack, email, etc.
        logging.critical(f"PRIVACY VIOLATION: {violation}")

# FastAPI service for privacy-preserving ML
app = FastAPI(title="Privacy-Preserving ML Platform")
platform = PrivacyPreservingMLPlatform()

class FederatedJobRequest(BaseModel):
    job_name: str
    model_config: Dict
    num_clients: int
    privacy_budget: float
    data_sources: List[str]

class PrivateInferenceRequest(BaseModel):
    model_name: str
    input_data: List[float]
    privacy_level: str = "high"

@app.post("/federated-learning/start")
async def start_federated_learning(job_request: FederatedJobRequest):
    """Start federated learning job"""
    
    try:
        job_id = platform.deploy_federated_learning_job(job_request.dict())
        
        return {
            'job_id': job_id,
            'status': 'started',
            'message': f'Federated learning job {job_id} started with {job_request.num_clients} clients'
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/inference/private")
async def private_inference(inference_request: PrivateInferenceRequest):
    """Perform privacy-preserving inference"""
    
    try:
        # Get inference service
        service_name = inference_request.model_name
        
        if service_name not in platform.private_inference_services:
            # Load service
            platform.private_inference_services[service_name] = PrivateInferenceService(
                f"models/{service_name}.pt"
            )
            
        service = platform.private_inference_services[service_name]
        
        # Convert input
        input_tensor = torch.tensor([inference_request.input_data])
        
        # Perform private inference
        if inference_request.privacy_level == "high":
            result = service.private_predict(input_tensor)
        else:
            # Standard inference for lower privacy requirements
            with torch.no_grad():
                result = service.model(input_tensor)
                
        return {
            'prediction': result.tolist(),
            'privacy_level': inference_request.privacy_level,
            'model_name': service_name
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/audit/privacy/{job_name}")
async def get_privacy_audit(job_name: str):
    """Get privacy audit report for job"""
    
    report = platform.privacy_auditor.generate_privacy_report(job_name)
    return report

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

class SecureDrugDiscoveryPlatform:
    def __init__(self):
        self.participating_organizations = {}
        self.shared_computation_engine = None
        
    def setup_multi_party_computation(self, organizations: List[str], 
                                    computation_threshold: int):
        """Setup MPC for collaborative drug discovery"""
        
        # Each organization contributes data without sharing raw data
        # Computation happens on encrypted shares
        
        for org in organizations:
            self.participating_organizations[org] = {
                'data_shares': None,
                'computation_shares': None,
                'privacy_budget': 5.0  # ε=5.0 for drug discovery
            }
            
        self.shared_computation_engine = SMCEngine(
            parties=len(organizations),
            threshold=computation_threshold
        )
        
    def collaborative_model_training(self, target_disease: str) -> Dict:
        """Train model across organizations without data sharing"""
        
        # Each organization computes on their local data
        # Results are combined using secure aggregation
        
        local_results = {}
        
        for org_name, org_info in self.participating_organizations.items():
            # Simulate local computation
            local_model_params = self._train_local_model(
                org_name, target_disease, org_info['privacy_budget']
            )
            
            # Convert to secret shares
            local_results[org_name] = self.shared_computation_engine.create_shares(
                local_model_params
            )
            
        # Aggregate using secure computation
        global_model = self.shared_computation_engine.secure_aggregate(
            local_results
        )
        
        return {
            'global_model': global_model,
            'participating_organizations': len(self.participating_organizations),
            'privacy_preserved': True,
            'target_disease': target_disease
        }