import pytest
import numpy as np
from privacy_tuner.core.utility.tstr_classifier import TSTRClassifier


def test_tstr_classifier_basic(classification_data):
    """测试TSTR分类器基本功能"""
    original = classification_data
    # 分离特征和目标
    X = original[:, :-1]
    y = original[:, -1].astype(int)
    # 对特征加噪声
    X_perturbed = X + np.random.normal(0, 0.1, size=X.shape)
    perturbed = np.column_stack([X_perturbed, y])

    evaluator = TSTRClassifier(target_column=-1, random_state=42)
    score = evaluator.evaluate(original, perturbed)
    assert 0.0 <= score <= 1.0


def test_tstr_classifier_perfect_synthetic(classification_data):
    """测试使用完美合成数据（即原始数据本身）时准确率应为1.0左右"""
    original = classification_data
    X = original[:, :-1]
    y = original[:, -1].astype(int)
    # 完美合成数据：特征不加噪声
    X_perturbed = X.copy()
    perturbed = np.column_stack([X_perturbed, y])

    evaluator = TSTRClassifier(target_column=-1, random_state=42)
    score = evaluator.evaluate(original, perturbed)
    assert score > 0.9  # 应该接近1.0


def test_tstr_classifier_custom_target(classification_data):
    """测试指定不同的目标列"""
    original = classification_data
    # 目标列在原数据最后一列，我们要将其移到第一列
    y = original[:, -1].astype(int)
    X = original[:, :-1]
    original_swapped = np.column_stack([y, X])
    # 对特征（除第一列外）加噪声
    X_swapped = original_swapped[:, 1:]
    X_perturbed = X_swapped + np.random.normal(0, 0.1, size=X_swapped.shape)
    perturbed = np.column_stack([original_swapped[:, 0], X_perturbed])

    evaluator = TSTRClassifier(target_column=0, random_state=42)
    score = evaluator.evaluate(original_swapped, perturbed, task={'target_column': 0})
    assert 0.0 <= score <= 1.0