rustytorch/demos/rtx-fedmed-demo/src/privacy.rs

//! Differential privacy mechanisms for federated learning.
//!
//! This module implements privacy-preserving mechanisms including gradient
//! clipping, noise addition, and privacy budget accounting.

use rand::Rng;

use fedmed_shared::{PrivacyAccountantType, PrivacyConfig};

/// Differential privacy mechanism for adding noise to gradients.
#[derive(Debug, Clone)]
pub struct DifferentialPrivacy {
    /// Privacy configuration.
    config: PrivacyConfig,
    /// Random number generator seed (reserved for reproducibility).
    #[allow(dead_code)]
    seed: u64,
}

impl DifferentialPrivacy {
    /// Create a new differential privacy mechanism.
    #[must_use]
    pub fn new(config: PrivacyConfig) -> Self {
        Self { config, seed: 42 }
    }

    /// Add Gaussian noise to gradients for differential privacy.
    #[must_use]
    pub fn add_noise(&self, gradients: &[f32]) -> Vec<f32> {
        let mut rng = rand::thread_rng();
        let sigma = self.config.noise_multiplier as f32 * self.config.clip_norm as f32;

        gradients
            .iter()
            .map(|g| {
                // Box-Muller transform for Gaussian noise
                let u1: f32 = rng.r#gen::<f32>();
                let u2: f32 = rng.r#gen::<f32>();
                let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
                g + noise * sigma
            })
            .collect()
    }

    /// Add noise using local differential privacy (for client-side).
    #[must_use]
    pub fn add_local_noise(&self, gradients: &[f32]) -> Vec<f32> {
        if !self.config.local_dp {
            return gradients.to_vec();
        }

        let mut rng = rand::thread_rng();
        // Local DP typically uses higher noise
        let sigma = self.config.noise_multiplier as f32 * self.config.clip_norm as f32 * 2.0;

        gradients
            .iter()
            .map(|g| {
                let u1: f32 = rng.r#gen::<f32>();
                let u2: f32 = rng.r#gen::<f32>();
                let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
                g + noise * sigma
            })
            .collect()
    }

    /// Get the noise standard deviation.
    #[must_use]
    pub fn noise_std(&self) -> f64 {
        self.config.noise_multiplier * self.config.clip_norm
    }

    /// Get the privacy configuration.
    #[must_use]
    pub fn config(&self) -> &PrivacyConfig {
        &self.config
    }

    /// Update the noise multiplier.
    pub fn set_noise_multiplier(&mut self, multiplier: f64) {
        self.config.noise_multiplier = multiplier;
    }
}

/// Gradient clipper for bounding sensitivity.
#[derive(Debug, Clone)]
pub struct GradientClipper {
    /// Maximum L2 norm for gradients.
    clip_norm: f64,
    /// Whether to use per-sample clipping (reserved for advanced use).
    #[allow(dead_code)]
    per_sample: bool,
}

impl Default for GradientClipper {
    fn default() -> Self {
        Self::new(1.0)
    }
}

impl GradientClipper {
    /// Create a new gradient clipper.
    #[must_use]
    pub fn new(clip_norm: f64) -> Self {
        Self {
            clip_norm,
            per_sample: false,
        }
    }

    /// Create a per-sample gradient clipper.
    #[must_use]
    pub fn per_sample(clip_norm: f64) -> Self {
        Self {
            clip_norm,
            per_sample: true,
        }
    }

    /// Clip gradients to bound their L2 norm.
    #[must_use]
    pub fn clip(&self, gradients: &[f32]) -> Vec<f32> {
        let norm = self.l2_norm(gradients);

        if norm <= self.clip_norm as f32 {
            return gradients.to_vec();
        }

        let scale = self.clip_norm as f32 / norm;
        gradients.iter().map(|g| g * scale).collect()
    }

    /// Clip gradients with adaptive clipping.
    #[must_use]
    pub fn adaptive_clip(&self, gradients: &[f32], quantile: f32) -> Vec<f32> {
        // Compute per-coordinate absolute values
        let mut abs_grads: Vec<f32> = gradients.iter().map(|g| g.abs()).collect();
        abs_grads.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        // Find the quantile threshold
        let idx = ((abs_grads.len() as f32 * quantile) as usize).min(abs_grads.len() - 1);
        let threshold = abs_grads[idx];

        // Clip to threshold
        gradients
            .iter()
            .map(|g| g.clamp(-threshold, threshold))
            .collect()
    }

