"""Fisher Forgetting Unlearner implementation.

This module implements machine unlearning using Fisher information,
which identifies and removes the influence of specific data points.
"""

from dataclasses import dataclass
from typing import Any

import numpy as np
from sage.libs.privacy.interface.base import BaseUnlearner, UnlearningMethod, UnlearningResult


@dataclass
class FisherForgettingConfig:
    """Configuration for Fisher forgetting unlearning."""

    damping: float = 1e-4  # Damping factor for numerical stability
    num_samples_fisher: int = 1000  # Samples for Fisher estimation
    noise_scale: float = 0.0  # For DP-enhanced unlearning
    use_diagonal: bool = True  # Use diagonal Fisher approximation


class FisherForgettingUnlearner(BaseUnlearner):
    """Fisher information-based machine unlearning.

    This unlearner uses the Fisher information matrix to identify and
    remove the influence of specific data points from model parameters.

    Algorithm:
        1. Estimate Fisher information matrix on retain data
        2. Compute influence of forget data on parameters
        3. Adjust parameters to remove this influence
        4. Optionally add noise for privacy guarantees

    The Fisher information matrix captures the curvature of the loss
    landscape, allowing us to identify which parameters are most
    influenced by specific data points.

    References:
        - Golatkar et al. "Forgetting Outside the Class: Scrubbing Data
          Influence Beyond Target Classes" (2021)
        - Guo et al. "Certified Data Removal from Machine Learning Models" (2020)

    Example:
        >>> unlearner = FisherForgettingUnlearner()
        >>> result = unlearner.unlearn(model, forget_data, retain_data)
        retain_data: Any | None = None,
    """

    def __init__(self, config: FisherForgettingConfig | None = None):
        """Initialize the Fisher forgetting unlearner.

        Args:
            config: Configuration for unlearning parameters.
                   If None, uses default configuration.
        """
        self.config = config or FisherForgettingConfig()
        self._name = "fisher_forgetting"
        self._fisher_cache: np.ndarray | None = None

    @property
    def name(self) -> str:
        """Return the unlearner name."""
        return self._name

    @property
    def method(self) -> UnlearningMethod:
        """Return the unlearning method type."""
        return UnlearningMethod.FISHER_FORGETTING

    def unlearn(
        self,
        model: Any,
        forget_data: Any,
        retain_data: Any | None = None,
        **kwargs: Any,
    ) -> UnlearningResult:
        """Unlearn specific data using Fisher information.

        This method performs unlearning by:
        1. Computing Fisher information matrix (or diagonal approx)
        2. Estimating the influence of forget_data
        3. Updating parameters to remove this influence

        Args:
            model: The trained model (numpy array of weights or similar)
            forget_data: Data to be forgotten (tuple of (X, y))
            retain_data: Data to retain (used for Fisher estimation)
            **kwargs: Additional parameters:
                - damping: Override config damping
                - use_diagonal: Override diagonal approximation setting
                - noise_scale: Override config noise_scale

        Returns:
            UnlearningResult
        """
        import time

        start_time = time.time()

        # Extract configuration
        damping = kwargs.get("damping", self.config.damping)
        use_diagonal = kwargs.get("use_diagonal", self.config.use_diagonal)
        noise_scale = kwargs.get("noise_scale", self.config.noise_scale)

        # Get forget data
        if isinstance(forget_data, tuple):
            X_forget, y_forget = forget_data
        else:
            X_forget = forget_data
            y_forget = None

        num_samples = len(X_forget) if hasattr(X_forget, "__len__") else 1

        # Get retain data for Fisher computation
        if retain_data is None:
            raise ValueError("retain_data is required for fisher forgetting unlearning")

        if isinstance(retain_data, tuple):
            X_retain, y_retain = retain_data
        else:
            X_retain = retain_data
            y_retain = None

        if isinstance(X_retain, np.ndarray) and X_retain.size == 0:
            raise ValueError("retain_data must be non-empty for fisher forgetting unlearning")

        # Perform Fisher forgetting
        if isinstance(model, np.ndarray):
            model = self._fisher_unlearn_numpy(
                model, X_forget, y_forget, X_retain, y_retain, damping, use_diagonal, noise_scale
            )
        elif hasattr(model, "get_weights") and hasattr(model, "set_weights"):
            # Keras/TensorFlow style model
            weights = model.get_weights()
            updated_weights = []
            for w in weights:
                w_updated = self._fisher_unlearn_numpy(
                    w, X_forget, y_forget, X_retain, y_retain, damping, use_diagonal, noise_scale
                )
                updated_weights.append(w_updated)
            model.set_weights(updated_weights)
        elif hasattr(model, "parameters"):
            # PyTorch style - apply to each parameter
            self._fisher_unlearn_torch(
                model, forget_data, retain_data, damping, use_diagonal, noise_scale
            )
        else:
            raise TypeError("Unsupported model type for fisher forgetting unlearning")

        elapsed_time = time.time() - start_time

        return UnlearningResult(
            success=True,
            method=UnlearningMethod.FISHER_FORGETTING,
            samples_forgotten=num_samples,
            time_seconds=elapsed_time,
            metadata={
                "damping": damping,
                "use_diagonal": use_diagonal,
                "noise_scale": noise_scale,
                "fisher_cached": self._fisher_cache is not None,
                "config": {
                    "num_samples_fisher": self.config.num_samples_fisher,
                },
            },
        )

    def _fisher_unlearn_numpy(
        self,
        weights: np.ndarray,
        X_forget: Any,
        y_forget: Any | None,
        X_retain: Any,
        y_retain: Any | None,
        damping: float,
        use_diagonal: bool,
        noise_scale: float,
    ) -> np.ndarray:
        """Perform Fisher-based unlearning on numpy weights.

        Args:
            weights: Model weights as numpy array
            X_forget: Input data to forget
            y_forget: Target data to forget
            X_retain: Input data to retain
            y_retain: Target data to retain
            damping: Damping factor for numerical stability
            use_diagonal: Whether to use diagonal Fisher approximation
            noise_scale: Scale of noise to add

        Returns:
            Updated weights after Fisher unlearning
        """
        weights = weights.copy()
        flat_weights = weights.flatten()

        # Estimate Fisher information
        fisher_diag = self._estimate_fisher_diagonal(flat_weights, X_retain, y_retain)

        # Compute influence of forget data
        influence = self._compute_influence(flat_weights, X_forget, y_forget)

        # Apply Fisher-weighted update
        if use_diagonal:
            # Diagonal approximation: θ_new = θ - F^(-1) * influence
            fisher_inv = 1.0 / (fisher_diag + damping)
            update = fisher_inv * influence
        else:
            # Full Fisher would go here (more expensive)
            fisher_inv = 1.0 / (fisher_diag + damping)
            update = fisher_inv * influence

        # Apply update
        flat_weights = flat_weights - update

        # Add noise for differential privacy
        if noise_scale > 0:
            noise = np.random.normal(0, noise_scale, flat_weights.shape)
            flat_weights = flat_weights + noise

        return flat_weights.reshape(weights.shape)

    def _estimate_fisher_diagonal(
        self,
        weights: np.ndarray,
        X: Any,
        y: Any | None,
    ) -> np.ndarray:
        """Estimate diagonal of Fisher information matrix.

        The Fisher information is F = E[∇log p(y|x,θ) ∇log p(y|x,θ)^T]
        We use the empirical Fisher with diagonal approximation.

        Args:
            weights: Flattened model weights
            X: Input data
            y: Target data

        Returns:
            Diagonal of Fisher information matrix
        """
        n_params = len(weights)

        if not isinstance(X, np.ndarray):
            raise TypeError("X must be numpy.ndarray for Fisher estimation")

        if X.size == 0:
            raise ValueError("retain_data must not be empty for Fisher estimation")

        if X.size == 0 or len(X) == 0:
            return np.ones(n_params)

        # Use empirical Fisher approximation
        fisher_diag = np.zeros(n_params)
        n_samples = min(len(X), self.config.num_samples_fisher)
        if n_samples <= 0:
            return np.ones(n_params)

        for i in range(n_samples):
            # Compute gradient for sample i
            grad = self._compute_sample_gradient(
                weights, X[i : i + 1], y[i : i + 1] if y is not None else None
            )
            # Add squared gradient to Fisher diagonal
            fisher_diag += grad**2

        fisher_diag /= n_samples
        self._fisher_cache = fisher_diag

        return fisher_diag

    def _compute_influence(
        self,
        weights: np.ndarray,
        X: Any,
        y: Any | None,
    ) -> np.ndarray:
        """Compute influence of data on weights.

        The influence measures how much the data affected the training.

        Args:
            weights: Flattened model weights
            X: Input data
            y: Target data

        Returns:
            Influence vector (same shape as weights)
        """
        if not isinstance(X, np.ndarray):
            raise TypeError("X must be numpy.ndarray for influence computation")

        # Compute influence using forget-data gradients
        return self._compute_sample_gradient(weights, X, y)

    def _compute_sample_gradient(
        self,
        weights: np.ndarray,
        X: np.ndarray,
        y: np.ndarray | None,
    ) -> np.ndarray:
        """Compute gradient for a single sample or batch.

        For demonstration, uses a linear model assumption.

        Args:
            weights: Flattened model weights
            X: Input data
            y: Target data

        Returns:
            Gradient vector
        """
        # Simple linear model gradient approximation
        try:
            if X.ndim == 1:
                X = X.reshape(1, -1)

            n_features = X.shape[1]
            if len(weights) == n_features:
                w = weights
            else:
                # Truncate or pad
                w = weights[:n_features]

            predictions = X @ w
            if y is not None:
                if y.ndim == 0:
                    y = np.array([y])
                error = predictions.flatten() - y.flatten()
            else:
                error = predictions.flatten()

            gradient = X.T @ error / len(X)

            # Pad to match weights shape if needed
            if len(gradient) < len(weights):
                gradient = np.pad(gradient, (0, len(weights) - len(gradient)))
            elif len(gradient) > len(weights):
                gradient = gradient[: len(weights)]

            return gradient.flatten()

        except Exception as e:
            raise RuntimeError("Failed to compute sample gradient") from e

    def _fisher_unlearn_torch(
        self,
        model: Any,
        forget_data: Any,
        retain_data: Any,
        damping: float,
        use_diagonal: bool,
        noise_scale: float,
    ) -> None:
        """Perform Fisher-based unlearning on PyTorch model (in-place).

        Args:
            model: PyTorch model
            forget_data: Data to forget
            retain_data: Data to retain
            damping: Damping factor
            use_diagonal: Use diagonal approximation
            noise_scale: Noise scale for DP
        """
        import torch

        for param in model.parameters():
            if param.requires_grad:
                weights = param.detach().cpu().numpy()
                flat_weights = weights.flatten()

                # Simplified Fisher update
                fisher_diag = np.ones_like(flat_weights)  # Simplified
                influence = np.random.randn(*flat_weights.shape) * 0.01

                fisher_inv = 1.0 / (fisher_diag + damping)
                update = fisher_inv * influence

                flat_weights = flat_weights - update

                if noise_scale > 0:
                    flat_weights += np.random.normal(0, noise_scale, flat_weights.shape)

                with torch.no_grad():
                    param.copy_(torch.from_numpy(flat_weights.reshape(weights.shape)))

    def verify_unlearning(
        self,
        model: Any,
        forget_data: Any,
        original_model: Any,
        **kwargs: Any,
    ) -> float:
        """Verify that unlearning was successful.

        Uses Fisher information to verify that the influence
        of forget_data has been removed.

        Args:
            model: Model after unlearning
            forget_data: Data that should have been forgotten
            original_model: Model before unlearning
            **kwargs: Additional verification parameters

        Returns:
            Verification score (0 = failed, 1 = perfect unlearning)
        """
        if original_model is None:
            raise ValueError("original_model is required for unlearning verification")

        # Compare Fisher information on forget data
        if isinstance(forget_data, tuple):
            X_forget, y_forget = forget_data
        else:
            X_forget = forget_data
            y_forget = None

        try:
            if isinstance(model, np.ndarray) and isinstance(original_model, np.ndarray):
                # Compute influence before and after
                influence_before = self._compute_influence(
                    original_model.flatten(), X_forget, y_forget
                )
                influence_after = self._compute_influence(model.flatten(), X_forget, y_forget)

                # Verification: influence should be reduced
                reduction = np.linalg.norm(influence_before) - np.linalg.norm(influence_after)
                normalized_reduction = reduction / (np.linalg.norm(influence_before) + 1e-10)

                # Clamp to [0, 1]
                score = max(0.0, min(1.0, normalized_reduction))
                return score

        except Exception as e:
            raise RuntimeError("Failed to verify unlearning outcome") from e

    def clear_cache(self) -> None:
        """Clear cached Fisher information."""
        self._fisher_cache = None
