"""Find the lead-vocal pitch line in a stereo mix, using three cues that need no training:

  1. centre:   the lead vocal is panned dead centre, so keep time-frequency bins where L ≈ R
  2. vibrato:  sung notes are frequency-modulated at 4–8 Hz; instruments aren't (pianoify.vibrato_score)
  3. range:    a singer stays in one register, so frames an octave away from the local register
               are folded back (and pitches outside 75–1100 Hz are rejected)
Also estimates the voice type from the register.
"""
import os
import sys

import numpy as np
from scipy.ndimage import median_filter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "pianoify"))
import pianoify as PZ  # noqa: E402

SR = 16000


def centre_signal(g, F, power=80):
    L, R = F.fwd(g[0]), F.fwd(g[1])
    psi = np.abs(L + R) ** 2 / (2 * (np.abs(L) ** 2 + np.abs(R) ** 2) + 1e-12)
    return F.adj((psi ** power) * (L + R) / 2)


def fold_to_register(hz, win_frames=200, thresh=10 / 12):
    """Range consistency: move each pitch by whole octaves toward the running median of the line."""
    out = hz.copy()
    v = hz > 0
    if v.sum() < 10:
        return out
    m = np.log2(np.where(v, hz, np.nan))
    med = np.copy(m)
    # running median over voiced frames only (nan-aware, via index trick)
    idx = np.where(v)[0]
    mv = median_filter(m[idx], size=min(win_frames, len(idx)) | 1, mode="nearest")
    med[idx] = mv
    d = med[idx] - m[idx]                               # octaves away from the register
    # only fold frames that are nearly a whole octave off: real melodies can span an octave,
    # so anything within ~10 semitones of the register is left alone
    k = np.where(np.abs(d) > thresh, np.round(d), 0)
    out[idx] = hz[idx] * 2.0 ** np.clip(k, -2, 2)
    return out


def voice_type(hz):
    v = hz[hz > 0]
    if not len(v):
        return "unknown", (0, 0)
    lo, med, hi = np.percentile(v, [5, 50, 95])
    # typical sung ranges (Hz): bass 82–330, baritone 98–392, tenor 130–523, alto 175–700, soprano 262–1047
    names = [("bass", 82, 330), ("baritone", 98, 392), ("tenor", 130, 523), ("alto", 175, 700), ("soprano", 262, 1047)]
    best = min(names, key=lambda n: abs(np.log2(med) - np.log2(np.sqrt(n[1] * n[2]))))
    return best[0], (float(lo), float(hi))


def track(g, F, centre_power=80, vibrato=3.0, voicing=0.03, fold=True):
    import librosa
    c = centre_signal(g, F, centre_power) if centre_power else g.mean(axis=0)
    y = librosa.resample(c.astype(np.float64), orig_sr=SR, target_sr=PZ.SR)
    t, midi, _ = PZ.melody_contour(y, vibrato=vibrato, voicing=voicing)
    hz = np.where(np.isfinite(midi), 440 * 2 ** ((midi - 69) / 12), 0.0)
    hz = hz[np.clip(np.searchsorted(t, F.frame_times()), 0, len(t) - 1)]
    hz = np.where((hz >= 75) & (hz <= 1100), hz, 0.0)  # outside any human singing range → not the voice
    if fold:
        hz = fold_to_register(hz)
    return hz