import pytest
import numpy as np
from typing import Dict, Any, Optional
from privacy_tuner.core.optimizers.genetic import GeneticOptimizer
from privacy_tuner.core.privacy_mechanism import PrivacyMechanism
from privacy_tuner.core.risk_meter import RiskMeter
from privacy_tuner.core.utility_evaluator import UtilityEvaluator


# 用于测试的简单机制、风险、效用
class TestMechanism(PrivacyMechanism):
    def apply(self, data: np.ndarray, params: Dict[str, Any]) -> np.ndarray:
        noise = params.get('noise', 0.0)
        return data + np.random.normal(0, noise, size=data.shape)


class TestRiskMeter(RiskMeter):
    def measure(self,
                original: np.ndarray,
                perturbed: np.ndarray,
                params: Optional[Dict[str, Any]] = None) -> float:
        return float(np.mean((original - perturbed) ** 2))


class TestUtility(UtilityEvaluator):
    def evaluate(self,
                 original: np.ndarray,
                 perturbed: np.ndarray,
                 task: Optional[Dict] = None) -> float:
        return -float(np.mean((original - perturbed) ** 2))


def test_genetic_optimizer_basic():
    opt = GeneticOptimizer(population_size=5, generations=3, random_state=42)
    mechanism = TestMechanism()
    risk = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(10, 2)

    param_space = {
        'noise': (0.0, 1.0, 'float')
    }
    risk_constraint = (0.0, np.inf)

    best_params, best_utility, best_risk = opt.search(
        data, mechanism, risk, utility, param_space, risk_constraint, max_iter=20
    )

    assert 'noise' in best_params
    assert 0.0 <= best_params['noise'] <= 1.0
    assert best_utility <= 0
    assert best_risk >= 0


def test_genetic_optimizer_with_constraint():
    opt = GeneticOptimizer(population_size=5, generations=3, random_state=42)
    mechanism = TestMechanism()
    risk = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(10, 2)

    param_space = {
        'noise': (0.0, 1.0, 'float')
    }
    risk_constraint = (0.0, 0.2)

    best_params, best_utility, best_risk = opt.search(
        data, mechanism, risk, utility, param_space, risk_constraint, max_iter=20
    )

    if best_params:  # 可能有可行解
        assert best_risk <= 0.2


def test_genetic_optimizer_mixed_space():
    opt = GeneticOptimizer(population_size=5, generations=3, random_state=42)
    mechanism = TestMechanism()
    risk = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(10, 2)

    param_space = {
        'noise': (0.0, 1.0, 'float'),
        'type': ['a', 'b', 'c']
    }
    risk_constraint = (0.0, np.inf)

    best_params, best_utility, best_risk = opt.search(
        data, mechanism, risk, utility, param_space, risk_constraint, max_iter=20
    )

    assert 'noise' in best_params
    assert 'type' in best_params
    assert best_params['type'] in ['a', 'b', 'c']


def test_genetic_optimizer_integer():
    opt = GeneticOptimizer(population_size=5, generations=3, random_state=42)
    mechanism = TestMechanism()
    risk = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(10, 2)

    param_space = {
        'k': (1, 5, 'int')  # 整数参数，但机制中未使用，仅测试
    }
    risk_constraint = (0.0, np.inf)

    best_params, best_utility, best_risk = opt.search(
        data, mechanism, risk, utility, param_space, risk_constraint, max_iter=20
    )

    assert 'k' in best_params
    assert isinstance(best_params['k'], (int, np.integer))
    assert 1 <= best_params['k'] <= 5