"""Tests for unlearning implementations."""

import numpy as np
import pytest

from sage_libs.sage_privacy.unlearning import FisherForgettingUnlearner, GradientAscentUnlearner
from sage_libs.sage_privacy.unlearning.fisher_forgetting import FisherForgettingConfig
from sage_libs.sage_privacy.unlearning.gradient_ascent import GradientAscentConfig


def get_result_value(result, key: str):
    """Get value from result, handling both dict and dataclass."""
    if isinstance(result, dict):
        return result.get(key)
    return getattr(result, key, None)


class TestGradientAscentUnlearner:
    """Tests for GradientAscentUnlearner."""

    def test_init_default(self):
        """Test initialization with default config."""
        unlearner = GradientAscentUnlearner()
        assert unlearner.name == "gradient_ascent"
        assert unlearner.config.learning_rate == 0.01
        assert unlearner.config.num_steps == 100

    def test_init_custom_config(self):
        """Test initialization with custom config."""
        config = GradientAscentConfig(
            learning_rate=0.1,
            num_steps=50,
            batch_size=64,
            noise_scale=0.01,
        )
        unlearner = GradientAscentUnlearner(config=config)
        assert unlearner.config.learning_rate == 0.1
        assert unlearner.config.num_steps == 50

    def test_unlearn_numpy_weights(self):
        """Test unlearning on numpy weight matrix."""
        unlearner = GradientAscentUnlearner()

        # Create simple model weights
        weights = np.random.randn(10, 5)

        # Create forget data
        X_forget = np.random.randn(20, 10)
        y_forget = np.random.randn(20, 5)

        result = unlearner.unlearn(weights, (X_forget, y_forget), num_steps=10, learning_rate=0.01)

        assert get_result_value(result, "success") is True
        assert get_result_value(result, "samples_forgotten") == 20
        assert get_result_value(result, "time_seconds") is not None
        assert get_result_value(result, "metadata") is not None

    def test_unlearn_with_retain_data(self):
        """Test unlearning with retain data."""
        unlearner = GradientAscentUnlearner()

        weights = np.random.randn(5)
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)
        X_retain = np.random.randn(100, 5)
        y_retain = np.random.randn(100)

        result = unlearner.unlearn(
            weights, (X_forget, y_forget), retain_data=(X_retain, y_retain), num_steps=5
        )

        assert get_result_value(result, "success") is True

    def test_unlearn_with_noise(self):
        """Test unlearning with differential privacy noise."""
        unlearner = GradientAscentUnlearner()

        weights = np.random.randn(5)
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)

        result = unlearner.unlearn(weights, (X_forget, y_forget), noise_scale=0.1, num_steps=5)

        assert get_result_value(result, "success") is True
        metadata = get_result_value(result, "metadata")
        assert metadata["noise_scale"] == 0.1

    def test_verify_unlearning(self):
        """Test verification of unlearning."""
        unlearner = GradientAscentUnlearner()

        original_model = np.random.randn(5)
        modified_model = original_model + np.random.randn(5) * 0.5
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)

        score = unlearner.verify_unlearning(modified_model, (X_forget, y_forget), original_model)

        assert 0.0 <= score <= 1.0


class TestFisherForgettingUnlearner:
    """Tests for FisherForgettingUnlearner."""

    def test_init_default(self):
        """Test initialization with default config."""
        unlearner = FisherForgettingUnlearner()
        assert unlearner.name == "fisher_forgetting"
        assert unlearner.config.damping == 1e-4
        assert unlearner.config.use_diagonal is True

    def test_init_custom_config(self):
        """Test initialization with custom config."""
        config = FisherForgettingConfig(
            damping=1e-3,
            num_samples_fisher=500,
            use_diagonal=False,
        )
        unlearner = FisherForgettingUnlearner(config=config)
        assert unlearner.config.damping == 1e-3
        assert unlearner.config.num_samples_fisher == 500

    def test_unlearn_numpy_weights(self):
        """Test unlearning on numpy weight matrix."""
        unlearner = FisherForgettingUnlearner()

        # Create simple model weights
        weights = np.random.randn(10)

        # Create forget and retain data
        X_forget = np.random.randn(20, 10)
        y_forget = np.random.randn(20)
        X_retain = np.random.randn(100, 10)
        y_retain = np.random.randn(100)

        result = unlearner.unlearn(weights, (X_forget, y_forget), retain_data=(X_retain, y_retain))

        assert get_result_value(result, "success") is True
        assert get_result_value(result, "samples_forgotten") == 20
        assert get_result_value(result, "time_seconds") is not None

    def test_unlearn_without_retain_data_raises(self):
        """Test unlearning fails when retain data is missing."""
        unlearner = FisherForgettingUnlearner()

        weights = np.random.randn(5)
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)

        with pytest.raises(ValueError, match="retain_data is required"):
            unlearner.unlearn(weights, (X_forget, y_forget))

    def test_unlearn_with_noise(self):
        """Test unlearning with differential privacy noise."""
        unlearner = FisherForgettingUnlearner()

        weights = np.random.randn(5)
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)

        X_retain = np.random.randn(50, 5)
        y_retain = np.random.randn(50)

        result = unlearner.unlearn(
            weights,
            (X_forget, y_forget),
            retain_data=(X_retain, y_retain),
            noise_scale=0.05,
        )

        assert get_result_value(result, "success") is True
        metadata = get_result_value(result, "metadata")
        assert metadata["noise_scale"] == 0.05

    def test_verify_unlearning(self):
        """Test verification of unlearning."""
        unlearner = FisherForgettingUnlearner()

        original_model = np.random.randn(5)
        modified_model = original_model + np.random.randn(5) * 0.5
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)

        score = unlearner.verify_unlearning(modified_model, (X_forget, y_forget), original_model)

        assert 0.0 <= score <= 1.0

    def test_clear_cache(self):
        """Test clearing Fisher cache."""
        unlearner = FisherForgettingUnlearner()

        # Run unlearning to populate cache
        weights = np.random.randn(5)
        X_forget = np.random.randn(10, 5)
        y_forget = np.random.randn(10)
        X_retain = np.random.randn(20, 5)
        y_retain = np.random.randn(20)
        unlearner.unlearn(weights, (X_forget, y_forget), retain_data=(X_retain, y_retain))

        # Clear cache
        unlearner.clear_cache()
        assert unlearner._fisher_cache is None


class TestUnlearnerProperties:
    """Test unlearner property methods."""

    def test_gradient_ascent_properties(self):
        """Test GradientAscentUnlearner properties."""
        unlearner = GradientAscentUnlearner()
        assert isinstance(unlearner.name, str)
        # method can be string or enum
        assert unlearner.method is not None

    def test_fisher_forgetting_properties(self):
        """Test FisherForgettingUnlearner properties."""
        unlearner = FisherForgettingUnlearner()
        assert isinstance(unlearner.name, str)
        # method can be string or enum
        assert unlearner.method is not None
