#!/usr/bin/env python3
"""修复评测结果：移除同歌手且时长相近的负样本（它们实际上是同一母带的不同分发）。

判定条件（与 download_from_db.py 保持一致）：
    sample_class == negative
    AND singer_name == original_singer
    AND |cover_duration - main_duration| / main_duration < 0.05

用法：
python results/fix_same_singer_samples.py \
    --input  results/dna_eval.csv \
    --queries downloads_db/queries.csv \
    --output  results/dna_eval_fixed.csv \
    --threshold 0.05
"""

import argparse
import csv
import json
import re
from pathlib import Path

DURATION_THRESHOLD = 0.05


def _norm_singer(s: str) -> str:
    """归一化歌手名：去除标点/空格，便于跨平台名字格式比较。"""
    return re.sub(r"[\s,&\.。·\-]+", "", s).lower()


def _same_singer(singer: str, original: str) -> bool:
    sn, on = _norm_singer(singer), _norm_singer(original)
    return sn == on or sn in on or on in sn


def load_main_durations(queries_csv: str) -> dict[str, float]:
    """从 queries.csv 中读取主版本（variant=self）的时长，返回 {song_id: duration}。"""
    durations: dict[str, float] = {}
    with open(queries_csv, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            if row.get("variant") == "self" and row.get("duration"):
                try:
                    durations[str(row["song_id"])] = float(row["duration"])
                except (ValueError, TypeError):
                    pass
    return durations


def should_remove(row: dict, main_durations: dict[str, float], threshold: float) -> bool:
    """判断该行是否应从负样本中移除。"""
    if row.get("sample_class") != "negative":
        return False
    singer = (row.get("singer_name") or "").strip()
    original = (row.get("original_singer") or "").strip()
    if not singer or not original or singer != original:
        return False
    # 同歌手 → 再看时长
    song_id = str(row.get("query_song_id") or row.get("song_id") or "")
    main_dur = main_durations.get(song_id)
    if main_dur is None:
        return False  # 没有参照时长，保守保留
    # 从 audio_path 中取文件名，尝试获取 cover 时长（results CSV 没有duration列，用variant作后备）
    # 实际上 results CSV 没有duration，用 metadata.csv 中匹配或直接信任同歌手即可
    # 这里 fallback：同歌手且没有 cover duration 信息，就根据 variant 判断
    # 由于 results CSV 里没有 cover duration，从 queries.csv 查对应行
    return True  # 同歌手 → 标记为待移除（时长将在外层用完整 queries 数据判断）


def compute_summary(rows: list[dict], threshold: float, out_path: str) -> dict:
    """计算评测指标。"""
    tp = fp = tn = fn = 0
    upload_ms, poll_ms, total_ms = [], [], []
    by_variant: dict[str, dict] = {}

    for r in rows:
        pred = r.get("predicted_duplicate", "").lower() == "true"
        exp = r.get("expected_duplicate", "").lower() == "true"
        correct = r.get("correct", "").lower() == "true"
        sample_class = r.get("sample_class", "unknown")

        variant = r.get("variant", "unknown")
        if variant not in by_variant:
            by_variant[variant] = {"correct": 0, "total": 0, "sample_class": sample_class}
        by_variant[variant]["total"] += 1
        if correct:
            by_variant[variant]["correct"] += 1

        if exp and pred:
            tp += 1
        elif not exp and pred:
            fp += 1
        elif not exp and not pred:
            tn += 1
        elif exp and not pred:
            fn += 1

        try:
            upload_ms.append(float(r["upload_ms"]))
        except (KeyError, ValueError, TypeError):
            pass
        try:
            poll_ms.append(float(r["poll_ms"]))
        except (KeyError, ValueError, TypeError):
            pass
        try:
            total_ms.append(float(r["total_ms"]))
        except (KeyError, ValueError, TypeError):
            pass

    total = len(rows)
    accuracy = round((tp + tn) / total, 4) if total else 0
    precision = round(tp / (tp + fp), 4) if (tp + fp) else 0
    recall = round(tp / (tp + fn), 4) if (tp + fn) else 0
    f1 = round(2 * precision * recall / (precision + recall), 4) if (precision + recall) else 0

    def _stats(vals):
        if not vals:
            return {"avg": 0, "p50": 0, "p95": 0, "p99": 0, "min": 0, "max": 0}
        s = sorted(vals)
        n = len(s)
        return {
            "avg": round(sum(s) / n, 1),
            "p50": round(s[int(n * 0.50)], 1),
            "p95": round(s[min(int(n * 0.95), n - 1)], 1),
            "p99": round(s[min(int(n * 0.99), n - 1)], 1),
            "min": round(s[0], 1),
            "max": round(s[-1], 1),
        }

    return {
        "total": total,
        "duplicate_threshold": threshold,
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "tp": tp, "fp": fp, "tn": tn, "fn": fn,
        "query_time_ms": _stats(total_ms),
        "upload_time_ms": _stats(upload_ms),
        "poll_time_ms": _stats(poll_ms),
        "by_sample_class": {
            cls: {
                "by_variant": {
                    v: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]}
                    for v, d in sorted(by_variant.items())
                    if d["sample_class"] == cls
                }
            }
            for cls in ("positive", "negative")
        },
        "by_variant": {
            v: {
                "sample_class": d["sample_class"],
                "accuracy": round(d["correct"] / d["total"], 4),
                "total": d["total"],
            }
            for v, d in sorted(by_variant.items())
        },
        "out": out_path,
    }


