import pytest
import numpy as np
from typing import Dict, Any, Optional
from privacy_tuner.core.optimizers.grid_search import GridSearchOptimizer
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_level = params.get('noise', 0.0)
        noise = np.random.normal(0, noise_level, size=data.shape)
        return data + noise


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_grid_search_basic():
    optimizer = GridSearchOptimizer()
    mechanism = TestMechanism()
    risk_meter = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(100, 5)

    param_space = {
        'noise': [0.0, 0.5, 1.0]
    }
    risk_constraint = (0.0, np.inf)  # 无约束，应选择效用最高的参数（即噪声最小的）

    best_params, best_utility, best_risk = optimizer.search(
        data, mechanism, risk_meter, utility, param_space, risk_constraint
    )

    # 最佳参数应为噪声最小的 0.0
    assert best_params['noise'] == 0.0
    assert best_utility >= 0  # 因为噪声为0，MSE=0，效用=0
    assert best_risk == 0.0


def test_grid_search_with_constraint():
    optimizer = GridSearchOptimizer()
    mechanism = TestMechanism()
    risk_meter = TestRiskMeter()
    utility = TestUtility()

    data = np.random.randn(100, 5)

    param_space = {
        'noise': [0.0, 0.5, 1.0]
    }
    risk_constraint = (0.5, 1.5)  # 只接受风险在0.5~1.5之间

    best_params, best_utility, best_risk = optimizer.search(
        data, mechanism, risk_meter, utility, param_space, risk_constraint
    )

    # 噪声为0的风险为0，不符合；噪声0.5的风险应在0.25左右（可能低于0.5）；噪声1.0的风险应在1.0左右。
    # 根据数据，可能只有噪声1.0满足约束，也可能噪声0.5略低于0.5，这里只检查返回参数不为空
    assert best_params is not None
    assert risk_constraint[0] <= best_risk <= risk_constraint[1]