"""Tests for differential privacy mechanisms."""

import numpy as np
import pytest

from sage_libs.sage_privacy.dp import GaussianMechanism, LaplaceMechanism


class TestGaussianMechanism:
    """Tests for GaussianMechanism."""

    def test_init_default(self):
        """Test initialization with default parameters."""
        mechanism = GaussianMechanism(epsilon=1.0)
        assert mechanism.epsilon == 1.0
        assert mechanism.delta == 1e-5
        assert mechanism.sensitivity == 1.0
        assert mechanism.name == "Gaussian"

    def test_init_custom(self):
        """Test initialization with custom parameters."""
        mechanism = GaussianMechanism(
            epsilon=0.5, delta=1e-6, sensitivity=2.0, name="CustomGaussian"
        )
        assert mechanism.epsilon == 0.5
        assert mechanism.delta == 1e-6
        assert mechanism.sensitivity == 2.0
        assert mechanism.name == "CustomGaussian"

    def test_init_invalid_epsilon(self):
        """Test that invalid epsilon raises error."""
        with pytest.raises(ValueError):
            GaussianMechanism(epsilon=0)
        with pytest.raises(ValueError):
            GaussianMechanism(epsilon=-1)

    def test_init_invalid_delta(self):
        """Test that invalid delta raises error."""
        with pytest.raises(ValueError):
            GaussianMechanism(epsilon=1.0, delta=0)
        with pytest.raises(ValueError):
            GaussianMechanism(epsilon=1.0, delta=1.0)
        with pytest.raises(ValueError):
            GaussianMechanism(epsilon=1.0, delta=1.5)

    def test_sigma_computation(self):
        """Test that sigma is computed correctly."""
        mechanism = GaussianMechanism(epsilon=1.0, delta=1e-5, sensitivity=1.0)
        # σ = Δf * sqrt(2 * ln(1.25/δ)) / ε
        expected_sigma = np.sqrt(2 * np.log(1.25 / 1e-5)) / 1.0
        assert np.isclose(mechanism.sigma, expected_sigma)

    def test_compute_noise(self):
        """Test noise computation."""
        mechanism = GaussianMechanism(epsilon=1.0)
        noise = mechanism.compute_noise()
        assert isinstance(noise, float)

    def test_compute_noise_vector(self):
        """Test vector noise computation."""
        mechanism = GaussianMechanism(epsilon=1.0)
        shape = (10, 5)
        noise = mechanism.compute_noise_vector(shape)
        assert noise.shape == shape

    def test_perturb_vector(self):
        """Test vector perturbation."""
        mechanism = GaussianMechanism(epsilon=1.0)
        vector = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        perturbed = mechanism.perturb_vector(vector)

        assert perturbed.shape == vector.shape
        assert not np.allclose(perturbed, vector)  # Should be different due to noise

    def test_perturb_vector_specific_indices(self):
        """Test perturbation of specific indices."""
        mechanism = GaussianMechanism(epsilon=1.0)
        vector = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        original = vector.copy()
        perturbed = mechanism.perturb_vector(vector, indices_to_perturb=[0, 2])

        # Indices 0 and 2 should be different
        assert not np.isclose(perturbed[0], original[0]) or not np.isclose(
            perturbed[2], original[2]
        )
        # Indices 1, 3, 4 should be same
        assert perturbed[1] == original[1]
        assert perturbed[3] == original[3]
        assert perturbed[4] == original[4]

    def test_perturb_matrix(self):
        """Test matrix perturbation."""
        mechanism = GaussianMechanism(epsilon=1.0)
        matrix = np.random.randn(5, 3)
        perturbed = mechanism.perturb_matrix(matrix)

        assert perturbed.shape == matrix.shape
        assert not np.allclose(perturbed, matrix)

    def test_privacy_cost(self):
        """Test privacy budget tracking."""
        mechanism = GaussianMechanism(epsilon=1.0, delta=1e-5)
        mechanism.reset_privacy_budget()

        # Initial cost should be 0
        eps, dlt = mechanism.privacy_cost()
        assert eps == 0.0
        assert dlt == 0.0

        # After one query
        mechanism.compute_noise()
        eps, dlt = mechanism.privacy_cost()
        assert eps == 1.0
        assert dlt == 1e-5

        # After multiple queries (basic composition)
        mechanism.compute_noise()
        mechanism.compute_noise()
        eps, dlt = mechanism.privacy_cost()
        assert eps == 3.0  # 3 queries

    def test_reset_privacy_budget(self):
        """Test resetting privacy budget."""
        mechanism = GaussianMechanism(epsilon=1.0)
        mechanism.compute_noise()
        mechanism.compute_noise()

        mechanism.reset_privacy_budget()
        eps, dlt = mechanism.privacy_cost()
        assert eps == 0.0

    def test_get_privacy_guarantee(self):
        """Test getting privacy guarantee."""
        mechanism = GaussianMechanism(epsilon=0.5, delta=1e-6)
        guarantee = mechanism.get_privacy_guarantee()

        assert guarantee["epsilon"] == 0.5
        assert guarantee["delta"] == 1e-6
        assert "sigma" in guarantee

    def test_repr(self):
        """Test string representation."""
        mechanism = GaussianMechanism(epsilon=1.0, delta=1e-5)
        repr_str = repr(mechanism)
        assert "Gaussian" in repr_str
        assert "ε=1.0" in repr_str
        assert "δ=1e-05" in repr_str


