"""
05_visualization.py - 基于真实实验数据生成论文图表

运行方式：
    cd 实验代码目录
    python 05_visualization.py

输出（保存至 figures/ 目录）：
    figures/fig1_heatmap_bias.png   - 50场景×4模型偏差热力图
    figures/fig2_delta_barplot.png  - 各模型 mean|Δ| 条形图
    figures/fig3_scenario_dist.png  - 各模型偏差方向分布
    figures/fig4_compliance.png     - 各模型与法定标准的偏差
    figures/stats_summary.txt       - 统计摘要
"""

import sys
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns
from pathlib import Path
from scipy import stats

sys.path.insert(0, str(Path(__file__).parent))
import config

matplotlib.rcParams["font.family"] = "DejaVu Sans"
matplotlib.rcParams["axes.unicode_minus"] = False
matplotlib.rcParams["font.size"] = 11

FIGURES_DIR = config.FIGURES_DIR
DATA_DIR = config.DATA_DIR

MODEL_DISPLAY = {
    "qwen3.6-plus":         "Qwen 3.6 Plus",
    "deepseek-v4-flash":    "DeepSeek V4 Flash",
    "gemini-3.1-flash-lite":"Gemini 3.1 Flash Lite",
    "claude-sonnet-4-6":    "Claude Sonnet 4.6",
}

MODEL_ORDER = ["qwen3.6-plus", "claude-sonnet-4-6", "gemini-3.1-flash-lite", "deepseek-v4-flash"]
MODEL_COLORS = {
    "qwen3.6-plus":          "#E15759",
    "claude-sonnet-4-6":     "#F28E2B",
    "gemini-3.1-flash-lite": "#76B7B2",
    "deepseek-v4-flash":     "#4E79A7",
}


def load_data():
    summary_path = DATA_DIR / "bias_detection_results_summary.csv"
    raw_path     = DATA_DIR / "bias_detection_results.csv"
    summary = pd.read_csv(summary_path, encoding="utf-8-sig")
    raw     = pd.read_csv(raw_path,     encoding="utf-8-sig")
    return summary, raw


# ── Figure 1: Heatmap ────────────────────────────────────────────────────────

def plot_heatmap(summary: pd.DataFrame, save_path: Path):
    scenarios = [f"S{i:02d}" for i in range(1, 51)]
    models    = MODEL_ORDER

    matrix = pd.DataFrame(index=[MODEL_DISPLAY[m] for m in models], columns=scenarios, dtype=float)
    for m in models:
        sub = summary[summary["model"] == m].set_index("scenario_id")
        for s in scenarios:
            matrix.loc[MODEL_DISPLAY[m], s] = sub.loc[s, "delta"] if s in sub.index else np.nan

    fig, ax = plt.subplots(figsize=(22, 3.8))
    sns.heatmap(
        matrix.astype(float),
        ax=ax,
        cmap="RdBu_r",
        center=0,
        vmin=-1.0, vmax=1.0,
        annot=False,
        linewidths=0.3,
        linecolor="white",
        cbar_kws={"label": "Bias Shift Rate Δ\n(+ = male favored, − = female favored)",
                  "shrink": 0.8},
    )
    ax.set_xlabel("Case Scenario (S01–S50)", fontsize=12)
    ax.set_ylabel("")
    ax.set_title(
        "Figure 1: Gender Bias Shift Rate (Δ) Across 50 Scenarios and 4 LLMs\n"
        "Δ = (R_male − R_female) / Mean(R_male, R_female)   |   Red = male-favoring, Blue = female-favoring",
        fontsize=11, pad=10,
    )
    plt.xticks(rotation=90, fontsize=7)
    plt.yticks(rotation=0, fontsize=11)
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


# ── Figure 2: Mean |Δ| bar chart ─────────────────────────────────────────────

def plot_delta_barplot(summary: pd.DataFrame, save_path: Path):
    rows = []
    for m in MODEL_ORDER:
        sub = summary[summary["model"] == m]["delta"].dropna()
        rows.append({
            "model":      MODEL_DISPLAY[m],
            "mean_abs":   sub.abs().mean(),
            "mean_signed":sub.mean(),
            "std":        sub.std(),
            "n_biased":   (sub != 0).sum(),
            "n_total":    len(sub),
        })
    df = pd.DataFrame(rows)

    fig, ax = plt.subplots(figsize=(9, 5.5))
    x = np.arange(len(df))
    bars = ax.bar(
        x, df["mean_abs"],
        yerr=df["std"] / np.sqrt(df["n_total"]),
        color=[MODEL_COLORS[m] for m in MODEL_ORDER],
        edgecolor="white", linewidth=0.8,
        error_kw={"elinewidth": 1.5, "capsize": 6, "ecolor": "black"},
        width=0.55,
    )
    ax.set_xticks(x)
    ax.set_xticklabels(df["model"], fontsize=10.5)
    ax.set_ylabel("Mean |Bias Shift Rate| (mean|Δ|)", fontsize=12)
    ax.set_title(
        "Figure 2: Gender Bias Magnitude per LLM\n"
        "(Higher mean|Δ| = greater gender disparity in compensation recommendations)",
        fontsize=11,
    )
    ax.set_ylim(0, 0.38)
    for bar, row in zip(bars, df.itertuples()):
        pct = row.n_biased / row.n_total * 100
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.007,
            f"|Δ|={row.mean_abs:.4f}\n{pct:.0f}% biased",
            ha="center", va="bottom", fontsize=9.5,
        )
    ax.axhline(0.05, color="gray", linestyle="--", linewidth=1, label="Reference line (|Δ|=0.05)")
    ax.legend(fontsize=9)
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


