#!/usr/bin/env python3
"""高精度 BPM 实测：onset 包络 + 梳状滤波扫描（分数精度，亚帧分辨率）。

用法: .venv-tts/bin/python tools/bpm_precise.py <audio> [bpm_min] [bpm_max]
输出 JSON: {"bpm": ..., "score": ..., "harmonics": {...}}
"""
import json
import sys

import librosa
import numpy as np

path = sys.argv[1]
bpm_lo = float(sys.argv[2]) if len(sys.argv) > 2 else 70.0
bpm_hi = float(sys.argv[3]) if len(sys.argv) > 3 else 140.0

y, sr = librosa.load(path, sr=22050, mono=True)
onset = librosa.onset.onset_strength(y=y, sr=sr)
onset = onset - onset.mean()
fps = sr / 512.0  # hop 512


def comb_score(bpm):
    period = fps * 60.0 / bpm
    n_teeth = int((len(onset) - 1) / period)
    if n_teeth < 16:
        return 0.0
    # 相位扫描取最优
    best = 0.0
    for phase in np.linspace(0, period, 24, endpoint=False):
        idx = np.round(phase + period * np.arange(n_teeth)).astype(int)
        idx = idx[idx < len(onset)]
        s = float(onset[idx].sum()) / len(idx)
        if s > best:
            best = s
    return best


cands = np.arange(bpm_lo, bpm_hi, 0.05)
scores = np.array([comb_score(b) for b in cands])
best_i = int(np.argmax(scores))
best_bpm = float(cands[best_i])

# 抛物线插值细化
if 0 < best_i < len(cands) - 1:
    y0, y1, y2 = scores[best_i - 1], scores[best_i], scores[best_i + 1]
    denom = (y0 - 2 * y1 + y2)
    if abs(denom) > 1e-12:
        best_bpm += 0.05 * 0.5 * (y0 - y2) / denom

harm = {}
for k, name in ((0.5, "half"), (2.0, "double")):
    b = best_bpm * k
    if bpm_lo <= b <= bpm_hi:
        harm[name] = round(float(scores[np.argmin(np.abs(cands - b))]), 4)

print(json.dumps({
    "file": path,
    "bpm": round(best_bpm, 2),
    "score": round(float(scores[best_i]), 4),
    "harmonics": harm,
}))