"""Clean a voice-reference clip for Chatterbox cloning.

Takes a raw recording (mp3/wav, any length), finds the cleanest ~15s of
sustained speech, spectrally denoises the constant room air/hiss, high-passes
the rumble, normalizes, and writes a crisp 24kHz mono voices/jarvis_ref.wav
that Chatterbox clones from.

    python tools/clean_voice_ref.py "some_interview.mp3"
    python tools/clean_voice_ref.py in.wav --out voices/jarvis_ref.wav --seconds 15

Deps: soundfile, soxr, noisereduce, scipy (all in requirements-win.txt).
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import soundfile as sf
import soxr
import noisereduce as nr
from scipy.signal import butter, sosfilt

PROJECT = Path(__file__).resolve().parent.parent


def _frames(x: np.ndarray, frame: int) -> np.ndarray:
    n = len(x) // frame * frame
    return x[:n].reshape(-1, frame)


def clean(src: str, out: str, seconds: float = 15.0) -> None:
    audio, sr = sf.read(src, dtype="float32")
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if sr != 48000:
        audio = soxr.resample(audio, sr, 48000)
        sr = 48000

    frame = int(0.03 * sr)
    # skip likely intro/outro, then pick the cleanest sustained-speech window
    head = min(int(15 * sr), len(audio) // 4)
    tail = min(int(10 * sr), len(audio) // 4)
    a = audio[head: len(audio) - tail] if len(audio) > head + tail else audio
    win = int(seconds * sr)
    hop = int(1 * sr)
    best, best_s = 0, -1.0
    for st in range(0, max(1, len(a) - win), hop):
        seg = a[st:st + win]
        rms = np.sqrt((_frames(seg, frame) ** 2).mean(axis=1))
        voiced = (rms > 0.03).mean()
        floor = np.percentile(rms, 10)
        steady = 1 - min(1, rms.std() / max(rms.mean(), 1e-6) / 3)
        s = voiced * 0.6 + steady * 0.25 - floor * 4
        if s > best_s:
            best_s, best = s, st
    seg = a[best:best + win]

    # noise profile = the quietest ~0.8s (pure room air)
    rms = np.sqrt((_frames(seg, frame) ** 2).mean(axis=1))
    qi = np.argsort(rms)[:max(1, int(0.8 * sr // frame))]
    noise = np.concatenate([_frames(seg, frame)[i] for i in qi])

    # SINGLE, gentle pass — enough to lift the room air without the "musical
    # noise" hollowing that a heavier/two-pass gate leaves (that made the
    # cloned voice sound muffled/"chubby"). A touch of natural air > artifacts.
    clean_sig = nr.reduce_noise(y=seg, sr=sr, y_noise=noise,
                                stationary=True, prop_decrease=0.92)
    clean_sig = sosfilt(butter(4, 75, "highpass", fs=sr, output="sos"),
                        clean_sig).astype(np.float32)
    clean_sig = sosfilt(butter(4, 11000, "lowpass", fs=sr, output="sos"),
                        clean_sig).astype(np.float32)
    clean_sig /= max(np.abs(clean_sig).max(), 1e-6) / 0.9

    out_24k = soxr.resample(clean_sig, sr, 24000).astype("float32")
    Path(out).parent.mkdir(parents=True, exist_ok=True)
    sf.write(out, out_24k, 24000)
    print(f"picked {best / sr:.0f}s window; wrote {out} "
          f"({len(out_24k) / 24000:.0f}s, 24kHz mono, denoised)")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("--out", default=str(PROJECT / "voices" / "jarvis_ref.wav"))
    ap.add_argument("--seconds", type=float, default=15.0)
    args = ap.parse_args()
    clean(args.src, args.out, args.seconds)