trustformers-mobile 0.1.1

//! Federated Learning with Differential Privacy v2.0
//!
//! This module implements next-generation federated learning with advanced differential
//! privacy mechanisms, secure aggregation protocols, and cutting-edge cryptographic
//! techniques for privacy-preserving mobile AI training.

pub mod crypto;
pub mod privacy;

// Re-export main types
pub use crypto::*;
pub use privacy::*;

// For now, include remaining types directly until full refactoring
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use trustformers_core::{Tensor};
use trustformers_core::error::{CoreError, Result};

/// Federated Learning v2.0 configuration with advanced privacy features
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FederatedLearningV2Config {
    /// Privacy configuration
    pub privacy_config: AdvancedPrivacyConfig,
    /// Cryptographic protocols
    pub crypto_config: CryptographicConfig,
    /// Aggregation configuration
    pub aggregation_config: SecureAggregationConfig,
    /// Communication protocols
    pub communication_config: CommunicationProtocolConfig,
    /// Model training configuration
    pub training_config: FederatedTrainingConfig,
    /// Security and robustness settings
    pub security_config: SecurityConfig,
    /// Privacy accounting configuration
    pub accounting_config: PrivacyAccountingConfig,
}

/// Secure aggregation configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecureAggregationConfig {
    /// Minimum number of participants
    pub min_participants: u32,
    /// Maximum number of participants
    pub max_participants: u32,
    /// Dropout resilience threshold
    pub dropout_threshold: f64,
    /// Byzantine fault tolerance
    pub byzantine_fault_tolerance: bool,
    /// Aggregation function
    pub aggregation_function: AggregationFunction,
    /// Verification methods
    pub verification_methods: Vec<VerificationMethod>,
}

/// Aggregation functions
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AggregationFunction {
    /// Federated averaging
    FederatedAveraging,
    /// Weighted averaging
    WeightedAveraging,
    /// Median aggregation
    MedianAggregation,
    /// Trimmed mean
    TrimmedMean,
    /// Coordinate-wise median
    CoordinateWiseMedian,
    /// Geometric median
    GeometricMedian,
}

/// Verification methods
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum VerificationMethod {
    /// Digital signatures
    DigitalSignatures,
    /// Zero-knowledge proofs
    ZeroKnowledgeProofs,
    /// Commitment schemes
    CommitmentSchemes,
    /// Hash-based verification
    HashBased,
}

/// Communication protocol configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommunicationProtocolConfig {
    /// Transport security settings
    pub transport_security: TransportSecurity,
    /// Compression configuration
    pub compression: CompressionConfig,
    /// Bandwidth adaptation
    pub bandwidth_adaptation: BandwidthAdaptationConfig,
    /// Message routing
    pub message_routing: MessageRoutingConfig,
}

/// Transport security configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransportSecurity {
    /// Use TLS encryption
    pub use_tls: bool,
    /// TLS version
    pub tls_version: String,
    /// Certificate validation
    pub certificate_validation: bool,
}

/// Compression configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressionConfig {
    /// Compression algorithm
    pub algorithm: CompressionAlgorithm,
    /// Compression level (1-9)
    pub compression_level: u8,
    /// Adaptive compression
    pub adaptive_compression: bool,
}

/// Compression algorithms
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionAlgorithm {
    /// No compression
    None,
    /// GZIP compression
    Gzip,
    /// Brotli compression
    Brotli,
    /// ZSTD compression
    Zstd,
}

/// Bandwidth adaptation configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BandwidthAdaptationConfig {
    /// Adaptive batching
    pub adaptive_batching: bool,
    /// Quality of service priority
    pub qos_priority: u8,
    /// Congestion control
    pub congestion_control: bool,
}

/// Message routing configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageRoutingConfig {
    /// Use gossip protocol
    pub use_gossip: bool,
    /// Gossip fanout
    pub gossip_fanout: u32,
    /// Message TTL
    pub message_ttl: u32,
}

/// Federated training configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FederatedTrainingConfig {
    /// Number of training rounds
    pub num_rounds: u32,
    /// Client sampling strategy
    pub client_sampling: ClientSamplingStrategy,
    /// Model update frequency
    pub update_frequency: u32,
    /// Personalization settings
    pub personalization: PersonalizationConfig,
}

