# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Analyse the lab: variance vs separation, and verify a suspicious F0 reading."""
from __future__ import annotations

import json
from pathlib import Path

import librosa
import numpy as np
import soundfile as sf

LAB = Path("/private/var/folders/y8/cbvhds297918lrxt2wd3btsr0000gn/T/opencode/author_lab")
res = json.loads((LAB / "results.json").read_text())
rows = res["rows"]
REF_F0 = res["reference"]["f0_median"]
REF_ENV = res["reference"]["envelope_hz"]
print(f"reference: F0 {REF_F0} Hz  envelope {REF_ENV} Hz\n")


def st(v, base):
    return 12 * np.log2(v / base) if v and base else float("nan")


# ---------------------------------------------------------------- 1. variance
print("=== INTRA-VOICE VARIANCE (same voice, different seed/temperature) ===")
D = [r for r in rows if r["label"].startswith("D_")]
f0s = np.array([r["f0_median"] for r in D], dtype=float)
envs = np.array([r["envelope_hz"] for r in D], dtype=float)
print(f"n={len(D)}  seeds 11-55 x temp 0.3/0.7/0.9")
print(f"  F0       {f0s.min():.1f} - {f0s.max():.1f} Hz   spread {f0s.max()-f0s.min():.1f} Hz "
      f"= {st(f0s.max(), REF_F0)-st(f0s.min(), REF_F0):+.1f} st   sd {f0s.std():.1f} Hz")
print(f"  envelope {envs.min()} - {envs.max()} Hz      spread "
      f"{st(envs.max(), REF_ENV)-st(envs.min(), REF_ENV):+.1f} st   sd {envs.std():.0f} Hz")
F0_SD_ST = f0s.std() / REF_F0 * 12 * np.log2(np.e)  # semitone sd
ENV_SD_ST = envs.std() / REF_ENV * 12 * np.log2(np.e)
print(f"  => 1-sd noise floor: F0 {F0_SD_ST:.2f} st, envelope {ENV_SD_ST:.2f} st")
print(f"  => 2-sd (where a variant becomes reliably 'not just a re-roll'): "
      f"F0 {2*F0_SD_ST:.2f} st")

# ---------------------------------------------------------------- 2. separation
print("\n=== AXIS B: did the neural instruct move anything on a REAL HUMAN reference? ===")
print(f"{'variant':20} {'d_F0_st':>9} {'d_env_st':>10} {'WER':>6}")
for r in rows:
    if r["label"].startswith("B_") or r["label"].startswith("A00"):
        print(f"{r['label']:20} {r.get('d_f0_st'):>9} {r.get('d_env_st'):>10} {r.get('wer'):>6}")

# ---------------------------------------------------------------- 3. verify F0 artifact
print("\n=== VERIFYING C_monster_grit_full (F0 69.7 Hz but WER 0.04) ===")


def true_f0(y, sr):
    """Autocorrelation F0 over 60-400 Hz - independent of pyin's tracking."""
    y = y.astype(np.float32)
    f0, _, _ = librosa.pyin(y, fmin=60.0, fmax=400.0, sr=sr, frame_length=1024)
    v = f0[~np.isnan(f0)]
    ac = librosa.autocorrelate(y, max_size=int(sr / 60))
    lo, hi = int(sr / 400), int(sr / 60)
    peak = lo + int(np.argmax(ac[lo:hi]))
    return float(np.median(v)), sr / peak


for label in ("A00_baseline", "C_monster_grit_ring", "C_monster_grit_full",
              "C_monster_pitch-10"):
    p = LAB / "wav" / f"{label}.wav"
    if not p.exists():
        continue
    y, sr = sf.read(str(p), dtype="float32")
    if y.ndim > 1:
        y = y.mean(axis=1)
    pyin_f0, ac_f0 = true_f0(y, sr)
    # strong spectral peaks
    S = np.abs(librosa.stft(y, n_fft=4096, hop_length=1024))
    f = librosa.fft_frequencies(sr=sr, n_fft=4096)
    m = (f > 40) & (f < 500)
    amp = S[m].mean(axis=1)
    peaks = f[m][np.argsort(amp)[-4:]][::-1]
    print(f"  {label:24} pyin={pyin_f0:6.1f}  autocorr={ac_f0:6.1f} Hz  "
          f"low-freq peaks={[round(float(x)) for x in peaks]}")

# ---------------------------------------------------------------- 4. cast headroom
print("\n=== CAST HEADROOM on the pitch axis ===")
A = [r for r in rows if r["label"].startswith("A_pitch") and r.get("f0_median")]
sep = np.array([abs(st(r["f0_median"], REF_F0)) for r in A])
print(f"  reliable separation needs > 2-sd = {2*F0_SD_ST:.1f} st")
print(f"  usable DSP variants found (|dF0| > 2sd): "
      f"{sorted(round(float(x),1) for x in sep if x > 2*F0_SD_ST)}")
usable = [r for r in A if abs(st(r["f0_median"], REF_F0)) > 2 * F0_SD_ST]
lo = min(st(r["f0_median"], REF_F0) for r in usable)
hi = max(st(r["f0_median"], REF_F0) for r in usable)
step = 2 * F0_SD_ST
print(f"  span {lo:+.1f}..{hi:+.1f} st  /  {step:.1f} st per voice  =>  "
      f"~{int((hi-lo)//step)+1} reliably distinct voices on ONE axis")