"""
02_data_debiasing.py - 数据预处理与净化（镜像对齐策略）

使用说明：
    本脚本实现"镜像对齐策略"，将偏见探测实验产出的有偏训练数据转换为对齐的公平训练数据，
    并保存为 HuggingFace datasets 格式，供后续模型微调脚本直接读取。

核心逻辑：
    输入：(scenario_male, R_male), (scenario_female, R_female)   ← 来自01的偏见探测结果
    输出：(scenario_male, R_aligned), (scenario_female, R_aligned)
    其中 R_aligned = max(R_legal, (R_male + R_female) / 2)
         R_legal   = 按《劳动合同法》第87条计算的法定赔偿金

运行前提：
    1. 已运行 01_bias_detection.py，data/bias_detection_results.csv 已生成
       （若文件不存在，脚本将使用模拟数据演示流程）
    2. 工作目录为 实验代码/

运行方式：
    python 02_data_debiasing.py

输出：
    - data/debiased_dataset/   HuggingFace DatasetDict（含 train/validation 子集）
    - data/debiased_dataset_preview.csv  便于人工检查的 CSV 预览版本

作者：研究团队
日期：2025
"""

from __future__ import annotations

import sys
import logging
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
from datasets import Dataset, DatasetDict

sys.path.insert(0, str(Path(__file__).parent))
import config
from config import SCENARIOS, GENDERS, PROMPT_TEMPLATE

# ============================================================
# 日志配置
# ============================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)


# ============================================================
# 辅助函数
# ============================================================

def build_instruction(scenario: dict, gender: str) -> str:
    """构建微调用的 instruction 字段（Alpaca 格式）"""
    return PROMPT_TEMPLATE.format(
        gender=gender,
        position=scenario["position"],
        years=scenario["years"],
        salary=scenario["salary"],
        education=scenario["education"],
        family=scenario["family"],
        reason=scenario["reason"],
        evidence=scenario["evidence"],
    )


def build_output(r_aligned: float, scenario: dict) -> str:
    """构建微调用的 output 字段（标准化法律答案）"""
    years = scenario["years"]
    salary = scenario["salary"]
    economic_comp = years * salary          # 经济补偿金
    r_legal = economic_comp * config.COMPENSATION_MULTIPLIER  # 法定赔偿金
    return (
        f"根据《劳动合同法》第47条，经济补偿金为：{years}年 × {salary}元/月 = {economic_comp}元。\n"
        f"根据第87条，违法解除赔偿金 = 经济补偿金 × 2 = {r_legal}元。\n"
        f"建议赔偿金：{int(r_aligned)}元"
    )


def compute_aligned_compensation(
    r_male: Optional[float],
    r_female: Optional[float],
    r_legal: float,
) -> float:
    """
    计算对齐后的赔偿金 R_aligned。

    策略：直接使用法定标准值 R* = 工龄 × 月薪 × 2。

    原因：本任务存在唯一正确答案（法律公式），使用 R* 作为监督目标
    同时保障法律准确性与性别中立性。若使用模型输出均值，当模型系统性
    偏离公式时（如 Claude 法律符合率仅 2%）会引入错误标签。

    参数：
        r_male:   模型对男性场景的平均输出（保留参数供日志使用，不影响计算）
        r_female: 模型对女性场景的平均输出（同上）
        r_legal:  法律条文计算的标准赔偿金（R* = 工龄 × 月薪 × 2）

    返回：
        R_legal（法定标准值）
    """
    return r_legal


def load_bias_results() -> Optional[pd.DataFrame]:
    """
    尝试从 01_bias_detection.py 的输出文件加载原始偏见探测结果。
    若文件不存在，返回 None（调用方将使用模拟数据）。
    """
    csv_path = Path(config.BIAS_RESULT_CSV)
    if not csv_path.exists():
        logger.warning(
            f"偏见探测结果文件不存在: {csv_path}\n"
            "将使用模拟数据演示数据净化流程。"
        )
        return None

    df = pd.read_csv(csv_path, encoding="utf-8-sig")
    logger.info(f"已加载偏见探测结果: {len(df)} 条记录")
    return df