/// Client sampling strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ClientSamplingStrategy {
    /// Random sampling
    Random,
    /// Round-robin sampling
    RoundRobin,
    /// Performance-based sampling
    PerformanceBased,
    /// Resource-aware sampling
    ResourceAware,
}

/// Personalization configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersonalizationConfig {
    /// Enable personalization
    pub enabled: bool,
    /// Personalization strategy
    pub strategy: PersonalizationStrategy,
    /// Local update ratio
    pub local_update_ratio: f64,
}

/// Personalization strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PersonalizationStrategy {
    /// Local fine-tuning
    LocalFineTuning,
    /// Meta-learning
    MetaLearning,
    /// Multi-task learning
    MultiTaskLearning,
    /// Transfer learning
    TransferLearning,
}

/// Security configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityConfig {
    /// Attack detection enabled
    pub attack_detection_enabled: bool,
    /// Byzantine fault tolerance
    pub byzantine_fault_tolerance: bool,
    /// Robust aggregation methods
    pub robust_aggregation_methods: Vec<RobustAggregationMethod>,
    /// Anomaly detection thresholds
    pub anomaly_detection_thresholds: AnomalyDetectionThresholds,
}

/// Robust aggregation methods
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RobustAggregationMethod {
    /// Trimmed mean
    TrimmedMean,
    /// Median
    Median,
    /// Krum
    Krum,
    /// Bulyan
    Bulyan,
    /// FoolsGold
    FoolsGold,
}

/// Anomaly detection thresholds
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AnomalyDetectionThresholds {
    /// Gradient norm threshold
    pub gradient_norm_threshold: f64,
    /// Model accuracy threshold
    pub model_accuracy_threshold: f64,
    /// Participation rate threshold
    pub participation_rate_threshold: f64,
}

// Default implementations
impl Default for FederatedLearningV2Config {
    fn default() -> Self {
        Self {
            privacy_config: AdvancedPrivacyConfig::default(),
            crypto_config: CryptographicConfig::default(),
            aggregation_config: SecureAggregationConfig::default(),
            communication_config: CommunicationProtocolConfig::default(),
            training_config: FederatedTrainingConfig::default(),
            security_config: SecurityConfig::default(),
            accounting_config: PrivacyAccountingConfig::default(),
        }
    }
}

impl Default for SecureAggregationConfig {
    fn default() -> Self {
        Self {
            min_participants: 3,
            max_participants: 1000,
            dropout_threshold: 0.1,
            byzantine_fault_tolerance: true,
            aggregation_function: AggregationFunction::FederatedAveraging,
            verification_methods: vec![VerificationMethod::DigitalSignatures],
        }
    }
}

impl Default for CommunicationProtocolConfig {
    fn default() -> Self {
        Self {
            transport_security: TransportSecurity::default(),
            compression: CompressionConfig::default(),
            bandwidth_adaptation: BandwidthAdaptationConfig::default(),
            message_routing: MessageRoutingConfig::default(),
        }
    }
}

impl Default for TransportSecurity {
    fn default() -> Self {
        Self {
            use_tls: true,
            tls_version: "1.3".to_string(),
            certificate_validation: true,
        }
    }
}

impl Default for CompressionConfig {
    fn default() -> Self {
        Self {
            algorithm: CompressionAlgorithm::Gzip,
            compression_level: 6,
            adaptive_compression: true,
        }
    }
}

impl Default for BandwidthAdaptationConfig {
    fn default() -> Self {
        Self {
            adaptive_batching: true,
            qos_priority: 3,
            congestion_control: true,
        }
    }
}

impl Default for MessageRoutingConfig {
    fn default() -> Self {
        Self {
            use_gossip: false,
            gossip_fanout: 3,
            message_ttl: 100,
        }
    }
}

impl Default for FederatedTrainingConfig {
    fn default() -> Self {
        Self {
            num_rounds: 100,
            client_sampling: ClientSamplingStrategy::Random,
            update_frequency: 1,
            personalization: PersonalizationConfig::default(),
        }
    }
}

impl Default for PersonalizationConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            strategy: PersonalizationStrategy::LocalFineTuning,
            local_update_ratio: 0.1,
        }
    }
}

impl Default for SecurityConfig {
    fn default() -> Self {
        Self {
            attack_detection_enabled: true,
            byzantine_fault_tolerance: true,
            robust_aggregation_methods: vec![RobustAggregationMethod::TrimmedMean],
            anomaly_detection_thresholds: AnomalyDetectionThresholds::default(),
        }
    }
}

