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

"""Phase-insensitive fidelity: how faithfully does the matrix reproduce the reference?

Frame-wise waveform correlation and LSD are dominated by phase, which a 12 Hz
neural codec does not preserve. Identity lives in the long-term average spectrum
(tonal balance), the pitch contour, and the mel envelope — all phase-insensitive.
"""
from __future__ import annotations

from pathlib import Path

import librosa
import numpy as np
import soundfile as sf

OUT = Path("/private/var/folders/y8/cbvhds297918lrxt2wd3btsr0000gn/T/opencode/probe_out")
A4 = Path.home() / "Documents/Audacity4"
dec, sr = sf.read(str(OUT / "author_ref_DECODED_FROM_MATRIX.wav"), dtype="float32")
src, _ = sf.read(str(A4 / "BDMurphy_Voice_NOBG_v2.wav"), dtype="float32")
if src.ndim > 1:
    src = src.mean(axis=1)

BANDS = [(80, 160), (160, 315), (315, 630), (630, 1250), (1250, 2500),
         (2500, 5000), (5000, 10000), (10000, 12000)]


def ltas(y):
    S = np.abs(librosa.stft(y.astype(np.float32), n_fft=2048, hop_length=512))
    f = librosa.fft_frequencies(sr=sr, n_fft=2048)
    out = []
    for lo, hi in BANDS:
        m = (f >= lo) & (f < hi)
        out.append(10 * np.log10((S[m] ** 2).sum() + 1e-12))
    return np.array(out)


def mel_cep(y):
    m = librosa.feature.melspectrogram(y=y.astype(np.float32), sr=sr, n_mels=40,
                                        n_fft=1024, hop_length=256)
    return np.log10(m + 1e-10)


def f0_track(y):
    f0, _, _ = librosa.pyin(y.astype(np.float32), fmin=60.0, fmax=400.0, sr=sr)
    return f0


print(f"{'':28} {'decoded':>9} {'reference':>10} {'delta':>8}")
print("--- long-term average spectrum, per band (dB, energy-normalised) ---")
a, b = ltas(dec), ltas(src)
a_n = a - a.mean()
b_n = b - b.mean()
for (lo, hi), x, y in zip(BANDS, a_n, b_n):
    print(f"{lo:>6}-{hi:<6} Hz{'':10} {x:9.2f} {y:10.2f} {x - y:+8.2f}")
print(f"{'RMS deviation':28} {'':9} {'':10} {float(np.sqrt((a_n - b_n) ** 2).mean()):+8.2f} dB")

print("\n--- mel envelope (40 bands, log) ---")
ma, mb = mel_cep(dec), mel_cep(src)
n = min(ma.shape[1], mb.shape[1])
ma_n = ma[:, :n] - ma[:, :n].mean(axis=0, keepdims=True)
mb_n = mb[:, :n] - mb[:, :n].mean(axis=0, keepdims=True)
print(f"mean |log-mel diff| = {np.abs(ma_n - mb_n).mean():.4f}  (log10 units)")

print("\n--- pitch contour ---")
fa, fb = f0_track(dec), f0_track(src)
va, vb = fa[~np.isnan(fa)], fb[~np.isnan(fb)]
print(f"median F0: decoded {np.median(va):.1f} Hz   reference {np.median(vb):.1f} Hz")
print(f"voiced frames: decoded {len(va)}  reference {len(vb)}")

# The real fidelity test: re-encode the DECODED audio. If the codec round-trip is
# stable, decoding twice converges, i.e. matrix -> wav -> codes -> same codes.
import sys

sys.path.insert(0, "/Volumes/1TB_APL_Storage/UmbraAudio")
print("\n(note: full re-encode check runs in matrix_roundtrip.py)")