#!/usr/bin/env python3
"""
Phase 2: 加载预提取特征 → 秒级搜索昔涟。
前提: 先跑完 extract_features.py
用法: python search_cyrene_v2.py
"""
import os, sys, shutil, json, time, logging
import numpy as np

FEAT_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\features"
OUT_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cyrene_round2"
FEAT_FILE = os.path.join(FEAT_DIR, "features_all.npz")

# ── 用户确认的昔涟样本 (53 个) ──
CONFIRMED = [
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0358_02ff22a9.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0369_0315d90a.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0370_0317b5bf.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0391_033cbeea.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0392_033d9e2e.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0398_0346fa30.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0417_0367cbc5.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0330_02cb864d.wav",
    "External_del_3.5_chapter_2/External_del_3.5_chapter_2_0357_02fd4ab1.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0106_009dfe25.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0107_009fba53.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0112_00a89eb2.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0124_00c08c68.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0128_00c6ba19.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0365_026cf0de.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0371_0278f79f.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0374_02803ab5.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0409_02bcd547.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0410_02bec9f4.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0437_02e57ff0.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0440_02e7d98f.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0467_031daa9e.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0102_0095ba46.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0109_00a458cc.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0126_00c44990.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0130_00c95298.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0359_02670f55.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0364_026b8023.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0366_026f3922.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0369_0274d62a.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0372_027b67ef.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0392_029e92a6.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0413_02c1a418.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0415_02c466f1.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0417_02c8ffc0.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0430_02dbad24.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0431_02dda165.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0434_02e216cf.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0443_02eb0bc4.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0455_03048430.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0462_0316bcd9.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0463_03181c56.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0464_03195142.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0468_031e8bd6.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0469_03208c8a.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0484_03339e50.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0488_033a7664.wav",
    "External_del_3.4_chapter_0/External_del_3.4_chapter_0_0489_033b5c44.wav",
]


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    LOG_FILE = os.path.join(OUT_DIR, "search.log")
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s  %(message)s", datefmt="%H:%M:%S",
        handlers=[logging.FileHandler(LOG_FILE, encoding='utf-8'), logging.StreamHandler(sys.stdout)],
    )
    log = logging.getLogger("search")

    log.info("Phase 2: Cyrene Voice Search (Round 2)")
    log.info(f"  refs : {len(CONFIRMED)} confirmed samples")
    log.info(f"  input: {FEAT_FILE}")

    # 加载特征
    t0 = time.time()
    data = np.load(FEAT_FILE, allow_pickle=True)
    feats = data["feats"]
    paths = data["paths"]
    log.info(f"  loaded: {len(feats):,} features ({feats.nbytes/1024/1024:.0f} MB) in {time.time()-t0:.1f}s")

    # 构建参考模板
    ref_indices = []
    for i, p in enumerate(paths):
        for cf in CONFIRMED:
            if p.endswith(cf.replace("/", os.sep)):
                ref_indices.append(i)
                break

    log.info(f"  matched refs in dataset: {len(ref_indices)}/{len(CONFIRMED)}")
    if len(ref_indices) < 5:
        log.error("  too few refs matched, check paths!")
        return

    template = np.mean(feats[ref_indices], axis=0)
    log.info(f"  template pitch: {template[0]:.0f}Hz  dim: {len(template)}")

    # 全量比对 (向量化, 秒级)
    t1 = time.time()
    norm_feats = feats / (np.linalg.norm(feats, axis=1, keepdims=True) + 1e-8)
    norm_template = template / (np.linalg.norm(template) + 1e-8)
    sims = np.dot(norm_feats, norm_template)

    pitches = feats[:, 0]
    penalty = 1.0 / (1.0 + np.abs(pitches - template[0]) / 100)
    combos = sims * 0.6 + penalty * 0.4
    log.info(f"  compared {len(combos):,} vectors in {time.time()-t1:.1f}s")

    # 排序
    order = np.argsort(-combos)
    results = [(combos[i], sims[i], pitches[i], paths[i]) for i in order]
    log.info(f"  sorted in {time.time()-t1:.1f}s")

    # ── 输出 ──
    log.info(f"\n{'='*55}")
    log.info(f"Top 50 Candidates  (round 2)")
    log.info("=" * 55)
    for rank, (combo, sim, pitch, path) in enumerate(results[:50], 1):
        star = " *" if combo > 0.90 else ""
        d = os.path.basename(os.path.dirname(path))
        f = os.path.basename(path)
        log.info(f"  {rank:2d}. [{combo:.4f}]{star} {d}/{f}")

    # 来源分布
    log.info(f"\nSource distribution (combo > 0.85):")
    srcs = {}
    for combo, sim, pitch, path in results:
        if combo > 0.85:
            d = os.path.basename(os.path.dirname(path))
            srcs[d] = srcs.get(d, 0) + 1
    for d in sorted(srcs):
        log.info(f"  {d}: {srcs[d]}")

    # 分级导出
    tiers = [
        ("tier1_095_100", 0.95),
        ("tier2_092_095", 0.92),
        ("tier3_090_092", 0.90),
        ("tier4_085_090", 0.85),
    ]
    total_copied = 0
    for tier_name, threshold in tiers:
        tier_dir = os.path.join(OUT_DIR, tier_name)
        os.makedirs(tier_dir, exist_ok=True)
        n = 0
        for combo, sim, pitch, path in results:
            if combo >= threshold:
                dst = os.path.join(tier_dir, os.path.basename(path))
                if os.path.exists(path) and not os.path.exists(dst):
                    shutil.copy2(path, dst)
                    n += 1
            else:
                break  # results are sorted, stop when below threshold
        log.info(f"  {tier_name}: {n} files")
        total_copied += n

    # 保存结果
    rp = os.path.join(OUT_DIR, "results.json")
    with open(rp, 'w') as f:
        json.dump([(float(s), float(p), w) for s, _, p, w in results], f, ensure_ascii=False)

    log.info(f"\n  total exported: {total_copied}")
    log.info(f"  results JSON : {rp}")
    log.info(f"  log file     : {LOG_FILE}")
    log.info(f"  DONE ({time.time()-t0:.1f}s)")


if __name__ == "__main__":
    main()
