"""Neighbor compensation strategies for DP unlearning."""

import numpy as np


class NeighborCompensation:
    """Compensates neighboring vectors to prevent collateral unlearning."""

    def __init__(self, similarity_threshold: float = 0.8, max_neighbors: int = 10):
        self.similarity_threshold = similarity_threshold
        self.max_neighbors = max_neighbors

    def identify_neighbors(
        self, target_vector: np.ndarray, all_vectors: np.ndarray, all_ids: list[str]
    ) -> list[tuple[str, float]]:
        similarities = self._compute_cosine_similarities(target_vector, all_vectors)
        neighbors = [
            (all_ids[i], float(sim))
            for i, sim in enumerate(similarities)
            if sim >= self.similarity_threshold
        ]
        neighbors.sort(key=lambda item: item[1], reverse=True)
        return neighbors[: self.max_neighbors]

    def compute_compensation(
        self,
        original_vector: np.ndarray,
        perturbed_vector: np.ndarray,
        neighbor_vector: np.ndarray,
        neighbor_similarity: float,
    ) -> np.ndarray:
        perturbation = perturbed_vector - original_vector
        return -neighbor_similarity * perturbation

    def _compute_cosine_similarities(self, query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
        query_norm = query / (np.linalg.norm(query) + 1e-10)
        vectors_norm = vectors / (np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-10)
        return np.dot(vectors_norm, query_norm)

    def apply_compensation(
        self,
        original_vector: np.ndarray,
        perturbed_vector: np.ndarray,
        all_vectors: np.ndarray,
        all_ids: list[str],
    ) -> dict[str, np.ndarray]:
        compensated_vectors: dict[str, np.ndarray] = {}
        for neighbor_id, similarity in self.identify_neighbors(
            original_vector, all_vectors, all_ids
        ):
            neighbor_idx = all_ids.index(neighbor_id)
            neighbor_vector = all_vectors[neighbor_idx]
            compensation = self.compute_compensation(
                original_vector,
                perturbed_vector,
                neighbor_vector,
                similarity,
            )
            compensated_vectors[neighbor_id] = neighbor_vector + compensation
        return compensated_vectors