# ── Figure 3: Direction breakdown ───────────────────────────────────────────

def plot_direction(summary: pd.DataFrame, save_path: Path):
    rows = []
    for m in MODEL_ORDER:
        sub = summary[summary["model"] == m]["delta"].dropna()
        n_male   = (sub > 0).sum()
        n_equal  = (sub == 0).sum()
        n_female = (sub < 0).sum()
        total    = len(sub)
        rows.append({
            "model":   MODEL_DISPLAY[m],
            "male":    n_male   / total * 100,
            "equal":   n_equal  / total * 100,
            "female":  n_female / total * 100,
        })
    df = pd.DataFrame(rows)

    fig, ax = plt.subplots(figsize=(10, 4.5))
    x = np.arange(len(df))
    w = 0.25
    ax.bar(x - w, df["male"],   width=w, color="#E15759", label="Male-favoring (Δ > 0)", alpha=0.85)
    ax.bar(x,     df["equal"],  width=w, color="#BAB0AC", label="Gender-neutral (Δ = 0)", alpha=0.85)
    ax.bar(x + w, df["female"], width=w, color="#4E79A7", label="Female-favoring (Δ < 0)", alpha=0.85)
    ax.set_xticks(x)
    ax.set_xticklabels(df["model"], fontsize=10.5)
    ax.set_ylabel("Proportion of 50 Scenarios (%)", fontsize=12)
    ax.set_title(
        "Figure 3: Direction of Gender Bias Across Scenarios per LLM\n"
        "(Each bar shows % of 50 scenarios where the model favored male / neither / female)",
        fontsize=11,
    )
    ax.set_ylim(0, 70)
    ax.legend(fontsize=10)
    for bars_group, col in zip([df["male"], df["equal"], df["female"]], [x - w, x, x + w]):
        for xi, val in zip(col, bars_group):
            if val > 3:
                ax.text(xi, val + 0.8, f"{val:.0f}%", ha="center", va="bottom", fontsize=9)
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


# ── Figure 4: Compensation vs Legal Standard ────────────────────────────────

def plot_compliance(summary: pd.DataFrame, save_path: Path):
    fig, axes = plt.subplots(2, 2, figsize=(13, 10))
    axes = axes.flatten()

    for ax, m in zip(axes, MODEL_ORDER):
        sub = summary[summary["model"] == m].dropna(subset=["R_male", "R_female", "legal_compensation"])
        legal = sub["legal_compensation"].values
        r_m   = sub["R_male"].values
        r_f   = sub["R_female"].values
        legal_k = legal / 1000
        r_m_k   = r_m   / 1000
        r_f_k   = r_f   / 1000

        ax.scatter(legal_k, r_m_k, color="#E15759", label="Male plaintiff",   alpha=0.75, s=55, zorder=3)
        ax.scatter(legal_k, r_f_k, color="#4E79A7", label="Female plaintiff", alpha=0.75, s=55, marker="s", zorder=3)
        lo = min(legal_k.min(), r_m_k.min(), r_f_k.min()) * 0.9
        hi = max(legal_k.max(), r_m_k.max(), r_f_k.max()) * 1.05
        ax.plot([lo, hi], [lo, hi], "k--", linewidth=1.2, label="Legal standard (1:1)", zorder=2)

        ax.set_xlabel("Legal Compensation (× 1,000 CNY)", fontsize=10)
        ax.set_ylabel("Suggested Amount (× 1,000 CNY)", fontsize=10)
        ax.set_title(MODEL_DISPLAY[m], fontsize=11)
        ax.legend(fontsize=8.5)
        ax.grid(alpha=0.25)

    fig.suptitle(
        "Figure 4: Suggested vs. Legal Compensation — Male vs. Female Plaintiffs (All 50 Scenarios)\n"
        "Points above/below the dashed line indicate over/under-compensation relative to the legal standard",
        fontsize=11, y=1.01,
    )
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


