# source_parity.py [SRC_DIR] -- authors' RawNetLite.RawNetLite + audio_preprocessor.preprocess_audio + tester.py scoring
# (outputs = model(waveforms).squeeze() = P(fake)) vs this wrapper on 16 J-SPAW_LA clips (+ a short and a long clip).
# Checkpoint loaded with weights_only=True instead of tester.py's plain torch.load. Compares -logit(P_fake) (wrapper
# score) with -logit of the source's sigmoid output, and the source's P(fake) with sigmoid(-wrapper score).
import io, os, sys
import numpy as np, pyarrow.parquet as pq, soundfile as sf, torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import rawnetlite as W
src = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "..", "_ckpt_dl", "gh__adipiz99__RawNetLite")
sys.path.append(src)
import RawNetLite as SRC  # noqa: E402  (pinned reviewed clone)
from audio_preprocessor import preprocess_audio  # noqa: E402

torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
dev = "cuda"
ref = SRC.RawNetLite().to(dev)
ref.load_state_dict(torch.load(W._CKPT, map_location="cpu", weights_only=True))
ref.eval()
m = W.RawNetLite(); m.load()

d = os.path.join(HERE, "..", "J-SPAW_LA", "data")
t = pq.read_table(os.path.join(d, "test-00000-of-00001.parquet"), columns=["audio"]).slice(0, 16).column("audio").to_pylist()
auds = [sf.read(io.BytesIO(x["bytes"]), dtype="float32")[0] for x in t]
auds.append(auds[0][:20000]); auds.append(np.concatenate([auds[1]] * 4))
ours = np.array(m.score_batch(auds, [16000] * len(auds)))
p_src = []
with torch.no_grad():
    for a in auds:
        w = preprocess_audio(torch.from_numpy(a)[None, :], 16000)  # [1, 48000]
        p_src.append(float(ref(w[None].to(dev)).squeeze()))  # DataLoader adds the batch dim -> [1, 1, 48000]
p_src = np.array(p_src, dtype=np.float64)
p_ours = 1.0 / (1.0 + np.exp(ours))  # sigmoid(logit_fake), logit_fake = -score
print("lens", [len(a) for a in auds])
print("ours score (-logit_fake)", np.round(ours, 4))
print("source P(fake)          ", np.round(p_src, 5))
print("max_abs_diff P(fake)", float(np.max(np.abs(p_ours - p_src))))
ok = (p_src > 1e-6) & (p_src < 1 - 1e-6)
print("max_abs_diff score where source P not saturated (n=%d)" % ok.sum(),
      float(np.max(np.abs(ours[ok] - (-np.log(p_src[ok] / (1 - p_src[ok])))))) if ok.any() else "n/a")