"""Vector perturbation helpers for DP unlearning."""

import numpy as np

from .base_mechanism import BasePrivacyMechanism


class VectorPerturbation:
    """Applies differential privacy perturbation to embedding vectors."""

    def __init__(self, mechanism: BasePrivacyMechanism):
        self.mechanism = mechanism

    def perturb_single_vector(self, vector: np.ndarray, strategy: str = "uniform") -> np.ndarray:
        if strategy == "uniform":
            return self._uniform_perturbation(vector)
        if strategy == "selective":
            return self._selective_perturbation(vector)
        if strategy == "adaptive":
            return self._adaptive_perturbation(vector)
        raise ValueError(f"Unknown strategy: {strategy}")

    def _uniform_perturbation(self, vector: np.ndarray) -> np.ndarray:
        perturbed = vector.copy()
        for i in range(len(vector)):
            perturbed[i] += self.mechanism.compute_noise()
        return perturbed

    def _selective_perturbation(self, vector: np.ndarray) -> np.ndarray:
        perturbed = vector.copy()
        dim = len(vector)
        selected_dims = np.random.choice(dim, size=max(1, dim // 2), replace=False)
        for i in selected_dims:
            perturbed[i] += self.mechanism.compute_noise()
        return perturbed

    def _adaptive_perturbation(self, vector: np.ndarray) -> np.ndarray:
        perturbed = vector.copy()
        magnitudes = np.abs(vector)
        total_magnitude = np.sum(magnitudes) + 1e-10
        for i in range(len(vector)):
            importance = magnitudes[i] / total_magnitude
            adaptive_sensitivity = self.mechanism.sensitivity * (1 - importance)
            perturbed[i] += self.mechanism.compute_noise(sensitivity=adaptive_sensitivity)
        return perturbed

    def perturb_batch_vectors(self, vectors: np.ndarray, strategy: str = "uniform") -> np.ndarray:
        perturbed_batch = np.zeros_like(vectors)
        for i, vector in enumerate(vectors):
            perturbed_batch[i] = self.perturb_single_vector(vector, strategy)
        return perturbed_batch

    def measure_perturbation_impact(
        self, original: np.ndarray, perturbed: np.ndarray
    ) -> dict[str, float]:
        l2_distance = np.linalg.norm(original - perturbed)
        l1_distance = np.sum(np.abs(original - perturbed))
        cosine_similarity = np.dot(original, perturbed) / (
            np.linalg.norm(original) * np.linalg.norm(perturbed) + 1e-10
        )
        return {
            "l2_distance": float(l2_distance),
            "l1_distance": float(l1_distance),
            "cosine_similarity": float(cosine_similarity),
            "relative_change": float(l2_distance / (np.linalg.norm(original) + 1e-10)),
        }
