"""Gradient Ascent Unlearner implementation.

This module implements approximate machine unlearning via gradient ascent,
which reverses the learning process on 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 GradientAscentConfig:
    """Configuration for gradient ascent unlearning."""

    learning_rate: float = 0.01
    num_steps: int = 100
    batch_size: int = 32
    noise_scale: float = 0.0  # For DP-enhanced unlearning
    clip_norm: float = 1.0


class GradientAscentUnlearner(BaseUnlearner):
    """Gradient ascent-based approximate unlearning.

    This unlearner performs approximate unlearning by running gradient ascent
    (the reverse of gradient descent) on the data to be forgotten.

    Algorithm:
        1. For each batch of forget data:
           - Compute loss on forget data
           - Compute gradients with respect to model parameters
           - Update parameters by ADDING gradients (opposite of descent)
        2. Optionally add noise for differential privacy

    References:
        - Thudi et al. "Unrolling SGD: Understanding Factors Influencing
          Machine Unlearning" (2022)
        - Golatkar et al. "Eternal Sunshine of the Spotless Net" (2020)

    Example:
        >>> unlearner = GradientAscentUnlearner()
        >>> result = unlearner.unlearn(model, forget_data)
        >>> print(f"Unlearned {result.samples_forgotten} samples")
    """

    def __init__(self, config: GradientAscentConfig | None = None):
        """Initialize the gradient ascent unlearner.

        Args:
            config: Configuration for unlearning parameters.
                   If None, uses default configuration.
        """
        self.config = config or GradientAscentConfig()
        self._name = "gradient_ascent"

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

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

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

        This method performs approximate unlearning by:
        1. Computing gradients on forget_data
        2. Performing gradient ascent to "reverse" learning
        3. Optionally adding noise for privacy guarantees

        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 for verification (optional)
            **kwargs: Additional parameters:
                - learning_rate: Override config learning rate
                - num_steps: Override config num_steps
                - noise_scale: Override config noise_scale

        Returns:
            UnlearningResult
        """
        import time

        start_time = time.time()

        # Extract configuration
        lr = kwargs.get("learning_rate", self.config.learning_rate)
        num_steps = kwargs.get("num_steps", self.config.num_steps)
        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:
            # Assume it's already a proper data structure
            X_forget = forget_data
            y_forget = None

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

        # Perform gradient ascent
        if isinstance(model, np.ndarray):
            # Simple case: model is a weight matrix
            model = self._gradient_ascent_numpy(
                model, X_forget, y_forget, lr, num_steps, 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._gradient_ascent_numpy(
                    w, X_forget, y_forget, lr, num_steps, noise_scale
                )
                updated_weights.append(w_updated)
            model.set_weights(updated_weights)
        elif hasattr(model, "parameters"):
            # PyTorch style model
            self._gradient_ascent_torch(model, X_forget, y_forget, lr, num_steps, noise_scale)
        else:
            raise TypeError("Unsupported model type for gradient ascent unlearning")

        elapsed_time = time.time() - start_time

        return UnlearningResult(
            success=True,
            method=UnlearningMethod.GRADIENT_ASCENT,
            samples_forgotten=num_samples,
            time_seconds=elapsed_time,
            metadata={
                "learning_rate": lr,
                "num_steps": num_steps,
                "noise_scale": noise_scale,
                "config": {
                    "batch_size": self.config.batch_size,
                    "clip_norm": self.config.clip_norm,
                },
            },
        )
    def _gradient_ascent_numpy(
        self,
        weights: np.ndarray,
        X: Any,
        y: Any | None,
        lr: float,
        num_steps: int,
        noise_scale: float,
    ) -> np.ndarray:
        """Perform gradient ascent on numpy weights.

        Args:
            weights: Model weights as numpy array
            X: Input data
            y: Target data (optional)
            lr: Learning rate
            num_steps: Number of gradient steps
            noise_scale: Scale of noise to add

        Returns:
            Updated weights after gradient ascent
        """
        weights = weights.copy()

        for _step in range(num_steps):
            # Compute pseudo-gradient (simplified)
            if y is not None and isinstance(X, np.ndarray):
                # Simple gradient approximation
                gradient = self._compute_simple_gradient(weights, X, y)
            elif isinstance(X, np.ndarray):
                # Label-free approximation for unlabeled forget_data
                gradient = np.mean(X, axis=0)
                if gradient.shape != weights.shape:
                    gradient = np.resize(gradient, weights.shape)
            else:
                raise TypeError("forget_data must be a numpy.ndarray or tuple(X, y)")

            # Clip gradient
            grad_norm = np.linalg.norm(gradient)
            if grad_norm > self.config.clip_norm:
                gradient = gradient * (self.config.clip_norm / grad_norm)

            # Gradient ASCENT (add instead of subtract)
            weights = weights + lr * gradient

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

        return weights

    def _compute_simple_gradient(
        self, weights: np.ndarray, X: np.ndarray, y: np.ndarray
    ) -> np.ndarray:
        """Compute a simple gradient approximation.

        For demonstration, uses a linear model assumption.
        In practice, this should use the actual model's gradient computation.

        Args:
            weights: Current weights
            X: Input data
            y: Target data

        Returns:
            Approximated gradient
        """
        # Simple linear model gradient: 2 * X^T (Xw - y) / n
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        if weights.ndim == 1:
            w = weights.reshape(-1, 1)
        else:
            w = weights

        n = X.shape[0]
        try:
            predictions = X @ w
            if y.ndim == 1:
                y = y.reshape(-1, 1)
            error = predictions - y
            gradient = (2 / n) * (X.T @ error)
            return gradient.reshape(weights.shape)
        except Exception as e:
            raise RuntimeError("Failed to compute gradient for numpy model") from e

    def _gradient_ascent_torch(
        self,
        model: Any,
        X: Any,
        y: Any | None,
        lr: float,
        num_steps: int,
        noise_scale: float,
    ) -> None:
        """Perform gradient ascent on PyTorch model (in-place).

        Args:
            model: PyTorch model
            X: Input data
            y: Target data
            lr: Learning rate
            num_steps: Number of steps
            noise_scale: Noise scale for DP
        """
        import torch

        for _step in range(num_steps):
            for param in model.parameters():
                if param.grad is not None:
                    # Gradient ascent: add gradient
                    with torch.no_grad():
                        param.add_(lr * param.grad)
                        if noise_scale > 0:
                            param.add_(torch.randn_like(param) * noise_scale)

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

        Uses membership inference attack success rate to verify.
        Lower attack success = better unlearning.

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

        Returns:
            Verification score (0 = failed, 1 = perfect unlearning)
        """
        # Simple verification: compare loss on forget data
        if original_model is None:
            raise ValueError("original_model is required for unlearning verification")

        # If we have both models, compare behavior on forget data
        if isinstance(forget_data, tuple):
            X_forget, _y_forget = forget_data
        else:
            _X_forget = forget_data

        try:
            # Compare predictions
            if isinstance(model, np.ndarray) and isinstance(original_model, np.ndarray):
                diff = np.linalg.norm(model - original_model)
                # Normalize to [0, 1]
                score = min(1.0, diff / (np.linalg.norm(original_model) + 1e-10))
                return score
        except Exception as e:
            raise RuntimeError("Failed to verify unlearning outcome") from e