def generate_mock_bias_results() -> pd.DataFrame:
    """
    生成模拟偏见探测数据（用于在未配置API时测试净化逻辑）。

    模拟规律：
        - 男性平均比女性多获得 10% 的赔偿建议（体现偏见）
        - 加入随机噪声模拟模型不确定性
    """
    np.random.seed(config.RANDOM_SEED)
    records = []
    for scenario in SCENARIOS:
        r_legal = scenario["legal_compensation"]
        for gender in GENDERS:
            # 模拟偏见：男性建议值偏高约10%
            bias_factor = 1.10 if gender == "男" else 0.92
            for repeat in range(config.N_REPEAT):
                noise = np.random.normal(1.0, 0.05)
                comp = r_legal * bias_factor * noise
                records.append({
                    "model":              "mock_model",
                    "scenario_id":        scenario["id"],
                    "position":           scenario["position"],
                    "years":              scenario["years"],
                    "salary":             scenario["salary"],
                    "gender":             gender,
                    "repeat":             repeat + 1,
                    "raw_response":       f"建议赔偿金：{int(comp)}元",
                    "compensation":       comp,
                    "legal_compensation": r_legal,
                })
    return pd.DataFrame(records)


# ============================================================
# 核心净化逻辑
# ============================================================

def build_pivot_table(df: pd.DataFrame) -> pd.DataFrame:
    """
    将原始结果聚合为宽格式表（按模型×场景×性别均值）。

    返回列：[model, scenario_id, R_male, R_female, legal_compensation]
    """
    valid = df.dropna(subset=["compensation"])
    grouped = (
        valid
        .groupby(["model", "scenario_id", "gender"])["compensation"]
        .mean()
        .reset_index()
        .rename(columns={"compensation": "mean_compensation"})
    )
    pivot = grouped.pivot_table(
        index=["model", "scenario_id"],
        columns="gender",
        values="mean_compensation",
    ).reset_index()
    pivot.columns.name = None

    # 兼容性处理：确保男女列均存在
    if "男" not in pivot.columns:
        pivot["男"] = np.nan
    if "女" not in pivot.columns:
        pivot["女"] = np.nan
    pivot = pivot.rename(columns={"男": "R_male", "女": "R_female"})

    # 合并法律标准值
    legal_map = {s["id"]: s["legal_compensation"] for s in SCENARIOS}
    pivot["legal_compensation"] = pivot["scenario_id"].map(legal_map)

    return pivot


def create_debiased_pairs(pivot: pd.DataFrame) -> list[dict]:
    """
    基于透视表生成对齐后的训练数据对。

    每个场景产生 2 条训练样本（男性 + 女性），
    两者使用相同的 R_aligned 作为监督目标。

    返回：
        list of dict，每条含 instruction, input, output, scenario_id, gender, r_aligned
    """
    # 构建场景信息查找字典
    scenario_map = {s["id"]: s for s in SCENARIOS}
    samples: list[dict] = []

    for _, row in pivot.iterrows():
        scenario_id = row["scenario_id"]
        scenario = scenario_map.get(scenario_id)
        if scenario is None:
            logger.warning(f"未找到场景信息: {scenario_id}，跳过")
            continue

        r_male   = row.get("R_male")
        r_female = row.get("R_female")
        r_legal  = row.get("legal_compensation", scenario["legal_compensation"])

        # 处理可能的 NaN
        r_male   = None if (r_male is None or np.isnan(r_male))     else float(r_male)
        r_female = None if (r_female is None or np.isnan(r_female)) else float(r_female)

        r_aligned = compute_aligned_compensation(r_male, r_female, float(r_legal))

        # 记录对齐信息（用于后续可解释性分析）
        align_note = (
            f"R_male={r_male:.0f}, R_female={r_female:.0f}, "
            f"R_legal={r_legal:.0f}, R_aligned={r_aligned:.0f}"
            if r_male and r_female
            else f"R_legal={r_legal:.0f}, R_aligned={r_aligned:.0f}"
        )

        # 为男女各生成一条样本
        for gender in GENDERS:
            instruction = build_instruction(scenario, gender)
            output_text = build_output(r_aligned, scenario)
            samples.append({
                "instruction":   instruction,
                "input":         "",              # Alpaca 格式，input 留空
                "output":        output_text,
                "scenario_id":   scenario_id,
                "gender":        gender,
                "r_aligned":     r_aligned,
                "r_legal":       float(r_legal),
                "r_male_orig":   r_male   if r_male   is not None else np.nan,
                "r_female_orig": r_female if r_female is not None else np.nan,
                "align_note":    align_note,
            })

    logger.info(f"共生成 {len(samples)} 条对齐训练样本")
    return samples


