#!/usr/bin/env python3
"""把 embedding top-K match 估出来的整数 position 写回 golden_set.csv。

默认列名固定为 estimated_position（下游 AIPF warm-start 代码读这个名）。
默认输入/输出都指向 AIPF 流水线实际读取的那个 csv，**直接原地覆盖**，
省掉 `cp` 到 example/yss_ruler_eval/data/ 这一步。

每次跑都覆盖 estimated_position 列；用 --k 切换不同邻居数即可，
其它列保持不变。

用法：
  python3 add_estimated_position.py --k 100   # 默认 K
  python3 add_estimated_position.py --k 5
  python3 add_estimated_position.py --k 1

⚠️ 第一次跑前建议备份原 csv：
  cp /mnt/.../aipf_golden_set.csv /mnt/.../aipf_golden_set.csv.bak

可选：想同时存多列对比（不覆盖默认列），手动指定列名：
  python3 add_estimated_position.py --k 5 --new-col estimated_position_top5
"""
import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd


DEFAULTS = dict(
    csv    = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.csv",
    jsonl  = "golden_top100.jsonl",
    # 输出直接覆盖 AIPF 流水线实际读取的那个 csv，省掉 cp 那一步
    output = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.csv",
    id_col = "task_id",
)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--csv",      default=DEFAULTS["csv"])
    p.add_argument("--jsonl",    default=DEFAULTS["jsonl"])
    p.add_argument("--output",   default=DEFAULTS["output"])
    p.add_argument("--id-col",   default=DEFAULTS["id_col"])
    p.add_argument("--k",        type=int, default=100,
                   help="用前 K 个邻居算均值（K=1 即 top1 rank）")
    p.add_argument("--new-col",  default=None,
                   help="新列名；不传则自动用 estimated_position_top{K}")
    p.add_argument("--k-suffix", default=None,
                   help="（已弃用，等价于 --new-col estimated_position_{suffix}）")
    return p.parse_args()


def main():
    args = parse_args()
    K = args.k

    # 默认列名固定为 estimated_position（让下游代码不用改）
    # 想要带 K 后缀（同时存多列对比）就显式传 --new-col 或 --k-suffix
    if args.new_col:
        col_name = args.new_col
    elif args.k_suffix:
        col_name = f"estimated_position_{args.k_suffix}"
    else:
        col_name = "estimated_position"

    # 1) 读 jsonl，每条样本算 top-K rank 均值（K=1 时等价于最近邻 rank）
    print(f"[1] read jsonl: {args.jsonl}")
    estimated = {}
    skipped = 0
    with open(args.jsonl, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            r = json.loads(line)
            topk = r.get("topk") or []
            if len(topk) < K:
                skipped += 1
                continue
            ranks = np.array([t["rank"] for t in topk[:K]], dtype=float)
            estimated[str(r["task_id"])] = int(round(ranks.mean()))
    print(f"      -> {len(estimated)} cases estimated (K={K})")
    if skipped:
        print(f"      ⚠️ {skipped} 条 jsonl 邻居数 < K，跳过")

    # 2) 读 csv
    print(f"[2] read csv: {args.csv}")
    df = pd.read_csv(args.csv, keep_default_na=False)
    print(f"      -> {len(df)} rows, {len(df.columns)} cols")

    if args.id_col not in df.columns:
        raise ValueError(f"'{args.id_col}' not in csv. cols: {list(df.columns)}")

    # 3) 给每条 case 加新列
    print(f"[3] add column: {col_name}")
    df[col_name] = df[args.id_col].astype(str).map(estimated)

    n_total = len(df)
    n_matched = df[col_name].notna().sum()
    print(f"      matched: {n_matched}/{n_total}")
    if n_matched == n_total:
        df[col_name] = df[col_name].astype("Int64")

    # 4) 写
    print(f"[4] write: {args.output}")
    df.to_csv(args.output, index=False)
    print(f"      done. csv 现在有 {len(df.columns)} 列（含新列：{col_name}）")

    # 5) 预览 + stats
    print(f"\n[preview] 前 5 行（task_id / label / {col_name}）:")
    cols = [args.id_col, "label", col_name]
    cols = [c for c in cols if c in df.columns]
    print(df[cols].head(5).to_string(index=False))

    s = pd.to_numeric(df[col_name], errors="coerce")
    print(f"\n[stats] {col_name}:")
    print(f"  min={int(s.min())}  max={int(s.max())}  mean={s.mean():.2f}  median={int(s.median())}  std={s.std():.2f}")
    if "label" in df.columns:
        print("\n按 label 分组:")
        for lab, sub in df.groupby("label"):
            ss = pd.to_numeric(sub[col_name], errors="coerce")
            print(f"  label={lab!r:>4}  n={len(sub):>4}  mean={ss.mean():.2f}  median={int(ss.median())}  std={ss.std():.2f}")


if __name__ == "__main__":
    main()