#!/usr/bin/env python3
"""比较 csv 里多种来源的 position / score 的 F1 指标。

读取 golden_set_with_estimated.csv（带 estimated_position 列），
对以下 4 套 position / score 各算：
  - 默认阈值 F1（rank<106 / score>=44.72）
  - 扫阈值最优 F1
  - AUC

来源：
  1. estimated_position           （embedding top-100 mean rank）
  2. position_gpt_4.1 / score_gpt_4.1
  3. position_gemini_2.5_flash / score_gemini_2.5_flash

用法：
  python3 compare_positions.py
  python3 compare_positions.py --csv /path/to/golden_set_with_estimated.csv
"""
import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd


DEFAULT_CSV = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.csv"


def metrics(preds, gts):
    tp = int(((preds == 1) & (gts == 1)).sum())
    fp = int(((preds == 1) & (gts == 0)).sum())
    tn = int(((preds == 0) & (gts == 0)).sum())
    fn = int(((preds == 0) & (gts == 1)).sum())
    p = tp/(tp+fp) if tp+fp else 0.0
    r = tp/(tp+fn) if tp+fn else 0.0
    f = 2*p*r/(p+r) if p+r else 0.0
    a = (tp+tn)/len(preds)
    return tp, fp, tn, fn, p, r, f, a


def best_threshold(scores, gts):
    """scores 越大越严重时调用；返回 (f1, thr, p, r)"""
    valid = ~np.isnan(scores)
    s, g = scores[valid], gts[valid]
    cands = sorted(set(s.tolist()))
    best = (-1.0, None, None, None)
    for c in cands:
        preds = (s >= c).astype(int)
        _, _, _, _, p, r, f, _ = metrics(preds, g)
        if f > best[0]:
            best = (f, c, p, r)
    return best


def auc(scores, gts):
    valid = ~np.isnan(scores)
    s, g = scores[valid], gts[valid]
    if len(set(g.tolist())) < 2:
        return float("nan")
    try:
        from sklearn.metrics import roc_auc_score
        return roc_auc_score(g, s)
    except ImportError:
        return float("nan")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--csv", default="golden_set_with_estimated.csv")
    p.add_argument("--positive-label", default="Y")
    p.add_argument("--boundary-rank",  type=int,   default=106)
    p.add_argument("--boundary-score", type=float, default=44.72)
    p.add_argument("--output-md", default="position_comparison.md")
    return p.parse_args()


def main():
    args = parse_args()
    print(f"[1] read csv: {args.csv}")
    df = pd.read_csv(args.csv, keep_default_na=True)
    print(f"      {len(df)} rows, cols: {len(df.columns)}")
    print(f"      label dist: {df['label'].value_counts().to_dict()}")

    gts = df["label"].astype(str).str.upper().eq(args.positive_label.upper()).astype(int).values

    # 4 套预测的"分数"（统一成"越大越严重"以便扫阈值）
    sources = []  # (name, score_array, default_threshold, default_pred_fn)
    if "estimated_position" in df.columns:
        pos = pd.to_numeric(df["estimated_position"], errors="coerce").values
        sources.append((
            "embedding estimated_position",
            -pos,                               # 取负 → 越大越严重
            -args.boundary_rank,                # 默认阈值也取负
            "rank<106",
        ))
    if "position_gpt_4.1" in df.columns:
        pos = pd.to_numeric(df["position_gpt_4.1"], errors="coerce").values
        sources.append(("LLM gpt_4.1 position", -pos, -args.boundary_rank, "rank<106"))
    if "score_gpt_4.1" in df.columns:
        sc = pd.to_numeric(df["score_gpt_4.1"], errors="coerce").values
        sources.append(("LLM gpt_4.1 score", sc, args.boundary_score, "score>=44.72"))
    if "position_gemini_2.5_flash" in df.columns:
        pos = pd.to_numeric(df["position_gemini_2.5_flash"], errors="coerce").values
        sources.append(("LLM gemini_2.5_flash position", -pos, -args.boundary_rank, "rank<106"))
    if "score_gemini_2.5_flash" in df.columns:
        sc = pd.to_numeric(df["score_gemini_2.5_flash"], errors="coerce").values
        sources.append(("LLM gemini_2.5_flash score", sc, args.boundary_score, "score>=44.72"))

    if not sources:
        raise SystemExit("没有可用的 position / score 列")

    # ---- 算指标 ----
    rows = []
    for name, s, default_thr, default_desc in sources:
        # 默认阈值
        valid = ~np.isnan(s)
        preds = np.zeros_like(s, dtype=int)
        preds[valid] = (s[valid] >= default_thr).astype(int)
        tp, fp, tn, fn, prec, rec, f1, acc = metrics(preds, gts)

        # 扫最优
        best_f1, best_thr, best_p, best_r = best_threshold(s, gts)
        a = auc(s, gts)

        rows.append({
            "source":     name,
            "default":    default_desc,
            "n_valid":    int(valid.sum()),
            "TP":         tp, "FP": fp, "TN": tn, "FN": fn,
            "P":          round(prec, 4),
            "R":          round(rec,  4),
            "F1":         round(f1,   4),
            "Acc":        round(acc,  4),
            "best F1":    round(best_f1, 4),
            "best thr":   best_thr,
            "AUC":        round(a, 4) if not np.isnan(a) else None,
        })

    out_df = pd.DataFrame(rows)
    print("\n" + out_df.to_string(index=False))

    # ---- 写 markdown ----
    with open(args.output_md, "w", encoding="utf-8") as f:
        f.write("# Position / Score 来源对比\n\n")
        f.write(f"数据：`{args.csv}` ({len(df)} rows)\n\n")
        f.write(f"GT：`label == '{args.positive_label}'` → 1\n\n")
        f.write(f"默认阈值：rank<{args.boundary_rank} 或 score≥{args.boundary_score}\n\n")
        f.write("## 完整对比\n\n")
        # 不依赖 tabulate，手写 markdown 表格
        cols = list(out_df.columns)
        f.write("| " + " | ".join(cols) + " |\n")
        f.write("|" + "|".join(["---"] * len(cols)) + "|\n")
        for _, row in out_df.iterrows():
            f.write("| " + " | ".join(str(v) for v in row.tolist()) + " |\n")
        f.write("\n\n## 关键观察\n\n")
        # F1 排名
        sorted_rows = sorted(rows, key=lambda r: -r["F1"])
        f.write(f"- 默认阈值下 F1 最高：**{sorted_rows[0]['source']}** ({sorted_rows[0]['F1']})\n")
        f.write(f"- 默认阈值下 F1 最低：**{sorted_rows[-1]['source']}** ({sorted_rows[-1]['F1']})\n")
        # AUC 排名
        valid_rows = [r for r in rows if r["AUC"] is not None]
        if valid_rows:
            sorted_rows2 = sorted(valid_rows, key=lambda r: -r["AUC"])
            f.write(f"- AUC 最高（分布可分性）：**{sorted_rows2[0]['source']}** ({sorted_rows2[0]['AUC']})\n")
        # embedding vs LLM gap
        emb = next((r for r in rows if r["source"].startswith("embedding")), None)
        gem = next((r for r in rows if "gemini" in r["source"] and "score" in r["source"]), None)
        if emb and gem:
            f.write(f"- embedding 估的 F1 比 Gemini 低 **{(gem['best F1']-emb['best F1']):.4f}**（best F1 对比）\n")

    print(f"\n[written] {args.output_md}")


if __name__ == "__main__":
    main()