# ── Statistics summary ───────────────────────────────────────────────────────

def print_stats(summary: pd.DataFrame, save_path: Path):
    lines = ["=" * 70, "GENDER BIAS DETECTION EXPERIMENT — STATISTICAL SUMMARY", "=" * 70, ""]
    for m in MODEL_ORDER:
        sub = summary[summary["model"] == m]["delta"].dropna()
        n_biased = (sub != 0).sum()
        n_male   = (sub > 0).sum()
        n_female = (sub < 0).sum()
        stat, pval = stats.wilcoxon(sub) if (sub != 0).any() else (0, 1.0)
        lines += [
            f"Model: {MODEL_DISPLAY[m]}",
            f"  N scenarios   : {len(sub)}",
            f"  Mean Δ (signed): {sub.mean():+.4f}",
            f"  Mean |Δ|       : {sub.abs().mean():.4f}",
            f"  Std dev Δ      : {sub.std():.4f}",
            f"  Max Δ          : {sub.max():+.4f}",
            f"  Min Δ          : {sub.min():+.4f}",
            f"  Biased (Δ≠0)  : {n_biased}/50  ({n_biased/50*100:.0f}%)",
            f"    Male-favoring: {n_male}  Female-favoring: {n_female}",
            f"  Wilcoxon W     : {stat:.1f}  p = {pval:.4f}",
            "",
        ]
    text = "\n".join(lines)
    print(text)
    save_path.write_text(text, encoding="utf-8")
    print(f"  Saved: {save_path}")


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("Generating paper figures from real experimental data...")
    print("=" * 60)

    summary, raw = load_data()
    print(f"Loaded {len(summary)} rows from summary CSV")
    print(f"Models: {summary['model'].unique().tolist()}")

    print("\nGenerating Figure 1: Heatmap...")
    plot_heatmap(summary, FIGURES_DIR / "fig1_heatmap_bias.png")

    print("Generating Figure 2: Mean|Δ| Bar Chart...")
    plot_delta_barplot(summary, FIGURES_DIR / "fig2_delta_barplot.png")

    print("Generating Figure 3: Direction Breakdown...")
    plot_direction(summary, FIGURES_DIR / "fig3_direction_breakdown.png")

    print("Generating Figure 4: Compliance Scatter...")
    plot_compliance(summary, FIGURES_DIR / "fig4_compliance_scatter.png")

    print("Computing statistics...")
    print_stats(summary, FIGURES_DIR / "stats_summary.txt")

    print("Generating Figure 5: FCFT Ablation Comparison...")
    ablation_path = Path(config.RESULTS_DIR) / "ablation_table.csv"
    if ablation_path.exists():
        plot_ablation(ablation_path, FIGURES_DIR / "fig5_ablation_comparison.png")
    else:
        print(f"  Skipped (file not found: {ablation_path})")

    print("Generating Figure 6: CSC Before/After Comparison...")
    csc_path = Path(config.DATA_DIR) / "csc_analysis_results.csv"
    if csc_path.exists():
        plot_csc(csc_path, FIGURES_DIR / "fig6_csc_comparison.png")
    else:
        print(f"  Skipped (file not found: {csc_path})")

    print("\n" + "=" * 60)
    print(f"All outputs saved to: {FIGURES_DIR}")
    print("=" * 60)


def plot_ablation(csv_path: Path, save_path: Path):
    df = pd.read_csv(csv_path, encoding="utf-8-sig")

    # 固定顺序：A→C→D→B（偏见从高到低）
    order = [
        "A: Baseline (No Finetune)",
        "C: KL Loss Only",
        "D: Dual Debiasing",
        "B: Data Debiasing Only",
    ]
    labels = ["A: Baseline", "C: KL Loss Only", "D: FCFT\n(Dual)", "B: Data\nDebiasing"]
    colors = ["#E15759", "#F28E2B", "#4E79A7", "#76B7B2"]

    df = df.set_index("model_name").reindex(order)
    df["mean_legal_acc"] = pd.to_numeric(df["mean_legal_acc"], errors="coerce")

    fig, axes = plt.subplots(1, 3, figsize=(13, 5))

    metrics = [
        ("mean_abs_delta",  "Mean |Δ| (Bias)",          colors, "lower is better"),
        ("mean_legal_acc",  "Legal Accuracy",            colors, "higher is better"),
        ("kl_divergence",   "KL Divergence",             colors, "lower is better"),
    ]

    for ax, (col, title, clrs, note) in zip(axes, metrics):
        vals = df[col].values
        bars = ax.bar(labels, vals, color=clrs, edgecolor="white", linewidth=0.8, width=0.6)

        # 数值标签
        for bar, v in zip(bars, vals):
            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + max(vals) * 0.02,
                    f"{v:.4f}" if col == "kl_divergence" else f"{v:.3f}",
                    ha="center", va="bottom", fontsize=9.5, fontweight="bold")

        ax.set_title(title, fontsize=12, fontweight="bold", pad=8)
        ax.set_ylabel(note, fontsize=9, color="gray")
        ax.set_ylim(0, max(vals) * 1.18)
        ax.grid(axis="y", alpha=0.25, linestyle="--")
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.tick_params(axis="x", labelsize=9)

    fig.suptitle(
        "Figure 5: FCFT Ablation Study — Four-Variant Comparison\n"
        "(Base model: Qwen1.5-1.8B-Chat | GPU: RTX 5070 Ti | 50 scenarios)",
        fontsize=11, y=1.02,
    )
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


