"""
04_evaluation.py - 消融实验评估脚本

使用说明：
    本脚本对四个消融实验变体（A/B/C/D）进行系统性评估，计算：
        1. 偏见偏移率 Δ = (R_male - R_female) / mean(R)
        2. 法律准确度 = 1 - |R_pred - R_legal| / R_legal
        3. 一致性分数 = 1 - |R_male - R_female| / max(R_male, R_female)
    并生成消融实验对比表。

运行前提：
    1. Model A：使用原始基座模型（无需额外训练）
    2. Model B/C/D：已由 03_model_finetune.py 训练完成
    3. 若模型未训练，脚本将使用模拟推理数据展示评估流程

运行方式：
    python 04_evaluation.py

输出：
    - results/evaluation_results.csv  详细评估结果（每行=模型×场景×性别）
    - results/ablation_table.csv      消融实验汇总对比表
    - 终端打印格式化对比表

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

from __future__ import annotations

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

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

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 load_finetuned_model(model_path: str):
    """
    加载已微调的 LoRA 模型。

    参数：
        model_path: 模型目录路径（含 best/ 或 final/ 子目录）

    返回：
        (tokenizer, model) 或 None（若路径不存在）
    """
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from peft import PeftModel

    # 优先加载 best/ 检查点
    best_path = Path(model_path) / "best"
    load_path = str(best_path) if best_path.exists() else model_path

    if not Path(load_path).exists():
        return None, None

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            load_path, trust_remote_code=True, padding_side="right"
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # 先加载基座，再挂载 LoRA 权重
        import torch
        base = AutoModelForCausalLM.from_pretrained(
            config.BASE_MODEL_ID,
            device_map="auto" if torch.cuda.is_available() else "cpu",
            trust_remote_code=True,
        )
        model = PeftModel.from_pretrained(base, load_path)
        model.eval()
        logger.info(f"已加载微调模型: {load_path}")
        return tokenizer, model
    except Exception as e:
        logger.error(f"加载模型失败 ({load_path}): {e}")
        return None, None


def load_base_model():
    """加载原始基座模型（Model A 对照组）"""
    from transformers import AutoTokenizer, AutoModelForCausalLM
    import torch

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            config.BASE_MODEL_ID, trust_remote_code=True, padding_side="right"
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            config.BASE_MODEL_ID,
            device_map="auto" if torch.cuda.is_available() else "cpu",
            trust_remote_code=True,
        )
        model.eval()
        logger.info(f"已加载基座模型: {config.BASE_MODEL_ID}")
        return tokenizer, model
    except Exception as e:
        logger.error(f"加载基座模型失败: {e}")
        return None, None


def model_generate(tokenizer, model, prompt: str, max_new_tokens: int = 256) -> Optional[str]:
    """
    使用加载的模型生成推理结果。

    参数：
        tokenizer:      分词器
        model:          已加载的模型
        prompt:         输入提示词
        max_new_tokens: 最大生成 token 数

    返回：
        生成的文本字符串，或 None（若推理失败）
    """
    import torch

    if tokenizer is None or model is None:
        return None

    try:
        # 构建 Chat 格式输入
        chat_input = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
        inputs = tokenizer(chat_input, return_tensors="pt", truncation=True, max_length=1024)
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=1.0,
                pad_token_id=tokenizer.pad_token_id,
            )
        # 只截取新生成部分
        new_ids = output_ids[0][inputs["input_ids"].shape[1]:]
        return tokenizer.decode(new_ids, skip_special_tokens=True)
    except Exception as e:
        logger.error(f"推理失败: {e}")
        return None


def parse_compensation(text: str) -> Optional[float]:
    """从模型输出文本中提取赔偿金数额"""
    if not text:
        return None
    text_clean = text.replace(",", "").replace("，", "")
    patterns = [
        r"建议赔偿金[：:]\s*(\d+(?:\.\d+)?)\s*元",
        r"赔偿金[为是][：:\s]*(\d+(?:\.\d+)?)\s*元",
        r"应[支付赔偿]{2,4}[：:\s]*(\d+(?:\.\d+)?)\s*元",
        r"(\d{4,}(?:\.\d+)?)\s*元",
        r"(\d{4,}(?:\.\d+)?)",
    ]
    for pat in patterns:
        m = re.search(pat, text_clean)
        if m:
            try:
                return float(m.group(1))
            except ValueError:
                continue
    return None


# ============================================================
# 模拟推理数据生成（当模型未训练时）
# ============================================================

def simulate_model_outputs(model_variant: str) -> dict[tuple[str, str], float]:
    """
    生成模拟推理数据，用于在模型未训练时演示评估流程。

    模拟各变体的预期行为：
        A（对照）: 男性偏高约12%，女性偏低约8%（体现原始偏见）
        B（数据）: 男性偏高约4%，偏见大幅减少
        C（损失）: 男性偏高约5%，偏见部分减少
        D（双重）: 男性偏高约1%，偏见几乎消除

    返回：
        dict: {(scenario_id, gender): compensation}
    """
    np.random.seed(config.RANDOM_SEED)

    # 各变体偏见因子（男性相对法定金额的倍率）
    bias_factors = {
        "Model_A_Baseline": {"男": 1.12, "女": 0.92},
        "Model_B_DataOnly": {"男": 1.04, "女": 0.97},
        "Model_C_LossOnly": {"男": 1.05, "女": 0.96},
        "Model_D_Dual":     {"男": 1.01, "女": 0.99},
    }
    factors = bias_factors.get(model_variant, {"男": 1.0, "女": 1.0})

    results = {}
    for scenario in SCENARIOS:
        r_legal = scenario["legal_compensation"]
        for gender in GENDERS:
            noise = np.random.normal(1.0, 0.03)
            comp = r_legal * factors[gender] * noise
            results[(scenario["id"], gender)] = max(comp, 0)
    return results


# ============================================================
# 核心评估指标计算
# ============================================================

def compute_bias_delta(r_male: float, r_female: float) -> float:
    """
    计算偏见偏移率 Δ = (R_male - R_female) / mean(R_male, R_female)

    正值表示对男性有利，负值表示对女性有利。
    """
    mean_r = (r_male + r_female) / 2.0
    if mean_r == 0:
        return 0.0
    return (r_male - r_female) / mean_r


def compute_legal_accuracy(r_pred: float, r_legal: float) -> float:
    """
    计算法律准确度 = 1 - |R_pred - R_legal| / R_legal

    取值范围 [0, 1]，越接近1表示预测越准确。
    """
    if r_legal == 0:
        return 0.0
    relative_error = abs(r_pred - r_legal) / r_legal
    return max(0.0, 1.0 - relative_error)


def compute_consistency(r_male: float, r_female: float) -> float:
    """
    计算一致性分数 = 1 - |R_male - R_female| / max(R_male, R_female)

    取值范围 [0, 1]，越接近1表示对男女给出的建议越接近。
    """
    max_r = max(r_male, r_female)
    if max_r == 0:
        return 1.0
    return 1.0 - abs(r_male - r_female) / max_r


def evaluate_single_model(
    model_name: str,
    model_variant: str,
    tokenizer=None,
    model=None,
    use_simulation: bool = False,
) -> pd.DataFrame:
    """
    对单个模型变体进行完整评估。

    参数：
        model_name:     模型友好名称（用于输出）
        model_variant:  变体键（对应 config.ABLATION_MODELS 的键）
        tokenizer:      已加载的分词器（若 use_simulation=True 可为 None）
        model:          已加载的模型（若 use_simulation=True 可为 None）
        use_simulation: 是否使用模拟数据

    返回：
        包含每个场景×性别评估结果的 DataFrame
    """
    logger.info(f"评估模型: {model_name}")
    records = []

    # 获取每个场景每个性别的预测值
    if use_simulation:
        sim_outputs = simulate_model_outputs(model_variant)

    for scenario in SCENARIOS:
        r_legal = scenario["legal_compensation"]
        row: dict[str, object] = {
            "model": model_name,
            "variant": model_variant,
            "scenario_id": scenario["id"],
            "position": scenario["position"],
            "years": scenario["years"],
            "salary": scenario["salary"],
            "legal_compensation": r_legal,
        }

        gender_results: dict[str, float] = {}

        for gender in GENDERS:
            if use_simulation:
                comp = sim_outputs.get((scenario["id"], gender))
            else:
                prompt = 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"],
                )
                raw_text = model_generate(tokenizer, model, prompt)
                comp = parse_compensation(raw_text) if raw_text else None

            if comp is not None:
                gender_results[gender] = comp
                legal_acc = compute_legal_accuracy(comp, r_legal)
            else:
                gender_results[gender] = np.nan
                legal_acc = np.nan

            row[f"compensation_{gender}"] = gender_results[gender]
            row[f"legal_accuracy_{gender}"] = legal_acc

        # 计算跨性别指标
        r_male   = gender_results.get("男", np.nan)
        r_female = gender_results.get("女", np.nan)

        if not np.isnan(r_male) and not np.isnan(r_female):
            row["delta"]       = compute_bias_delta(r_male, r_female)
            row["consistency"] = compute_consistency(r_male, r_female)
        else:
            row["delta"]       = np.nan
            row["consistency"] = np.nan

        records.append(row)

    return pd.DataFrame(records)


# ============================================================
# 消融实验汇总表生成
# ============================================================

def compute_kl_divergence_metric(df: pd.DataFrame, model_name: str) -> float:
    """
    基于评估结果计算该模型男女输出分布的 KL 散度。
    """
    sub = df[df["model"] == model_name].copy()
    male_vals   = sub["compensation_男"].dropna().values
    female_vals = sub["compensation_女"].dropna().values

    if len(male_vals) == 0 or len(female_vals) == 0:
        return float("nan")

    all_vals = np.concatenate([male_vals, female_vals])
    min_v, max_v = all_vals.min(), all_vals.max()
    if min_v == max_v:
        return 0.0

    bins = np.linspace(min_v, max_v, 21)
    p_male,   _ = np.histogram(male_vals,   bins=bins, density=False)
    p_female, _ = np.histogram(female_vals, 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 build_ablation_table(all_results: pd.DataFrame) -> pd.DataFrame:
    """
    生成消融实验汇总对比表。

    列含义：
        model_name:        模型名称
        mean_delta:        平均偏见偏移率（负方向=对女性有利，正=对男性有利）
        mean_abs_delta:    平均|Δ|（衡量偏见幅度，越小越好）
        mean_legal_acc:    平均法律准确度（越大越好，范围[0,1]）
        mean_consistency:  平均一致性分数（越大越好，范围[0,1]）
        kl_divergence:     KL散度D_KL(P_male||P_female)（越小越好）
    """
    rows = []
    for model_name in all_results["model"].unique():
        sub = all_results[all_results["model"] == model_name]
        kl  = compute_kl_divergence_metric(all_results, model_name)

        # 聚合法律准确度（取男女均值）
        legal_acc_m = sub["legal_accuracy_男"].mean()
        legal_acc_f = sub["legal_accuracy_女"].mean()
        mean_legal_acc = (legal_acc_m + legal_acc_f) / 2

        rows.append({
            "model_name":       model_name,
            "mean_delta":       sub["delta"].mean(),
            "mean_abs_delta":   sub["delta"].abs().mean(),
            "mean_legal_acc":   mean_legal_acc,
            "mean_consistency": sub["consistency"].mean(),
            "kl_divergence":    kl,
        })

    df_table = pd.DataFrame(rows)
    # 按偏见幅度升序排列
    df_table = df_table.sort_values("mean_abs_delta").reset_index(drop=True)
    return df_table


def print_ablation_table(df_table: pd.DataFrame) -> None:
    """格式化打印消融实验对比表"""
    logger.info("\n" + "=" * 80)
    logger.info("消融实验对比表（Ablation Study Results）")
    logger.info("=" * 80)
    header = (
        f"{'Model':<25} {'Mean Δ':>10} {'Mean |Δ|':>10} "
        f"{'Legal Acc':>10} {'Consistency':>12} {'KL Div':>10}"
    )
    logger.info(header)
    logger.info("-" * 80)
    for _, row in df_table.iterrows():
        line = (
            f"{row['model_name']:<25} "
            f"{row['mean_delta']:>10.4f} "
            f"{row['mean_abs_delta']:>10.4f} "
            f"{row['mean_legal_acc']:>10.4f} "
            f"{row['mean_consistency']:>12.4f} "
            f"{row['kl_divergence']:>10.4f}"
        )
        logger.info(line)
    logger.info("=" * 80)
    logger.info("说明: Mean |Δ| 越小越好 | Legal Acc 越大越好 | Consistency 越大越好 | KL Div 越小越好")


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

def main() -> None:
    logger.info("=" * 60)
    logger.info("消融实验评估脚本 启动")
    logger.info("=" * 60)

    all_dfs: list[pd.DataFrame] = []

    # 定义四个变体的配置
    eval_configs = [
        {
            "model_name":    "A: Baseline (No Finetune)",
            "model_variant": "Model_A_Baseline",
            "model_path":    config.BASE_MODEL_ID,  # 原始基座模型
            "is_base":       True,
        },
        {
            "model_name":    "B: Data Debiasing Only",
            "model_variant": "Model_B_DataOnly",
            "model_path":    config.ABLATION_MODELS["Model_B_DataOnly"],
            "is_base":       False,
        },
        {
            "model_name":    "C: KL Loss Only",
            "model_variant": "Model_C_LossOnly",
            "model_path":    config.ABLATION_MODELS["Model_C_LossOnly"],
            "is_base":       False,
        },
        {
            "model_name":    "D: Dual Debiasing",
            "model_variant": "Model_D_Dual",
            "model_path":    config.ABLATION_MODELS["Model_D_Dual"],
            "is_base":       False,
        },
    ]

    for cfg in eval_configs:
        model_path = cfg["model_path"]
        use_sim = False
        tokenizer, model = None, None

        # 尝试加载模型
        if cfg["is_base"]:
            tokenizer, model = load_base_model()
        else:
            if Path(model_path).exists():
                tokenizer, model = load_finetuned_model(model_path)
            else:
                logger.warning(
                    f"模型路径不存在: {model_path}\n"
                    f"  → 使用模拟数据评估 {cfg['model_name']}"
                )

        if tokenizer is None or model is None:
            use_sim = True
            logger.info(f"[{cfg['model_name']}] 使用模拟推理数据")

        # 执行评估
        df_result = evaluate_single_model(
            model_name=cfg["model_name"],
            model_variant=cfg["model_variant"],
            tokenizer=tokenizer,
            model=model,
            use_simulation=use_sim,
        )
        all_dfs.append(df_result)

        # 及时释放显存
        if model is not None:
            import torch
            del model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    # 合并结果
    all_results = pd.concat(all_dfs, ignore_index=True)

    # 保存详细评估结果
    all_results.to_csv(config.EVAL_RESULT_CSV, index=False, encoding="utf-8-sig")
    logger.info(f"\n详细评估结果已保存至: {config.EVAL_RESULT_CSV}")

    # 生成并保存消融汇总表
    ablation_table = build_ablation_table(all_results)
    ablation_path  = str(Path(config.RESULTS_DIR) / "ablation_table.csv")
    ablation_table.to_csv(ablation_path, index=False, encoding="utf-8-sig")
    logger.info(f"消融实验对比表已保存至: {ablation_path}")

    # 打印汇总表
    print_ablation_table(ablation_table)

    logger.info("\n评估完成！")


if __name__ == "__main__":
    main()
