import os
import json
import numpy as np
from sklearn.datasets import fetch_openml, load_breast_cancer, load_digits
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

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.validation.attack import MembershipInferenceAttack


class CaseLibrary:
    """案例库，管理预运行结果，并包含攻击验证。"""

    def __init__(self, storage_dir='case_results'):
        self.storage_dir = storage_dir
        os.makedirs(storage_dir, exist_ok=True)

    def _run_case(self, name, X, y, attack=True):
        """通用运行逻辑，返回结果字典"""
        # 合并数据和标签
        data = np.column_stack([X, y])

        mechanism = CopulaMechanism()
        risk_meter = 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, 2.0],
            'noise_scale': [0.0, 0.1, 0.5]
        }
        risk_constraint = (0.0, 1.0)

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

        result = {
            'dataset': name,
            'best_params': best_params,
            'best_utility': best_utility,
            'best_risk': best_risk,
            'data_shape': data.shape
        }

        # 如果需要进行攻击验证
        if attack:
            # 划分成员和非成员（训练集和测试集）
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)
            # 训练目标模型（使用最佳参数生成的扰动数据？这里直接用原始数据训练一个随机森林作为目标模型）
            target_model = RandomForestClassifier(n_estimators=10, random_state=42)
            target_model.fit(X_train, y_train)

            # 攻击验证
            attack_tool = MembershipInferenceAttack(random_state=42)
            attack_tool.fit_attack_model(X, y)  # 使用全部数据训练影子模型
            acc, auc = attack_tool.evaluate_target_model(target_model, X_train, X_test)

            result['attack_accuracy'] = acc
            result['attack_auc'] = auc
            result['risk_meter_value'] = best_risk

        # 保存结果
        with open(os.path.join(self.storage_dir, f'{name}.json'), 'w') as f:
            json.dump(result, f, indent=2)

        print(f"案例 {name} 运行完成，结果已保存。")
        return result

    def run_adult_case(self, attack=True):
        """运行Adult数据集案例"""
        adult = fetch_openml(data_id=1590, as_frame=True)
        df = adult.data.select_dtypes(include=[np.number]).dropna()
        y = (adult.target == '>50K').astype(int).values
        X = df.values
        return self._run_case("adult", X, y, attack=attack)

    def run_breast_cancer_case(self, attack=True):
        """运行乳腺癌数据集案例"""
        bc = load_breast_cancer()
        X, y = bc.data, bc.target
        return self._run_case("breast_cancer", X, y, attack=attack)

    def run_digits_case(self, attack=True):
        """运行手写数字数据集案例"""
        digits = load_digits()
        X, y = digits.data, digits.target
        # 二值化（简单处理：奇数类为1，偶数为0）
        y = (y % 2).astype(int)
        return self._run_case("digits", X, y, attack=attack)

    def list_cases(self):
        """列出已运行的案例，并显示攻击结果（如果有）"""
        files = [f for f in os.listdir(self.storage_dir) if f.endswith('.json')]
        for f in files:
            with open(os.path.join(self.storage_dir, f), 'r') as fp:
                data = json.load(fp)
                attack_info = f", 攻击准确率={data.get('attack_accuracy', 'N/A'):.4f}" if 'attack_accuracy' in data else ""
                print(f"{data['dataset']}: 效用={data['best_utility']:.4f}, 风险={data['best_risk']:.4f}{attack_info}")

    def generate_report(self, output_path='case_report.md'):
        lines = ["# 案例库运行报告\n"]
        files = [f for f in os.listdir(self.storage_dir) if f.endswith('.json')]
        for f in files:
            with open(os.path.join(self.storage_dir, f), 'r', encoding='utf-8') as fp:
                data = json.load(fp)
            lines.append(f"## {data['dataset']}\n")
            lines.append(f"- 数据形状: {data['data_shape']}")
            lines.append(f"- 最优参数: {data['best_params']}")
            lines.append(f"- 最优效用: {data['best_utility']:.4f}")
            lines.append(f"- 对应风险: {data['best_risk']:.4f}")
            if 'attack_accuracy' in data:
                lines.append(f"- 攻击准确率: {data['attack_accuracy']:.4f}")
                lines.append(f"- 攻击 AUC: {data['attack_auc']:.4f}")
                diff = data['attack_accuracy'] - data['best_risk']
                lines.append(f"- 风险度量与攻击差异: {diff:.4f}")
            lines.append("")

        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(lines))
        print(f"综合报告已生成: {output_path}")