    /// Compute L2 norm of gradients.
    #[must_use]
    pub fn l2_norm(&self, gradients: &[f32]) -> f32 {
        gradients.iter().map(|g| g * g).sum::<f32>().sqrt()
    }

    /// Compute L1 norm of gradients.
    #[must_use]
    pub fn l1_norm(&self, gradients: &[f32]) -> f32 {
        gradients.iter().map(|g| g.abs()).sum()
    }

    /// Compute infinity norm of gradients.
    #[must_use]
    pub fn linf_norm(&self, gradients: &[f32]) -> f32 {
        gradients.iter().map(|g| g.abs()).fold(0.0_f32, f32::max)
    }

    /// Get the clip norm.
    #[must_use]
    pub fn clip_norm(&self) -> f64 {
        self.clip_norm
    }

    /// Set the clip norm.
    pub fn set_clip_norm(&mut self, clip_norm: f64) {
        self.clip_norm = clip_norm;
    }
}

/// Privacy accountant for tracking privacy budget.
#[derive(Debug, Clone)]
pub struct PrivacyAccountant {
    /// Privacy configuration.
    config: PrivacyConfig,
    /// Total privacy spent (epsilon).
    epsilon_spent: f64,
    /// Number of composition steps.
    num_steps: usize,
    /// Sample rate for subsampling amplification.
    sample_rate: f64,
    /// Accountant type.
    accountant_type: PrivacyAccountantType,
}

impl PrivacyAccountant {
    /// Create a new privacy accountant.
    #[must_use]
    pub fn new(config: PrivacyConfig) -> Self {
        Self {
            accountant_type: config.accountant,
            config,
            epsilon_spent: 0.0,
            num_steps: 0,
            sample_rate: 1.0,
        }
    }

    /// Create an accountant with a specific sample rate.
    #[must_use]
    pub fn with_sample_rate(config: PrivacyConfig, sample_rate: f64) -> Self {
        Self {
            accountant_type: config.accountant,
            config,
            epsilon_spent: 0.0,
            num_steps: 0,
            sample_rate: sample_rate.clamp(0.0, 1.0),
        }
    }

    /// Record a privacy-consuming step and return the epsilon spent.
    pub fn step(&mut self, num_clients: usize) -> f64 {
        self.num_steps += 1;

        // Compute epsilon for this step based on accountant type
        let step_epsilon = match self.accountant_type {
            PrivacyAccountantType::Moments => self.moments_epsilon(num_clients),
            PrivacyAccountantType::RDP => self.rdp_epsilon(num_clients),
            PrivacyAccountantType::GDP => self.gdp_epsilon(num_clients),
            PrivacyAccountantType::PLD => self.pld_epsilon(num_clients),
        };

        self.epsilon_spent += step_epsilon;
        step_epsilon
    }

    /// Compute epsilon using moments accountant.
    fn moments_epsilon(&self, _num_clients: usize) -> f64 {
        // Simplified moments accountant
        let sigma = self.config.noise_multiplier;
        let q = self.sample_rate;

        // Basic composition: epsilon per step
        if sigma > 0.0 {
            q * (1.0 / sigma.powi(2))
        } else {
            f64::INFINITY
        }
    }

    /// Compute epsilon using Renyi differential privacy.
    fn rdp_epsilon(&self, _num_clients: usize) -> f64 {
        // Simplified RDP accountant
        let sigma = self.config.noise_multiplier;
        let q = self.sample_rate;
        let alpha = 2.0; // RDP order

        if sigma > 0.0 {
            // RDP bound with subsampling: q * alpha / (2 * sigma^2)
            let rdp = q * alpha / (2.0 * sigma.powi(2));
            // Convert to (epsilon, delta)-DP
            let delta = self.config.delta;
            rdp + (1.0 / delta).ln() / (alpha - 1.0)
        } else {
            f64::INFINITY
        }
    }

    /// Compute epsilon using Gaussian differential privacy.
    fn gdp_epsilon(&self, _num_clients: usize) -> f64 {
        // Simplified GDP accountant
        let sigma = self.config.noise_multiplier;

        if sigma > 0.0 {
            // Central limit theorem approximation
            1.0 / sigma
        } else {
            f64::INFINITY
        }
    }

    /// Compute epsilon using privacy loss distribution.
    fn pld_epsilon(&self, _num_clients: usize) -> f64 {
        // Simplified PLD accountant (falls back to RDP for simplicity)
        self.rdp_epsilon(_num_clients)
    }

    /// Get total epsilon spent.
    #[must_use]
    pub fn epsilon_spent(&self) -> f64 {
        self.epsilon_spent
    }

