analyze.py

"""
analyze — the AI layer for Voice Harvester.

Makes a mixed recording usable for cloning by breaking it into per-speaker,
per-utterance segments you can pick from:

  1. Demucs-isolate the voice (reuses engine.py).
  2. Whisper transcribe with timestamps (local, faster-whisper if available).
  3. Split into segments at natural pauses + transcript boundaries, and tag each
     with a rough "voice signature" (pitch + energy) so different speakers cluster
     into groups (Speaker A / B / C). You confirm which group is the one you want.

So from a clip of mom + dad + you, you get labeled segments and can export just
one person's — exactly what voice cloning needs.

Outputs are JSON so any UI (the cross-platform GUI, CLI, automation) can drive it.
Local + private; nothing leaves the machine.
"""

from __future__ import annotations

import json
import math
import os
import subprocess
import sys
import tempfile
import wave
from dataclasses import dataclass, asdict
from typing import Any, Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import engine  # noqa: E402


@dataclass
class Segment:
    index: int
    start: float
    end: float
    text: str
    pitch_hz: float           # mean fundamental-ish pitch
    energy: float             # RMS loudness 0..1
    speaker: str = "?"        # assigned group label, e.g. "A"
    @property
    def duration(self) -> float:
        return round(self.end - self.start, 2)