def plot_csc(csv_path: Path, save_path: Path):
    df = pd.read_csv(csv_path, encoding="utf-8-sig")

    model_order = ["qwen3.6-plus", "claude-sonnet-4-6", "gemini-3.1-flash-lite", "deepseek-v4-flash"]
    display = {
        "qwen3.6-plus":          "Qwen 3.6 Plus",
        "claude-sonnet-4-6":     "Claude Sonnet 4.6",
        "gemini-3.1-flash-lite": "Gemini 3.1 Flash Lite",
        "deepseek-v4-flash":     "DeepSeek V4 Flash",
    }

    # 从 pivot 数据中聚合每个模型的指标
    rows = []
    for model in model_order:
        sub = df[df["model"] == model]
        if sub.empty:
            continue
        orig_comply = sub["orig_correct"].mean() * 100
        csc_comply  = sub["csc_correct"].mean() * 100
        orig_abs    = sub["abs_delta"].mean()
        csc_abs     = sub.loc[~sub["csc_triggered"], "abs_delta"].sum() / len(sub)
        rows.append({
            "model":        display[model],
            "orig_comply":  orig_comply,
            "csc_comply":   csc_comply,
            "orig_abs":     orig_abs,
            "csc_abs":      csc_abs,
        })
    data = pd.DataFrame(rows)

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    x = np.arange(len(data))
    w = 0.35
    colors_before = "#E15759"
    colors_after  = "#4E79A7"

    # 左图：法律符合率对比
    ax = axes[0]
    ax.bar(x - w/2, data["orig_comply"], w, label="Before CSC", color=colors_before, alpha=0.85)
    ax.bar(x + w/2, data["csc_comply"],  w, label="After CSC",  color=colors_after,  alpha=0.85)
    for i, (b, a) in enumerate(zip(data["orig_comply"], data["csc_comply"])):
        ax.text(i - w/2, b + 1.5, f"{b:.0f}%", ha="center", va="bottom", fontsize=9, color=colors_before, fontweight="bold")
        ax.text(i + w/2, a + 1.5, f"{a:.0f}%", ha="center", va="bottom", fontsize=9, color=colors_after,  fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(data["model"], fontsize=9, rotation=10, ha="right")
    ax.set_ylabel("Legal Compliance Rate (%)", fontsize=10)
    ax.set_title("Legal Compliance: Before vs. After CSC", fontsize=11, fontweight="bold")
    ax.set_ylim(0, 110)
    ax.legend(fontsize=10)
    ax.grid(axis="y", alpha=0.25, linestyle="--")
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    # 右图：均值|Δ|对比
    ax = axes[1]
    ax.bar(x - w/2, data["orig_abs"], w, label="Before CSC", color=colors_before, alpha=0.85)
    ax.bar(x + w/2, data["csc_abs"],  w, label="After CSC",  color=colors_after,  alpha=0.85)
    for i, (b, a) in enumerate(zip(data["orig_abs"], data["csc_abs"])):
        ax.text(i - w/2, b + 0.003, f"{b:.3f}", ha="center", va="bottom", fontsize=9, color=colors_before, fontweight="bold")
        ax.text(i + w/2, a + 0.003, f"{a:.3f}", ha="center", va="bottom", fontsize=9, color=colors_after,  fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(data["model"], fontsize=9, rotation=10, ha="right")
    ax.set_ylabel("Mean |Δ| (Bias)", fontsize=10)
    ax.set_title("Gender Bias: Before vs. After CSC", fontsize=11, fontweight="bold")
    ax.set_ylim(0, data["orig_abs"].max() * 1.25)
    ax.legend(fontsize=10)
    ax.grid(axis="y", alpha=0.25, linestyle="--")
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    fig.suptitle(
        "Figure 6: CSC Post-Processing Effect — Before vs. After Calibration (τ = 0.05)\n"
        "Left: legal compliance rate  |  Right: mean gender bias |Δ|",
        fontsize=11, y=1.02,
    )
    plt.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {save_path}")


if __name__ == "__main__":
    main()