    /// Get number of steps.
    #[must_use]
    pub fn num_steps(&self) -> usize {
        self.num_steps
    }

    /// Check if privacy budget is exhausted.
    #[must_use]
    pub fn is_budget_exhausted(&self) -> bool {
        if let Some(target) = self.config.target_epsilon {
            self.epsilon_spent >= target
        } else {
            false
        }
    }

    /// Get remaining privacy budget.
    #[must_use]
    pub fn remaining_budget(&self) -> Option<f64> {
        self.config
            .target_epsilon
            .map(|target| (target - self.epsilon_spent).max(0.0))
    }

    /// Estimate number of remaining steps.
    #[must_use]
    pub fn estimate_remaining_steps(&self) -> Option<usize> {
        if self.num_steps == 0 || self.epsilon_spent <= 0.0 {
            return None;
        }

        let epsilon_per_step = self.epsilon_spent / self.num_steps as f64;
        self.remaining_budget()
            .map(|remaining| (remaining / epsilon_per_step) as usize)
    }

    /// Reset the accountant.
    pub fn reset(&mut self) {
        self.epsilon_spent = 0.0;
        self.num_steps = 0;
    }

    /// Get the target epsilon.
    #[must_use]
    pub fn target_epsilon(&self) -> Option<f64> {
        self.config.target_epsilon
    }

    /// Get delta.
    #[must_use]
    pub fn delta(&self) -> f64 {
        self.config.delta
    }
}

/// Secure noise generation for distributed settings.
#[derive(Debug)]
pub struct SecureNoiseGenerator {
    /// Seed for reproducibility.
    seed: u64,
    /// Number of parties.
    num_parties: usize,
}

impl SecureNoiseGenerator {
    /// Create a new secure noise generator.
    #[must_use]
    pub fn new(seed: u64, num_parties: usize) -> Self {
        Self { seed, num_parties }
    }

    /// Generate correlated noise that sums to Gaussian noise.
    /// Each party generates their share independently.
    #[must_use]
    pub fn generate_noise_share(&self, party_id: usize, size: usize, sigma: f32) -> Vec<f32> {
        use rand::SeedableRng;
        let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed.wrapping_add(party_id as u64));

