import numpy as np
import pandas as pd
from typing import Dict, Any
from copulas.multivariate import GaussianMultivariate
from copulas.univariate import GaussianUnivariate  # 强制使用高斯分布
from ..privacy_mechanism import PrivacyMechanism


class CopulaMechanism(PrivacyMechanism):
    """
    基于 Copula 的生成式隐私机制。
    学习原始数据的分布并生成合成数据。
    参数:
        - sampling_ratio: float, 生成样本数量与原始样本数量的比例 (0, +∞)
        - noise_scale: float, 可选，添加到合成数据的噪声尺度 (0, 1]，0 表示无噪声
    """

    def __init__(self):
        self.model = None
        self.fitted = False

    def apply(self, data: np.ndarray, params: Dict[str, Any]) -> np.ndarray:
        """
        应用 Copula 机制，返回合成数据。

        Parameters
        ----------
        data : np.ndarray
            原始数据，形状 (n_samples, n_features)
        params : Dict[str, Any]
            必须包含:
                - sampling_ratio: float
                - noise_scale: float (可选，默认为 0)

        Returns
        -------
        np.ndarray
            合成数据，形状 (int(n_samples * sampling_ratio), n_features)
        """
        # 提取参数
        sampling_ratio = params.get('sampling_ratio', 1.0)
        noise_scale = params.get('noise_scale', 0.0)

        n_samples = data.shape[0]
        n_synth = int(n_samples * sampling_ratio)
        if n_synth <= 0:
            return np.empty((0, data.shape[1]))

        # 将数据转换为 DataFrame 以便 copulas 处理
        df = pd.DataFrame(data)

        # 强制使用高斯分布拟合，避免自动选择分布导致的错误
        self.model = GaussianMultivariate(distribution=GaussianUnivariate)
        self.model.fit(df)
        self.fitted = True

        # 生成合成数据
        synthetic_df = self.model.sample(n_synth)
        synthetic = synthetic_df.values

        # 添加噪声（可选）
        if noise_scale > 0:
            # 为每列计算标准差
            stds = np.std(data, axis=0)
            # 生成噪声：正态分布，均值为0，标准差为 noise_scale * std
            noise = np.random.normal(0, noise_scale * stds, size=synthetic.shape)
            synthetic = synthetic + noise

        return synthetic