def main():
    parser = argparse.ArgumentParser(description="修复评测结果：移除同歌手且时长相近的负样本")
    parser.add_argument("--input", default="results/dna_eval.csv")
    parser.add_argument("--queries", default="downloads_db/queries.csv")
    parser.add_argument("--output", default="results/dna_eval_fixed.csv")
    parser.add_argument("--threshold", type=float, default=DURATION_THRESHOLD,
                        help="时长差异容忍比例，默认 0.05（5%%）")
    args = parser.parse_args()

    # 从 queries.csv 建两张表：
    # 1. song_id → 主版本时长
    # 2. (song_id, audio_path) → cover 时长（用于精确判断）
    main_durations: dict[str, float] = {}
    cover_durations: dict[str, float] = {}  # key = audio_path
    with open(args.queries, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            dur_str = row.get("duration", "")
            try:
                dur = float(dur_str) if dur_str else None
            except (ValueError, TypeError):
                dur = None
            if row.get("variant") == "self" and dur is not None:
                main_durations[str(row["song_id"])] = dur
            if dur is not None:
                cover_durations[row.get("audio_path", "")] = dur

    print(f"主版本时长索引: {len(main_durations)} 首")

    # 读原始结果
    with open(args.input, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames
        rows = list(reader)

    print(f"原始行数: {len(rows)}")

    removed = []
    kept = []
    for r in rows:
        if r.get("sample_class") != "negative":
            kept.append(r)
            continue

        singer = (r.get("singer_name") or "").strip()
        original = (r.get("original_singer") or "").strip()
        if not singer or not original or not _same_singer(singer, original):
            kept.append(r)
            continue

        # 同歌手 → 检查时长
        song_id = str(r.get("query_song_id") or "")
        main_dur = main_durations.get(song_id)
        cover_dur = cover_durations.get(r.get("audio_path", ""))

        if main_dur and cover_dur:
            diff_ratio = abs(cover_dur - main_dur) / main_dur
            if diff_ratio < args.threshold:
                removed.append(r)
                continue
        elif main_dur:
            # 没有 cover 时长信息但同歌手，保守移除
            removed.append(r)
            continue

        kept.append(r)

    print(f"移除行数: {len(removed)}（同歌手且时长相近）")
    print(f"保留行数: {len(kept)}")

    # 写修复后的 CSV（去掉冗余列）
    DROP_COLS = {"topk_hit", "expected_duplicate", "predicted_duplicate", "sample_class"}
    out_fieldnames = [f for f in fieldnames if f not in DROP_COLS]
    out_path = Path(args.output)
    with out_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=out_fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(kept)
    print(f"已写入: {out_path}")

    # 重新计算 summary
    summary = compute_summary(kept, threshold=args.threshold, out_path=str(out_path))
    summary_path = out_path.with_suffix(".summary.json")
    with summary_path.open("w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    print(f"已写入: {summary_path}")

    print()
    print("=== 修复后指标 ===")
    print(f"总样本:    {summary['total']}")
    print(f"Accuracy:  {summary['accuracy']}")
    print(f"Precision: {summary['precision']}")
    print(f"Recall:    {summary['recall']}")
    print(f"F1:        {summary['f1']}")
    print(f"TP={summary['tp']} FP={summary['fp']} TN={summary['tn']} FN={summary['fn']}")

    for cls in ("positive", "negative"):
        variants = summary["by_sample_class"][cls]["by_variant"]
        if not variants:
            continue
        print(f"\n[{cls}]")
        for v, d in variants.items():
            print(f"  {v:20s} acc={d['accuracy']:.4f}  n={d['total']}")

    # 打印被移除的样本明细
    if removed:
        print(f"\n移除的 {len(removed)} 条负样本:")
        for r in removed:
            print(f"  song_id={r.get('query_song_id')} variant={r.get('variant')} "
                  f"singer={r.get('singer_name')} correct_was={r.get('correct')}")


if __name__ == "__main__":
    main()