import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from typing import Dict, Any, Optional
from ..utility_evaluator import UtilityEvaluator


class TSTRClassifier(UtilityEvaluator):
    """
    Train on Synthetic, Test on Real (TSTR) 分类效用评估器。
    在合成数据上训练分类器，在原始数据上测试，返回准确率。

    参数:
        target_column: int, 目标列所在的索引（默认-1，即最后一列）
        classifier: sklearn 分类器实例，默认使用 RandomForestClassifier
        random_state: int, 随机种子（默认42）
    """

    def __init__(self, target_column: int = -1,
                 classifier=None,
                 random_state: int = 42):
        self.target_column = target_column
        self.classifier = classifier if classifier else RandomForestClassifier(n_estimators=10,
                                                                               random_state=random_state)
        self.random_state = random_state

    def evaluate(self, original: np.ndarray, perturbed: np.ndarray,
                 task: Optional[Dict[str, Any]] = None) -> float:
        """
        评估 TSTR 分类准确率。

        Parameters
        ----------
        original : np.ndarray
            原始数据，最后一列为目标列（除非 task 中指定 target_column）。
        perturbed : np.ndarray
            合成数据，应与 original 具有相同的列数。
        task : dict, optional
            可包含 'target_column' 覆盖默认值。

        Returns
        -------
        float
            分类准确率，范围 [0, 1]。
        """
        # 确定目标列索引
        target_col = self.target_column
        if task and 'target_column' in task:
            target_col = task['target_column']

        # 分离特征和标签
        # 原始数据
        X_original = np.delete(original, target_col, axis=1)
        y_original = original[:, target_col]

        # 合成数据
        X_synth = np.delete(perturbed, target_col, axis=1)
        y_synth = perturbed[:, target_col]

        # 对标签进行离散化处理（四舍五入到最近整数，确保是 0/1）
        y_synth = np.round(y_synth).astype(int)
        y_original = np.round(y_original).astype(int)  # 原始已经是整数，但确保

        # 在合成数据上训练，在原始数据上测试
        self.classifier.fit(X_synth, y_synth)
        y_pred = self.classifier.predict(X_original)
        acc = accuracy_score(y_original, y_pred)

        return acc