# ---------- pitch (autocorrelation, stdlib only) ----------
def _mean_pitch(samples: list[float], sr: int) -> float:
    """Rough mean pitch (Hz) of a voiced chunk via autocorrelation. 0 if unvoiced."""
    n = len(samples)
    if n < sr // 20:
        return 0.0
    # work in 40ms frames, take the median of voiced frames
    frame = int(sr * 0.04)
    hop = frame // 2
    pitches: list[float] = []
    lo = int(sr / 350)   # 350 Hz max
    hi = int(sr / 70)    # 70 Hz min
    i = 0
    while i + frame <= n:
        f = samples[i:i + frame]
        # remove DC
        m = sum(f) / len(f)
        f = [x - m for x in f]
        e = sum(x * x for x in f)
        if e > 1e-3:
            best_lag, best = 0, 0.0
            for lag in range(lo, min(hi, frame - 1)):
                s = 0.0
                for k in range(0, frame - lag, 2):  # stride 2 for speed
                    s += f[k] * f[k + lag]
                if s > best:
                    best, best_lag = s, lag
            if best_lag:
                pitches.append(sr / best_lag)
        i += hop
    if not pitches:
        return 0.0
    pitches.sort()
    return round(pitches[len(pitches) // 2], 1)


def _read_wav_mono(path: str) -> tuple[list[float], int]:
    w = wave.open(path, "rb")
    sr = w.getframerate()
    n = w.getnframes()
    raw = w.readframes(n)
    w.close()
    import array
    a = array.array("h")
    a.frombytes(raw)
    ch = w.getnchannels()
    data = list(a)
    if ch == 2:
        data = [(data[i] + data[i + 1]) / 2 for i in range(0, len(data) - 1, 2)]
    return [x / 32768.0 for x in data], sr


# ---------- transcription ----------
def _transcribe(path: str) -> list[dict[str, Any]]:
    """Return [{start,end,text}] via faster-whisper if available, else a single
    block (so segmentation still works on pitch/pauses alone)."""
    try:
        from faster_whisper import WhisperModel
        model = WhisperModel("base", device="auto", compute_type="int8")
        segs, _ = model.transcribe(path, beam_size=1, vad_filter=True)
        return [{"start": float(s.start), "end": float(s.end), "text": s.text.strip()} for s in segs]
    except Exception:
        return []


# ---------- main analysis ----------
def analyze(src: str, *, use_demucs: Optional[bool] = None,
            log=lambda m: None) -> dict[str, Any]:
    if not engine.have_ffmpeg():
        return {"ok": False, "error": "ffmpeg not found"}
    use_demucs = engine.have_demucs() if use_demucs is None else use_demucs

    with tempfile.TemporaryDirectory() as tmp:
        raw = os.path.join(tmp, "raw.wav")
        log("Extracting audio…")
        engine._extract_audio(src, raw)
        voice = raw
        if use_demucs:
            try:
                log("Isolating voice (Demucs)…")
                voice = engine._isolate_with_demucs(raw, tmp)
            except engine.ProcessingError:
                voice = raw
        # normalize to mono 16k for analysis
        clean = os.path.join(tmp, "clean.wav")
        engine._run(["ffmpeg", "-y", "-i", voice, "-ar", "16000", "-ac", "1",
                     "-c:a", "pcm_s16le", clean])

        log("Transcribing…")
        words = _transcribe(clean)
        data, sr = _read_wav_mono(clean)

        # build segments: from whisper if we have it, else fixed 3s windows
        spans = ([(w["start"], w["end"], w["text"]) for w in words] if words
                 else [(t, min(t + 3, len(data) / sr), "")
                       for t in _frange(0, len(data) / sr, 3)])

        segs: list[Segment] = []
        for i, (st, en, text) in enumerate(spans):
            a, b = int(st * sr), int(en * sr)
            chunk = data[a:b]
            if not chunk:
                continue
            rms = math.sqrt(sum(x * x for x in chunk) / len(chunk))
            segs.append(Segment(i, round(st, 2), round(en, 2), text,
                                _mean_pitch(chunk, sr), round(rms, 4)))

        log("Clustering speakers by voice…")
        _assign_speakers(segs, data, sr)
        return {"ok": True, "source": src, "duration": round(len(data) / sr, 1),
                "transcribed": bool(words),
                "segments": [_seg_json(s) for s in segs],
                "speakers": _speaker_summary(segs)}


def _frange(a: float, b: float, step: float):
    x = a
    while x < b:
        yield x
        x += step


def _assign_speakers(segs: list[Segment], data: list[float] | None = None,
                     sr: int = 16000) -> None:
    """Group segments by *speaker* using voice timbre, not just pitch.

    Pitch alone merges people who happen to speak in the same range (a mother and
    a daughter both read as "female-range"). Instead we build a small timbre
    fingerprint per segment — log-energy across mel-ish frequency bands plus pitch
    — and cluster those. Two speakers with the same pitch but different vocal
    tracts still separate, because their spectral shape (formants) differs.

    Falls back to the old pitch-band split when we have no audio to fingerprint.
    Transparent and dependency-light: numpy is used if present (faster), else a
    pure-Python path. """
    voiced = [s for s in segs if s.pitch_hz > 0]
    if not voiced:
        return
    if data is None:
        _assign_speakers_by_pitch(segs)
        return

    # 1) fingerprint each voiced segment (spectral-band energies + pitch)
    feats: list[list[float]] = []
    for s in voiced:
        a, b = int(s.start * sr), int(s.end * sr)
        feats.append(_voice_fingerprint(data[a:b], sr, s.pitch_hz))

    # 2) how many speakers? estimate from feature spread, capped at 3.
    k = _estimate_speakers(feats)
    if k <= 1:
        for s in voiced:
            s.speaker = "A"
        return

    # 3) cluster (k-means on standardized features), label by size (A=biggest)
    labels = _kmeans(feats, k)
    if labels is None:                     # clustering unavailable → pitch fallback
        _assign_speakers_by_pitch(segs)
        return
    order = _labels_by_frequency(labels)   # remap so A is the most common speaker
    for s, lab in zip(voiced, labels):
        s.speaker = order[lab]


def _assign_speakers_by_pitch(segs: list[Segment]) -> None:
    """The original transparent fallback: split voiced segments into up to 3
    pitch bands (A/B/C). Used when audio fingerprinting isn't available."""
    voiced = [s for s in segs if s.pitch_hz > 0]
    if not voiced:
        return
    pitches = sorted(s.pitch_hz for s in voiced)
    lo, hi = pitches[0], pitches[-1]
    if hi - lo < 25:
        for s in voiced:
            s.speaker = "A"
        return
    t1 = lo + (hi - lo) / 3
    t2 = lo + 2 * (hi - lo) / 3
    for s in voiced:
        s.speaker = "A" if s.pitch_hz < t1 else ("B" if s.pitch_hz < t2 else "C")


def _voice_fingerprint(samples: list[float], sr: int, pitch_hz: float) -> list[float]:
    """A small timbre vector: log energy in a handful of frequency bands (a cheap
    stand-in for MFCC formant structure) plus normalized pitch. Averaged over the
    segment's frames so it characterizes the *voice*, not the words."""
    if not samples:
        return [0.0] * 9
    frame = _next_pow2(int(sr * 0.032))     # ~32ms FFT frame
    hop = frame // 2
    # mel-ish band edges (Hz) across the vocal range
    edges = [80, 200, 400, 700, 1100, 1700, 2600, 4000, 6000]
    bands = [0.0] * (len(edges) - 1)
    nframes = 0
    i = 0
    while i + frame <= len(samples):
        mag = _fft_mag(samples[i:i + frame])
        # map FFT bins to Hz and accumulate energy per band
        for bi in range(len(bands)):
            f_lo, f_hi = edges[bi], edges[bi + 1]
            k_lo = max(1, int(f_lo * frame / sr))
            k_hi = min(len(mag) - 1, int(f_hi * frame / sr))
            e = sum(mag[k] for k in range(k_lo, k_hi + 1))
            bands[bi] += e
        nframes += 1
        i += hop
    if nframes:
        bands = [b / nframes for b in bands]
    # log-compress (perceptual) and append normalized pitch
    vec = [math.log(b + 1e-6) for b in bands]
    vec.append(pitch_hz / 200.0)
    return vec


def _next_pow2(n: int) -> int:
    p = 1
    while p < n:
        p <<= 1
    return p


def _fft_mag(frame: list[float]) -> list[float]:
    """Magnitude spectrum. Uses numpy if present (fast), else a compact recursive
    FFT. Applies a Hann window first to reduce leakage."""
    n = len(frame)
    win = [0.5 - 0.5 * math.cos(2 * math.pi * i / (n - 1)) for i in range(n)]
    x = [frame[i] * win[i] for i in range(n)]
    try:
        import numpy as np
        return list(np.abs(np.fft.rfft(np.asarray(x))))
    except Exception:
        spec = _fft(x)
        half = n // 2 + 1
        return [abs(spec[k]) for k in range(half)]


def _fft(x: list[float]) -> list[complex]:
    """Iterative radix-2 Cooley–Tukey FFT (n must be a power of 2)."""
    n = len(x)
    if n <= 1:
        return [complex(v) for v in x]
    a = [complex(v) for v in x]
    # bit-reversal permutation
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    length = 2
    while length <= n:
        ang = -2j * math.pi / length
        wlen = complex(math.cos(ang.imag), math.sin(ang.imag))
        for i in range(0, n, length):
            w = 1 + 0j
            for k in range(length // 2):
                u = a[i + k]
                v = a[i + k + length // 2] * w
                a[i + k] = u + v
                a[i + k + length // 2] = u - v
                w *= wlen
        length <<= 1
    return a


def _estimate_speakers(feats: list[list[float]]) -> int:
    """Guess speaker count (1–3) from how spread out the fingerprints are. If the
    features form one tight blob → 1 speaker; clear separation → 2 or 3."""
    if len(feats) < 4:
        return 1
    # standardized spread: mean pairwise distance vs. within-tolerance
    import statistics
    cols = list(zip(*feats))
    means = [statistics.fmean(c) for c in cols]
    stds = [statistics.pstdev(c) or 1.0 for c in cols]
    z = [[(v - means[j]) / stds[j] for j, v in enumerate(f)] for f in feats]
    # silhouette-ish: try k=2 and k=3, keep whichever separates clearly
    best_k = 1
    for k in (2, 3):
        labels = _kmeans(z, k)
        if labels is None or len(set(labels)) < k:
            continue
        if _cluster_separation(z, labels) > 1.15:   # inter/intra ratio threshold
            best_k = k
    return best_k


def _cluster_separation(feats: list[list[float]], labels: list[int]) -> float:
    """Ratio of between-cluster to within-cluster spread. >1 means real clusters."""
    import statistics
    groups: dict[int, list[list[float]]] = {}
    for f, l in zip(feats, labels):
        groups.setdefault(l, []).append(f)
    centroids = {l: [statistics.fmean(c) for c in zip(*g)] for l, g in groups.items()}
    overall = [statistics.fmean(c) for c in zip(*feats)]
    within = between = 0.0
    for l, g in groups.items():
        for f in g:
            within += _dist(f, centroids[l])
        between += len(g) * _dist(centroids[l], overall)
    within = within / len(feats) or 1e-6
    between = between / len(feats)
    return between / within


def _dist(a: list[float], b: list[float]) -> float:
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


def _kmeans(feats: list[list[float]], k: int, iters: int = 25) -> list[int] | None:
    """Standardized k-means. numpy if available, else pure-Python. Deterministic
    init (spread-out seeds) so results are stable run-to-run."""
    if len(feats) < k:
        return None
    import statistics
    cols = list(zip(*feats))
    means = [statistics.fmean(c) for c in cols]
    stds = [statistics.pstdev(c) or 1.0 for c in cols]
    z = [[(v - means[j]) / stds[j] for j, v in enumerate(f)] for f in feats]
    # k-means++ style deterministic seeds: first point, then farthest each time
    centers = [z[0]]
    while len(centers) < k:
        far, far_d = z[0], -1.0
        for p in z:
            d = min(_dist(p, c) for c in centers)
            if d > far_d:
                far, far_d = p, d
        centers.append(far)
    labels = [0] * len(z)
    for _ in range(iters):
        changed = False
        for i, p in enumerate(z):
            best, bd = 0, float("inf")
            for ci, c in enumerate(centers):
                d = _dist(p, c)
                if d < bd:
                    best, bd = ci, d
            if labels[i] != best:
                labels[i] = best
                changed = True
        for ci in range(k):
            members = [z[i] for i in range(len(z)) if labels[i] == ci]
            if members:
                centers[ci] = [statistics.fmean(c) for c in zip(*members)]
        if not changed:
            break
    return labels


def _labels_by_frequency(labels: list[int]) -> dict[int, str]:
    """Map raw cluster ids → A/B/C so 'A' is always the most-talked speaker."""
    from collections import Counter
    counts = Counter(labels).most_common()
    names = "ABCDEFG"
    return {lab: names[i] for i, (lab, _n) in enumerate(counts)}


def _speaker_summary(segs: list[Segment]) -> list[dict[str, Any]]:
    out: dict[str, dict[str, Any]] = {}
    for s in segs:
        if s.speaker == "?":
            continue
        d = out.setdefault(s.speaker, {"speaker": s.speaker, "segments": 0,
                                       "seconds": 0.0, "pitches": []})
        d["segments"] += 1
        d["seconds"] = round(d["seconds"] + s.duration, 1)
        d["pitches"].append(s.pitch_hz)
    for d in out.values():
        ps = sorted(d.pop("pitches"))
        d["avg_pitch"] = round(sum(ps) / len(ps), 0) if ps else 0
        # a friendly hint based on typical ranges
        d["likely"] = ("lower / male-range" if d["avg_pitch"] < 145
                       else "higher / female-range" if d["avg_pitch"] > 175
                       else "mid-range")
    return sorted(out.values(), key=lambda d: -d["seconds"])


def _seg_json(s: Segment) -> dict[str, Any]:
    j = asdict(s)
    j["duration"] = s.duration
    return j


def export_speaker(src: str, speaker: str, out_path: str, *,
                   use_demucs: Optional[bool] = None,
                   pick_indices: Optional[list[int]] = None,
                   log=lambda m: None) -> dict[str, Any]:
    """Re-extract + concatenate just the chosen speaker's (or chosen segments')
    audio into one clean wav, ready for cloning."""
    res = analyze(src, use_demucs=use_demucs, log=log)
    if not res.get("ok"):
        return res
    segs = res["segments"]
    chosen = [s for s in segs
              if (pick_indices is not None and s["index"] in pick_indices)
              or (pick_indices is None and s["speaker"] == speaker)]
    if not chosen:
        return {"ok": False, "error": "no segments matched"}

    with tempfile.TemporaryDirectory() as tmp:
        raw = os.path.join(tmp, "raw.wav")
        engine._extract_audio(src, raw)
        voice = raw
        if (engine.have_demucs() if use_demucs is None else use_demucs):
            try:
                voice = engine._isolate_with_demucs(raw, tmp)
            except engine.ProcessingError:
                voice = raw
        parts = []
        for i, s in enumerate(chosen):
            part = os.path.join(tmp, f"p{i}.wav")
            engine._run(["ffmpeg", "-y", "-i", voice, "-ss", str(s["start"]),
                         "-to", str(s["end"]), "-ar", "22050", "-ac", "1",
                         "-c:a", "pcm_s16le", part])
            parts.append(part)
        engine.merge_wavs(parts, out_path)
        # final clean-up pass (silence strip + normalize)
        tmp_out = out_path + ".tmp.wav"
        engine._clean_normalize(out_path, tmp_out)
        os.replace(tmp_out, out_path)
    dur = engine.probe_duration(out_path)
    return {"ok": True, "out": out_path, "segments": len(chosen),
            "duration": round(dur, 1)}


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("--export-speaker")
    ap.add_argument("--out", default="speaker.wav")
    a = ap.parse_args()
    if a.export_speaker:
        print(json.dumps(export_speaker(a.src, a.export_speaker, a.out, log=lambda m: print(m, file=sys.stderr)), indent=2))
    else:
        print(json.dumps(analyze(a.src, log=lambda m: print(m, file=sys.stderr)), indent=2))