#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""稳健人声基频(F0)统计：自相关 + 周期性门限 + 八度纠错 + 能量门限。
输出中位 F0、分位数与估计的歌唱音域。用于"男声/女声"的客观对比（同一把尺子前后测量）。
用法: f0-stats.py <vocals.wav> <标签>
"""
from __future__ import annotations
import sys
import numpy as np
import soundfile as sf

FMIN, FMAX = 80.0, 600.0     # 人声基频搜索范围
WIN, HOP = 1024, 320          # 16kHz 下约 64ms 窗 / 20ms 跳
CORR_MIN = 0.45               # 周期性门限（清音/噪声会被滤掉）
RMS_FLOOR = 0.02              # 相对全局 RMS 的能量门限


def frame_f0(x: np.ndarray, sr: int):
    win = np.hanning(WIN).astype(np.float32)
    lag_min, lag_max = int(sr / FMAX), int(sr / FMIN)
    global_rms = float(np.sqrt((x ** 2).mean())) + 1e-12
    out = []
    for start in range(0, len(x) - WIN, HOP):
        fr = x[start:start + WIN]
        r = float(np.sqrt((fr ** 2).mean()))
        if r < RMS_FLOOR * global_rms:
            continue
        fr = (fr - fr.mean()) * win
        n = 1
        while n < 2 * WIN:
            n *= 2
        spec = np.fft.rfft(fr, n)
        ac = np.fft.irfft(spec * np.conj(spec), n)[:WIN]
        if ac[0] <= 0:
            continue
        ac = ac / ac[0]
        seg = ac[lag_min:lag_max + 1]
        if seg.size == 0:
            continue
        lag = int(np.argmax(seg)) + lag_min
        peak = float(ac[lag])
        if peak < CORR_MIN:
            continue
        # 八度纠错：若 2 倍 lag 仍在范围内且相关性接近，则取更低频（真基频）
        for mult in (2, 3):
            l2 = lag * mult
            if l2 <= lag_max and ac[l2] > 0.9 * peak:
                lag, peak = l2, float(ac[l2])
        out.append(sr / lag)
    return np.array(out)


def main():
    path, tag = sys.argv[1], sys.argv[2]
    x, sr = sf.read(path, always_2d=True)
    mono = x.mean(axis=1).astype(np.float32)
    if sr != 16000:  # 简单线性重采样（只用于基频估计）
        idx = np.linspace(0, len(mono) - 1, int(len(mono) * 16000 / sr))
        mono = np.interp(idx, np.arange(len(mono)), mono).astype(np.float32)
        sr = 16000
    f0 = frame_f0(mono, sr)
    if f0.size < 20:
        print(f"{tag}: 有效清音帧过少({f0.size})，无法判定")
        return
    q = np.percentile(f0, [5, 10, 25, 50, 75, 90, 95])
    names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
    def nname(f):
        m = int(round(69 + 12 * np.log2(f / 440.0)))
        return f"{names[m % 12]}{m // 12 - 1}"
    print(f"{tag}: 有效清音帧 {f0.size}")
    print(f"  F0 分位 5/25/50/75/95 = {q[0]:.0f} / {q[2]:.0f} / {q[3]:.0f} / {q[4]:.0f} / {q[6]:.0f} Hz")
    print(f"  中位 F0 = {q[3]:.0f} Hz ≈ {nname(q[3])}")
    print(f"  歌唱音域(5-95%) = {nname(q[0])} ~ {nname(q[6])}")
    lo, hi = 100, 155
    verdict = "女声区" if q[3] >= 175 else ("男声区" if q[3] <= 155 else "中间区")
    print(f"  判定(中位): {verdict}   [男声典型 100-155Hz / 女声典型 175-260Hz / 童声更高]")


if __name__ == "__main__":
    main()