impl Default for AnomalyDetectionThresholds {
    fn default() -> Self {
        Self {
            gradient_norm_threshold: 10.0,
            model_accuracy_threshold: 0.5,
            participation_rate_threshold: 0.8,
        }
    }
}

/// Main federated learning v2.0 implementation
pub struct FederatedLearningV2 {
    config: FederatedLearningV2Config,
    privacy_budget_tracker: PrivacyBudgetTracker,
    key_manager: CryptographicKeyManager,
    participants: HashMap<String, ParticipantInfo>,
    round_number: u32,
}

/// Participant information
#[derive(Debug, Clone)]
pub struct ParticipantInfo {
    pub id: String,
    pub public_key: Vec<u8>,
    pub last_seen: std::time::SystemTime,
    pub trust_score: f64,
    pub contribution_history: Vec<f64>,
}

impl FederatedLearningV2 {
    /// Create new federated learning instance
    pub fn new(config: FederatedLearningV2Config) -> Self {
        let privacy_budget_tracker = PrivacyBudgetTracker::new(
            config.accounting_config.max_epsilon,
            config.accounting_config.max_delta,
            config.accounting_config.accounting_method,
        );

        Self {
            config,
            privacy_budget_tracker,
            key_manager: CryptographicKeyManager::new(),
            participants: HashMap::new(),
            round_number: 0,
        }
    }

    /// Add participant to the federation
    pub fn add_participant(&mut self, participant_info: ParticipantInfo) {
        self.participants.insert(participant_info.id.clone(), participant_info);
    }

    /// Remove participant from the federation
    pub fn remove_participant(&mut self, participant_id: &str) {
        self.participants.remove(participant_id);
        self.key_manager.remove_participant_key(participant_id);
    }

    /// Start federated training round
    pub fn start_training_round(&mut self) -> Result<()> {
        self.round_number += 1;

        // Check privacy budget
        let epsilon_per_round = self.config.privacy_config.epsilon / self.config.training_config.num_rounds as f64;
        let delta_per_round = self.config.privacy_config.delta / self.config.training_config.num_rounds as f64;

        if !self.privacy_budget_tracker.can_consume(epsilon_per_round, delta_per_round) {
            return Err(TrustformersError::runtime_error("Privacy budget exceeded".into()).into());
        }

        // Consume privacy budget
        self.privacy_budget_tracker.consume(epsilon_per_round, delta_per_round)
            .map_err(|e| TrustformersError::runtime_error(e))?;

        Ok(())
    }

    /// Get current round number
    pub fn get_round_number(&self) -> u32 {
        self.round_number
    }

    /// Get participant count
    pub fn get_participant_count(&self) -> usize {
        self.participants.len()
    }

    /// Get privacy budget status
    pub fn get_privacy_budget_status(&self) -> (f64, f64) {
        self.privacy_budget_tracker.remaining_budget()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_federated_learning_v2_creation() {
        let config = FederatedLearningV2Config::default();
        let fl = FederatedLearningV2::new(config);
        assert_eq!(fl.get_round_number(), 0);
        assert_eq!(fl.get_participant_count(), 0);
    }

    #[test]
    fn test_participant_management() {
        let config = FederatedLearningV2Config::default();
        let mut fl = FederatedLearningV2::new(config);

        let participant = ParticipantInfo {
            id: "participant1".to_string(),
            public_key: vec![1, 2, 3, 4],
            last_seen: std::time::SystemTime::now(),
            trust_score: 0.9,
            contribution_history: vec![0.8, 0.85, 0.9],
        };

        fl.add_participant(participant);
        assert_eq!(fl.get_participant_count(), 1);

        fl.remove_participant("participant1");
        assert_eq!(fl.get_participant_count(), 0);
    }

    #[test]
    fn test_training_round_privacy_budget() {
        let config = FederatedLearningV2Config::default();
        let mut fl = FederatedLearningV2::new(config);

        // Should be able to start multiple rounds within budget
        for _ in 0..10 {
            assert!(fl.start_training_round().is_ok());
        }

        assert_eq!(fl.get_round_number(), 10);

        let (remaining_eps, remaining_delta) = fl.get_privacy_budget_status();
        assert!(remaining_eps >= 0.0);
        assert!(remaining_delta >= 0.0);
    }
}