import pytest
import numpy as np
import pandas as pd
import tempfile
import os

from privacy_tuner.input import CSVLoader
from privacy_tuner.core.mechanisms import CopulaMechanism
from privacy_tuner.core.risk import EmpiricalRiskMeter
from privacy_tuner.core.utility import TSTRClassifier
from privacy_tuner.core.optimizers import GridSearchOptimizer
from privacy_tuner.output import CSVExporter, SimpleReporter


def test_end_to_end():
    # 创建临时CSV数据
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
        f.write("age,income,label\n")
        f.write("25,50000,0\n")
        f.write("30,60000,0\n")
        f.write("35,70000,1\n")
        f.write("40,80000,1\n")
        f.write("45,90000,1\n")
        csv_path = f.name

    # 加载数据
    loader = CSVLoader()
    data_df = loader.load(csv_path)
    data = data_df.values

    # 设置组件
    mechanism = CopulaMechanism()
    risk = EmpiricalRiskMeter(random_state=42)
    utility = TSTRClassifier(target_column=-1, random_state=42)
    optimizer = GridSearchOptimizer(n_jobs=1, use_progress=False)

    param_space = {
        'sampling_ratio': [0.5, 1.0],
        'noise_scale': [0.0, 0.1]
    }
    risk_constraint = (0.0, 1.0)

    # 执行搜索
    best_params, best_utility, best_risk = optimizer.search(
        original_data=data,
        mechanism=mechanism,
        risk_meter=risk,
        utility_evaluator=utility,
        param_space=param_space,
        risk_constraint=risk_constraint,
        max_iter=10
    )

    assert best_params is not None
    assert best_utility != -np.inf
    assert best_risk is not None

    # 生成最终数据并导出
    final_perturbed = mechanism.apply(data, best_params)

    with tempfile.TemporaryDirectory() as tmpdir:
        exporter = CSVExporter()
        out_path = os.path.join(tmpdir, 'perturbed.csv')
        exporter.export(final_perturbed, out_path, index=False, header=False)
        assert os.path.exists(out_path)

        # 生成报告
        reporter = SimpleReporter()
        report = reporter.generate(
            original_info={"shape": data.shape},
            mechanism=mechanism,
            search_results={
                "best_params": best_params,
                "best_utility": best_utility,
                "best_risk": best_risk
            }
        )
        report_path = os.path.join(tmpdir, 'report.txt')
        with open(report_path, 'w') as f:
            f.write(report)
        assert os.path.exists(report_path)

    # 清理临时文件
    os.unlink(csv_path)