class TestLaplaceMechanism:
    """Tests for LaplaceMechanism."""

    def test_init_default(self):
        """Test initialization with default parameters."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        assert mechanism.epsilon == 1.0
        assert mechanism.delta is None
        assert mechanism.sensitivity == 1.0
        assert mechanism.name == "Laplace"

    def test_init_custom(self):
        """Test initialization with custom parameters."""
        mechanism = LaplaceMechanism(epsilon=2.0, sensitivity=0.5, name="CustomLaplace")
        assert mechanism.epsilon == 2.0
        assert mechanism.sensitivity == 0.5
        assert mechanism.name == "CustomLaplace"

    def test_init_invalid_epsilon(self):
        """Test that invalid epsilon raises error."""
        with pytest.raises(ValueError):
            LaplaceMechanism(epsilon=0)
        with pytest.raises(ValueError):
            LaplaceMechanism(epsilon=-0.5)

    def test_scale_computation(self):
        """Test that scale is computed correctly."""
        mechanism = LaplaceMechanism(epsilon=2.0, sensitivity=1.0)
        # b = Δf / ε
        expected_scale = 1.0 / 2.0
        assert np.isclose(mechanism.scale, expected_scale)

    def test_compute_noise(self):
        """Test noise computation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        noise = mechanism.compute_noise()
        assert isinstance(noise, float)

    def test_compute_noise_vector(self):
        """Test vector noise computation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        shape = (10, 5)
        noise = mechanism.compute_noise_vector(shape)
        assert noise.shape == shape

    def test_perturb_vector(self):
        """Test vector perturbation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        vector = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        perturbed = mechanism.perturb_vector(vector)

        assert perturbed.shape == vector.shape
        assert not np.allclose(perturbed, vector)

    def test_perturb_histogram(self):
        """Test histogram perturbation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        histogram = np.array([10, 20, 30, 15, 25])
        perturbed = mechanism.perturb_histogram(histogram)

        assert perturbed.shape == histogram.shape
        assert np.all(perturbed >= 0)  # Non-negative

    def test_perturb_histogram_normalized(self):
        """Test normalized histogram perturbation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        histogram = np.array([10, 20, 30, 15, 25])
        perturbed = mechanism.perturb_histogram(histogram, normalize=True)

        assert perturbed.shape == histogram.shape
        assert np.isclose(perturbed.sum(), 1.0)

    def test_private_sum(self):
        """Test private sum computation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        values = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        result = mechanism.private_sum(values)

        assert isinstance(result, float)
        # With high probability, sum should be close to true sum (15)

    def test_private_sum_with_bounds(self):
        """Test private sum with clipping."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        values = np.array([1.0, 2.0, 100.0, 4.0, 5.0])  # 100 is outlier
        result = mechanism.private_sum(values, bounds=(0, 10))

        # With clipping, max true sum would be 5 * 10 = 50
        # Original sum without clipping would be 112
        assert isinstance(result, float)

    def test_private_mean(self):
        """Test private mean computation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        values = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        private_mean = mechanism.private_mean(values)

        assert isinstance(private_mean, float)

    def test_private_count(self):
        """Test private count computation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        predicate = np.array([True, True, False, True, False])
        private_count = mechanism.private_count(predicate)

        assert private_count >= 0  # Non-negative

    def test_privacy_cost_pure_dp(self):
        """Test that Laplace provides pure DP (delta=0)."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        mechanism.reset_privacy_budget()

        mechanism.compute_noise()
        eps, dlt = mechanism.privacy_cost()

        assert eps == 1.0
        assert dlt == 0.0  # Pure DP

    def test_get_privacy_guarantee(self):
        """Test getting privacy guarantee."""
        mechanism = LaplaceMechanism(epsilon=2.0)
        guarantee = mechanism.get_privacy_guarantee()

        assert guarantee["epsilon"] == 2.0
        assert guarantee["delta"] == 0.0  # Pure DP
        assert "scale" in guarantee

    def test_repr(self):
        """Test string representation."""
        mechanism = LaplaceMechanism(epsilon=1.0)
        repr_str = repr(mechanism)
        assert "Laplace" in repr_str
        assert "ε=1.0" in repr_str


class TestMechanismComparison:
    """Tests comparing Gaussian and Laplace mechanisms."""

    def test_noise_distributions(self):
        """Test that mechanisms produce different noise distributions."""
        gaussian = GaussianMechanism(epsilon=1.0)
        laplace = LaplaceMechanism(epsilon=1.0)

        # Generate many noise samples
        n_samples = 1000
        gaussian_noise = [gaussian.compute_noise() for _ in range(n_samples)]
        laplace_noise = [laplace.compute_noise() for _ in range(n_samples)]

        # Both should have mean around 0
        assert abs(np.mean(gaussian_noise)) < 0.5
        assert abs(np.mean(laplace_noise)) < 0.5

        # Laplace should have heavier tails (higher kurtosis)
        gaussian_kurtosis = np.mean((np.array(gaussian_noise) / np.std(gaussian_noise)) ** 4)
        laplace_kurtosis = np.mean((np.array(laplace_noise) / np.std(laplace_noise)) ** 4)
        # Gaussian kurtosis ~ 3, Laplace kurtosis ~ 6
        assert laplace_kurtosis > gaussian_kurtosis

    def test_utility_comparison(self):
        """Test utility comparison for same epsilon."""
        # For same epsilon, Gaussian with small delta should give better utility
        gaussian = GaussianMechanism(epsilon=1.0, delta=1e-5)
        laplace = LaplaceMechanism(epsilon=1.0)

        # Gaussian sigma for (1, 1e-5)-DP
        # Laplace scale for 1-DP
        # For many queries, Gaussian provides better composition
        assert gaussian.sigma > 0
        assert laplace.scale > 0
