"""
01_bias_detection.py - 性别偏见探测实验

使用说明：
    本脚本构建"1×2"反事实案例矩阵（同一场景，只改变性别为男/女），
    向多个主流LLM API发送请求，解析赔偿金建议，计算偏见偏移率Δ及KL散度。

运行前提：
    1. 已安装所有依赖：pip install -r requirements.txt
    2. 已在环境变量中配置API密钥（详见 config.py）
    3. 工作目录为 实验代码/

运行方式：
    python 01_bias_detection.py

输出：
    - data/bias_detection_results.csv  原始结果（每行=一个模型×场景×性别）
    - data/bias_summary.csv            汇总表（每行=一个模型×场景，含Δ值）
    - 终端打印各模型的平均偏见偏移率及KL散度

注意：
    - 如某个API密钥未配置，该模型将被自动跳过并打印警告
    - 每次API调用之间有1秒延迟，避免触发速率限制
    - 若要快速测试框架逻辑，可将 config.N_REPEAT 设为1

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

from __future__ import annotations

import re
import sys
import time
import json
import logging
import traceback
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import requests
from scipy.special import rel_entr
from tqdm import tqdm

# 将脚本所在目录加入路径，确保能导入 config
sys.path.insert(0, str(Path(__file__).parent))
import config

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


# ============================================================
# 反事实案例矩阵 & 性别列表（从 config 集中导入，保证各脚本一致）
# ============================================================
from config import SCENARIOS, GENDERS


# ============================================================
# API 调用函数
# ============================================================

def call_openai(prompt: str, model_cfg: dict) -> Optional[str]:
    """调用 OpenAI GPT-4 API，返回模型原始文本响应"""
    try:
        from openai import OpenAI
        client = OpenAI(api_key=config.OPENAI_API_KEY)
        response = client.chat.completions.create(
            model=model_cfg["model_id"],
            messages=[{"role": "user", "content": prompt}],
            temperature=model_cfg["temperature"],
            max_tokens=model_cfg["max_tokens"],
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"OpenAI API 调用失败: {e}")
        return None


def call_anthropic(prompt: str, model_cfg: dict) -> Optional[str]:
    """调用 Anthropic Claude API，返回模型原始文本响应"""
    try:
        import anthropic
        client = anthropic.Anthropic(api_key=config.ANTHROPIC_API_KEY)
        message = client.messages.create(
            model=model_cfg["model_id"],
            max_tokens=model_cfg["max_tokens"],
            messages=[{"role": "user", "content": prompt}],
        )
        return message.content[0].text
    except Exception as e:
        logger.error(f"Anthropic API 调用失败: {e}")
        return None


def call_dashscope(prompt: str, model_cfg: dict) -> Optional[str]:
    """调用阿里云 DashScope Qwen API，返回模型原始文本响应"""
    try:
        import dashscope
        from dashscope import Generation
        dashscope.api_key = config.DASHSCOPE_API_KEY
        response = Generation.call(
            model=model_cfg["model_id"],
            messages=[{"role": "user", "content": prompt}],
            temperature=model_cfg["temperature"],
            max_tokens=model_cfg["max_tokens"],
            result_format="message",
        )
        if response.status_code == 200:
            return response.output.choices[0]["message"]["content"]
        else:
            logger.error(f"DashScope API 错误: {response.code} - {response.message}")
            return None
    except Exception as e:
        logger.error(f"DashScope API 调用失败: {e}")
        return None


def _get_ernie_access_token() -> Optional[str]:
    """获取百度 ERNIE API 的 access_token"""
    url = (
        f"https://aip.baidubce.com/oauth/2.0/token"
        f"?grant_type=client_credentials"
        f"&client_id={config.ERNIE_API_KEY}"
        f"&client_secret={config.ERNIE_SECRET_KEY}"
    )
    try:
        resp = requests.post(url, timeout=10)
        resp.raise_for_status()
        return resp.json().get("access_token")
    except Exception as e:
        logger.error(f"获取 ERNIE access_token 失败: {e}")
        return None


def call_ernie(prompt: str, model_cfg: dict) -> Optional[str]:
    """调用百度千帆 ERNIE API，返回模型原始文本响应"""
    try:
        token = _get_ernie_access_token()
        if not token:
            return None
        url = (
            f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/"
            f"{model_cfg['model_id']}?access_token={token}"
        )
        payload = {
            "messages": [{"role": "user", "content": prompt}],
            "temperature": model_cfg["temperature"],
            "max_output_tokens": model_cfg["max_tokens"],
        }
        resp = requests.post(url, json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        return data.get("result")
    except Exception as e:
        logger.error(f"ERNIE API 调用失败: {e}")
        return None


def call_deepseek(prompt: str, model_cfg: dict) -> Optional[str]:
    """调用 DeepSeek API（兼容 OpenAI 协议），返回模型原始文本响应"""
    try:
        from openai import OpenAI
        client = OpenAI(
            api_key=config.DEEPSEEK_API_KEY,
            base_url="https://api.deepseek.com",
        )
        response = client.chat.completions.create(
            model=model_cfg["model_id"],
            messages=[{"role": "user", "content": prompt}],
            temperature=model_cfg["temperature"],
            max_tokens=model_cfg["max_tokens"],
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"DeepSeek API 调用失败: {e}")
        return None


def call_dashscope_openai(prompt: str, model_cfg: dict) -> Optional[str]:
    """通过阿里云 DashScope OpenAI 兼容接口调用 Qwen 模型"""
    try:
        from openai import OpenAI
        client = OpenAI(
            api_key=config.DASHSCOPE_API_KEY,
            base_url=config.DASHSCOPE_BASE_URL,
        )
        response = client.chat.completions.create(
            model=model_cfg["model_id"],
            messages=[{"role": "user", "content": prompt}],
            temperature=model_cfg["temperature"],
            max_tokens=model_cfg["max_tokens"],
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"DashScope OpenAI 接口调用失败: {e}")
        return None


def call_ofox_openai(prompt: str, model_cfg: dict) -> Optional[str]:
    """通过 ofox.ai 代理调用 OpenAI 兼容模型（GPT / Qwen / DeepSeek）"""
    try:
        from openai import OpenAI
        client = OpenAI(
            api_key=config.OFOX_API_KEY,
            base_url=config.OFOX_OPENAI_BASE_URL,
        )
        response = client.chat.completions.create(
            model=model_cfg["model_id"],
            messages=[{"role": "user", "content": prompt}],
            temperature=model_cfg["temperature"],
            max_tokens=model_cfg["max_tokens"],
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"ofox OpenAI 代理调用失败: {e}")
        return None


def call_ofox_anthropic(prompt: str, model_cfg: dict) -> Optional[str]:
    """通过 ofox.ai 代理调用 Anthropic Claude 模型"""
    try:
        import anthropic
        client = anthropic.Anthropic(
            api_key=config.OFOX_API_KEY,
            base_url=config.OFOX_ANTHROPIC_BASE_URL,
        )
        message = client.messages.create(
            model=model_cfg["model_id"],
            max_tokens=model_cfg["max_tokens"],
            messages=[{"role": "user", "content": prompt}],
        )
        return message.content[0].text
    except Exception as e:
        logger.error(f"ofox Anthropic 代理调用失败: {e}")
        return None


# 提供者 → 调用函数的映射表
PROVIDER_FUNC = {
    "openai":           call_openai,
    "anthropic":        call_anthropic,
    "dashscope":        call_dashscope,
    "dashscope_openai": call_dashscope_openai,
    "ernie":            call_ernie,
    "deepseek":         call_deepseek,
    "ofox_openai":      call_ofox_openai,
    "ofox_anthropic":   call_ofox_anthropic,
}


# ============================================================
# 工具函数
# ============================================================

def build_prompt(scenario: dict, gender: str) -> str:
    """根据场景信息和性别生成提示词"""
    return config.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 parse_compensation(text: str) -> Optional[float]:
    """
    从模型响应文本中提取赔偿金数额。

    支持多种格式：
        - "建议赔偿金：30000元"
        - "赔偿金为 30,000 元"
        - "应支付30000元"
        - "30000"
    返回 float 类型金额，若解析失败返回 None。
    """
    if not text:
        return None
    # 去除千分位逗号，统一处理
    text_clean = text.replace(",", "").replace("，", "")
    # 优先匹配"建议赔偿金：XXXXX元"格式
    patterns = [
        r"建议赔偿金[：:]\s*(\d+(?:\.\d+)?)\s*元",
        r"赔偿金[为是][：:\s]*(\d+(?:\.\d+)?)\s*元",
        r"应[支付赔偿]{2,4}[：:\s]*(\d+(?:\.\d+)?)\s*元",
        r"(\d{4,}(?:\.\d+)?)\s*元",   # 回退：4位以上数字后跟"元"
        r"(\d{4,}(?:\.\d+)?)",         # 最后回退：纯数字（至少4位）
    ]
    for pat in patterns:
        m = re.search(pat, text_clean)
        if m:
            try:
                return float(m.group(1))
            except ValueError:
                continue
    return None


def compute_kl_divergence(
    dist_male: np.ndarray,
    dist_female: np.ndarray,
    n_bins: int = 20,
) -> float:
    """
    计算男性分布与女性分布之间的 KL 散度 D_KL(P_male || P_female)。

    参数：
        dist_male:   男性场景赔偿金数组
        dist_female: 女性场景赔偿金数组
        n_bins:      直方图分箱数

    返回：
        KL 散度值（float）
    """
    if len(dist_male) == 0 or len(dist_female) == 0:
        return float("nan")
    # 构建共同 bin 边界
    all_vals = np.concatenate([dist_male, dist_female])
    min_val, max_val = all_vals.min(), all_vals.max()
    if min_val == max_val:
        return 0.0
    bins = np.linspace(min_val, max_val, n_bins + 1)
    # 计算归一化直方图（概率分布）
    p_male, _   = np.histogram(dist_male,   bins=bins, density=False)
    p_female, _ = np.histogram(dist_female, bins=bins, density=False)
    # 平滑处理，避免除零
    p_male   = (p_male   + 1e-10) / (p_male.sum()   + 1e-10 * len(p_male))
    p_female = (p_female + 1e-10) / (p_female.sum() + 1e-10 * len(p_female))
    return float(np.sum(rel_entr(p_male, p_female)))


def check_api_key(provider: str) -> bool:
    """检查指定提供者的 API 密钥是否已配置"""
    key_map = {
        "openai":           config.OPENAI_API_KEY,
        "anthropic":        config.ANTHROPIC_API_KEY,
        "dashscope":        config.DASHSCOPE_API_KEY,
        "dashscope_openai": config.DASHSCOPE_API_KEY,
        "ernie":            config.ERNIE_API_KEY,
        "deepseek":         config.DEEPSEEK_API_KEY,
        "ofox_openai":      config.OFOX_API_KEY,
        "ofox_anthropic":   config.OFOX_API_KEY,
    }
    key = key_map.get(provider, "")
    return bool(key and key.strip())


# ============================================================
# 核心实验流程
# ============================================================

def run_bias_detection() -> pd.DataFrame:
    """
    执行完整偏见探测实验。

    返回：
        包含所有原始结果的 DataFrame，列为：
        [model, scenario_id, gender, repeat, raw_response, compensation, legal_compensation]
    """
    records: list[dict] = []

    for model_name, model_cfg in config.PROBE_MODELS.items():
        provider = model_cfg["provider"]

        # 检查 API 密钥
        if not check_api_key(provider):
            logger.warning(
                f"[{model_name}] API 密钥未配置，跳过该模型（提供者: {provider}）"
            )
            continue

        call_fn = PROVIDER_FUNC[provider]
        logger.info(f"===== 开始探测模型: {model_name} =====")

        for scenario in tqdm(SCENARIOS, desc=f"{model_name} 场景进度"):
            for gender in GENDERS:
                prompt = build_prompt(scenario, gender)

                for repeat_idx in range(config.N_REPEAT):
                    raw_text: Optional[str] = None
                    compensation: Optional[float] = None
                    try:
                        raw_text = call_fn(prompt, model_cfg)
                        compensation = parse_compensation(raw_text) if raw_text else None
                    except Exception:
                        logger.error(
                            f"[{model_name}] 场景{scenario['id']} 性别={gender} "
                            f"第{repeat_idx+1}次调用异常:\n{traceback.format_exc()}"
                        )
                    finally:
                        records.append({
                            "model":              model_name,
                            "scenario_id":        scenario["id"],
                            "position":           scenario["position"],
                            "years":              scenario["years"],
                            "salary":             scenario["salary"],
                            "gender":             gender,
                            "repeat":             repeat_idx + 1,
                            "raw_response":       raw_text,
                            "compensation":       compensation,
                            "legal_compensation": scenario["legal_compensation"],
                        })
                    # 速率限制保护
                    time.sleep(1.0)

    df = pd.DataFrame(records)
    return df


def compute_bias_summary(df: pd.DataFrame) -> pd.DataFrame:
    """
    基于原始结果计算每个（模型×场景）的偏见指标。

    返回列：
        [model, scenario_id, R_male, R_female, delta, legal_compensation]
    其中 delta = (R_male - R_female) / mean(R_male, R_female)
    """
    if df.empty:
        logger.warning("原始结果为空，无法计算偏见摘要")
        return pd.DataFrame()

    # 过滤掉解析失败的行
    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
    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)

    # 计算偏见偏移率 Δ
    mean_r = (pivot["R_male"] + pivot["R_female"]) / 2
    pivot["delta"] = (pivot["R_male"] - pivot["R_female"]) / mean_r.replace(0, np.nan)

    return pivot


def compute_model_kl(df: pd.DataFrame) -> pd.DataFrame:
    """
    按模型计算跨所有场景的 KL 散度（男性输出分布 vs 女性输出分布）。

    返回列：[model, kl_divergence, mean_abs_delta]
    """
    if df.empty:
        return pd.DataFrame()
    valid = df.dropna(subset=["compensation"])
    rows = []
    for model_name, sub in valid.groupby("model"):
        male_vals   = sub[sub["gender"] == "男"]["compensation"].values
        female_vals = sub[sub["gender"] == "女"]["compensation"].values
        kl = compute_kl_divergence(male_vals, female_vals)

        # 平均绝对偏见偏移率（跨场景均值）
        summary = compute_bias_summary(
            sub[sub["model"] == model_name] if "model" in sub.columns else sub
        )
        mean_abs_delta = (
            summary["delta"].abs().mean()
            if not summary.empty and "delta" in summary.columns
            else float("nan")
        )
        rows.append({
            "model":          model_name,
            "kl_divergence":  kl,
            "mean_abs_delta": mean_abs_delta,
        })
    return pd.DataFrame(rows)


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

def main() -> None:
    logger.info("=" * 60)
    logger.info("法律LLM性别偏见探测实验 启动")
    logger.info(f"场景数量: {len(SCENARIOS)}，性别: {GENDERS}，重复次数: {config.N_REPEAT}")
    logger.info("=" * 60)

    # 执行实验
    df_raw = run_bias_detection()

    if df_raw.empty:
        logger.error("未获得任何有效结果，请检查API密钥配置后重试")
        return

    # 保存原始结果
    df_raw.to_csv(config.BIAS_RESULT_CSV, index=False, encoding="utf-8-sig")
    logger.info(f"原始结果已保存至: {config.BIAS_RESULT_CSV}")

    # 计算偏见摘要
    df_summary = compute_bias_summary(df_raw)
    summary_path = config.BIAS_RESULT_CSV.replace(".csv", "_summary.csv")
    df_summary.to_csv(summary_path, index=False, encoding="utf-8-sig")
    logger.info(f"偏见摘要已保存至: {summary_path}")

    # 计算模型级 KL 散度
    df_kl = compute_model_kl(df_raw)
    kl_path = config.BIAS_RESULT_CSV.replace(".csv", "_kl.csv")
    df_kl.to_csv(kl_path, index=False, encoding="utf-8-sig")
    logger.info(f"KL散度结果已保存至: {kl_path}")

    # 打印汇总报告
    logger.info("\n" + "=" * 60)
    logger.info("模型偏见探测汇总报告")
    logger.info("=" * 60)
    if not df_kl.empty:
        for _, row in df_kl.iterrows():
            logger.info(
                f"  模型: {row['model']:30s} | "
                f"KL散度: {row['kl_divergence']:.4f} | "
                f"平均|Δ|: {row['mean_abs_delta']:.4f}"
            )

    logger.info("\n偏见探测实验完成！")


if __name__ == "__main__":
    main()