        (0..size)
            .map(|_| {
                let u1: f32 = rng.r#gen::<f32>();
                let u2: f32 = rng.r#gen::<f32>();
                let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
                noise * sigma / (self.num_parties as f32).sqrt()
            })
            .collect()
    }

    /// Combine noise shares from all parties.
    #[must_use]
    pub fn combine_shares(&self, shares: &[Vec<f32>]) -> Vec<f32> {
        if shares.is_empty() {
            return vec![];
        }

        let size = shares[0].len();
        let mut combined = vec![0.0_f32; size];

        for share in shares {
            for (i, &s) in share.iter().enumerate() {
                if i < combined.len() {
                    combined[i] += s;
                }
            }
        }

        combined
    }
}

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

    fn sample_config() -> PrivacyConfig {
        PrivacyConfig {
            epsilon: 8.0,
            delta: 1e-5,
            clip_norm: 1.0,
            noise_multiplier: 1.1,
            local_dp: false,
            accountant: PrivacyAccountantType::RDP,
            target_epsilon: Some(10.0),
        }
    }

    #[test]
    fn test_differential_privacy() {
        let config = sample_config();
        let dp = DifferentialPrivacy::new(config);

        let gradients = vec![0.1, 0.2, 0.3, 0.4, 0.5];
        let noisy = dp.add_noise(&gradients);

        assert_eq!(noisy.len(), gradients.len());
        // Noisy gradients should be different from original
        assert!(
            noisy
                .iter()
                .zip(gradients.iter())
                .any(|(n, g)| (n - g).abs() > 0.0)
        );
    }

    #[test]
    fn test_local_dp() {
        let mut config = sample_config();
        config.local_dp = true;
        let dp = DifferentialPrivacy::new(config);

        let gradients = vec![0.1, 0.2, 0.3];
        let noisy = dp.add_local_noise(&gradients);

        assert_eq!(noisy.len(), gradients.len());
    }

    #[test]
    fn test_gradient_clipping() {
        let clipper = GradientClipper::new(1.0);

        // Large gradients should be clipped
        let large_grads = vec![10.0, 10.0, 10.0];
        let clipped = clipper.clip(&large_grads);

        let clipped_norm = clipper.l2_norm(&clipped);
        assert!(clipped_norm <= 1.01); // Allow small floating point error

        // Small gradients should not be affected
        let small_grads = vec![0.1, 0.1, 0.1];
        let not_clipped = clipper.clip(&small_grads);
        assert_eq!(not_clipped, small_grads);
    }

    #[test]
    fn test_adaptive_clipping() {
        let clipper = GradientClipper::new(1.0);

        let gradients = vec![0.1, 0.5, 1.0, 2.0, 5.0];
        let clipped = clipper.adaptive_clip(&gradients, 0.5);

        // Median value should be the threshold
        assert!(clipped.iter().all(|&g| g.abs() <= 1.01));
    }

    #[test]
    fn test_norms() {
        let clipper = GradientClipper::new(1.0);
        let gradients = vec![3.0, 4.0];

        assert!((clipper.l2_norm(&gradients) - 5.0).abs() < 0.001);
        assert!((clipper.l1_norm(&gradients) - 7.0).abs() < 0.001);
        assert!((clipper.linf_norm(&gradients) - 4.0).abs() < 0.001);
    }

    #[test]
    fn test_privacy_accountant() {
        let config = sample_config();
        let mut accountant = PrivacyAccountant::new(config);

        assert_eq!(accountant.epsilon_spent(), 0.0);
        assert_eq!(accountant.num_steps(), 0);

        let epsilon = accountant.step(5);
        assert!(epsilon > 0.0);
        assert_eq!(accountant.num_steps(), 1);
        assert!(accountant.epsilon_spent() > 0.0);
    }

    #[test]
    fn test_budget_exhaustion() {
        let mut config = sample_config();
        config.target_epsilon = Some(0.1);
        config.noise_multiplier = 0.1; // Low noise = high epsilon per step

        let mut accountant = PrivacyAccountant::new(config);

        // Take many steps until budget is exhausted
        for _ in 0..1000 {
            accountant.step(5);
            if accountant.is_budget_exhausted() {
                break;
            }
        }

        assert!(accountant.is_budget_exhausted());
        assert!(accountant.remaining_budget().unwrap() <= 0.0);
    }

    #[test]
    fn test_remaining_steps_estimate() {
        let config = sample_config();
        let mut accountant = PrivacyAccountant::new(config);

        // Take some steps
        for _ in 0..10 {
            accountant.step(5);
        }

        let remaining = accountant.estimate_remaining_steps();
        assert!(remaining.is_some());
    }

    #[test]
    fn test_accountant_reset() {
        let config = sample_config();
        let mut accountant = PrivacyAccountant::new(config);

        accountant.step(5);
        assert!(accountant.epsilon_spent() > 0.0);

        accountant.reset();
        assert_eq!(accountant.epsilon_spent(), 0.0);
        assert_eq!(accountant.num_steps(), 0);
    }

    #[test]
    fn test_secure_noise_generator() {
        let generator = SecureNoiseGenerator::new(42, 3);

        let shares: Vec<_> = (0..3)
            .map(|party| generator.generate_noise_share(party, 10, 1.0))
            .collect();

        let combined = generator.combine_shares(&shares);
        assert_eq!(combined.len(), 10);
    }

    #[test]
    fn test_different_accountant_types() {
        for accountant_type in [
            PrivacyAccountantType::Moments,
            PrivacyAccountantType::RDP,
            PrivacyAccountantType::GDP,
            PrivacyAccountantType::PLD,
        ] {
            let mut config = sample_config();
            config.accountant = accountant_type;

            let mut accountant = PrivacyAccountant::new(config);
            let epsilon = accountant.step(5);

            assert!(epsilon > 0.0);
            assert!(epsilon < f64::INFINITY);
        }
    }

    #[test]
    fn test_sample_rate() {
        let config = sample_config();
        let mut accountant_full = PrivacyAccountant::new(config.clone());
        let mut accountant_sub = PrivacyAccountant::with_sample_rate(config, 0.1);

        let eps_full = accountant_full.step(5);
        let eps_sub = accountant_sub.step(5);

        // Subsampled should have lower epsilon (privacy amplification)
        assert!(eps_sub < eps_full);
    }

    #[test]
    fn test_clipper_set_norm() {
        let mut clipper = GradientClipper::new(1.0);
        assert_eq!(clipper.clip_norm(), 1.0);

        clipper.set_clip_norm(2.0);
        assert_eq!(clipper.clip_norm(), 2.0);
    }

    #[test]
    fn test_dp_set_noise_multiplier() {
        let config = sample_config();
        let mut dp = DifferentialPrivacy::new(config);

        dp.set_noise_multiplier(2.0);
        assert_eq!(dp.config().noise_multiplier, 2.0);
    }
}