def build_sft_text(sample: dict) -> str:
    """
    将 Alpaca 格式样本转换为 Qwen Chat 风格的完整文本，
    用于 SFT Trainer 的 dataset_text_field。
    """
    return (
        f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n"
        f"<|im_start|>assistant\n{sample['output']}<|im_end|>"
    )


# ============================================================
# 数据集构建与保存
# ============================================================

def build_and_save_dataset(samples: list[dict], train_ratio: float = 0.85) -> DatasetDict:
    """
    将样本列表构建为 HuggingFace DatasetDict 并保存到磁盘。

    参数：
        samples:     净化后的训练样本列表
        train_ratio: 训练集比例（其余为验证集）

    返回：
        DatasetDict，含 train / validation 两个子集
    """
    if not samples:
        raise ValueError("样本列表为空，无法构建数据集")

    # 添加 SFT 完整文本字段
    for s in samples:
        s["text"] = build_sft_text(s)

    # 随机打乱并分割
    np.random.seed(config.RANDOM_SEED)
    indices = np.random.permutation(len(samples))
    split_idx = int(len(samples) * train_ratio)
    train_indices = indices[:split_idx]
    val_indices   = indices[split_idx:]

    train_samples = [samples[i] for i in train_indices]
    val_samples   = [samples[i] for i in val_indices]

    # 转换为 HuggingFace Dataset
    train_ds = Dataset.from_list(train_samples)
    val_ds   = Dataset.from_list(val_samples)

    dataset_dict = DatasetDict({"train": train_ds, "validation": val_ds})

    # 保存到磁盘
    save_path = config.DEBIASED_DATASET_DIR
    dataset_dict.save_to_disk(save_path)
    logger.info(f"DatasetDict 已保存至: {save_path}")
    logger.info(f"  训练集: {len(train_ds)} 条")
    logger.info(f"  验证集: {len(val_ds)} 条")

    # 额外保存 CSV 预览
    preview_path = Path(config.DATA_DIR) / "debiased_dataset_preview.csv"
    pd.DataFrame(samples).to_csv(preview_path, index=False, encoding="utf-8-sig")
    logger.info(f"CSV 预览已保存至: {preview_path}")

    return dataset_dict


# ============================================================
# 主入口
# ============================================================

def main() -> None:
    logger.info("=" * 60)
    logger.info("数据预处理与净化（镜像对齐策略） 启动")
    logger.info("=" * 60)

    # 1. 加载偏见探测结果
    df_raw = load_bias_results()
    if df_raw is None:
        logger.info("使用模拟偏见数据进行演示...")
        df_raw = generate_mock_bias_results()

    logger.info(f"原始数据行数: {len(df_raw)}")
    logger.info(f"涉及模型: {df_raw['model'].unique().tolist()}")
    logger.info(f"涉及场景: {df_raw['scenario_id'].unique().tolist()}")

    # 2. 构建透视表
    pivot = build_pivot_table(df_raw)
    logger.info(f"透视表行数（模型×场景）: {len(pivot)}")

    # 3. 生成对齐训练样本
    samples = create_debiased_pairs(pivot)

    # 4. 构建并保存 DatasetDict
    dataset_dict = build_and_save_dataset(samples, train_ratio=0.85)

    # 5. 打印数据质量检验报告
    logger.info("\n" + "=" * 60)
    logger.info("数据净化质量报告")
    logger.info("=" * 60)
    df_samples = pd.DataFrame(samples)
    if not df_samples.empty:
        # 按场景检验对齐前后差距
        df_samples["r_male_orig"] = pd.to_numeric(df_samples["r_male_orig"], errors="coerce")
        df_samples["r_female_orig"] = pd.to_numeric(df_samples["r_female_orig"], errors="coerce")
        df_samples["gap_before"] = df_samples["r_male_orig"] - df_samples["r_female_orig"]
        df_samples["gap_after"]  = 0.0  # 对齐后男女输出相同，gap=0

        mean_gap_before = df_samples["gap_before"].abs().mean()
        logger.info(f"  对齐前平均性别赔偿差距: {mean_gap_before:.2f} 元")
        logger.info(f"  对齐后平均性别赔偿差距: 0.00 元（完全对齐）")
        logger.info(f"  总样本数: {len(df_samples)} 条")
        logger.info(f"  男性样本: {(df_samples['gender']=='男').sum()} 条")
        logger.info(f"  女性样本: {(df_samples['gender']=='女').sum()} 条")

    logger.info("\n数据预处理与净化完成！")


if __name__ == "__main__":
    main()
