"""Back-end parity: authors' models/XLSR2_Nes2Net_X_ATT.py Nested_Res2Net_TDNN (incl. LocalWindowAttention)
vs ours (_net.py), both loaded from epoch_5_1.422.pth, on SSL features of real In-the-Wild clips (eval mode).
The XLS-R front-end is the fairseq model in both (the authors' SSLModel only differs in how it is constructed).
usage: python source_parity.py <path to authors' XLSR2_Nes2Net_X_ATT.py> [n_clips]"""
import importlib.util, io, os, sys
import numpy as np, torch
HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
import xlsr_nes2net_la as W

n = int(sys.argv[2]) if len(sys.argv) > 2 else 12
spec = importlib.util.spec_from_file_location("authors_net", sys.argv[1]); A = importlib.util.module_from_spec(spec); spec.loader.exec_module(A)
m = W.XLSRNes2NetLA(); m.load()
ref = A.Nested_Res2Net_TDNN(Nes_ratio=[8, 8], input_channel=1024, n_output_logits=2, dilation=2, pool_func="mean", SE_ratio=[1])
pre = "Nested_Res2Net_TDNN."
sd = {k[len(pre):]: v for k, v in W.load_state().items() if k.startswith(pre)}
r = ref.load_state_dict(sd, strict=True); assert not r.missing_keys and not r.unexpected_keys
ref = ref.to(m.device).eval()

import pyarrow.parquet as pq, soundfile as sf
t = pq.read_table(os.path.join(HERE, "..", "InTheWild", "data", "test-00000-of-00008.parquet"), columns=["audio"]).slice(0, n).to_pylist()
x = torch.from_numpy(np.stack([W.pad_fixed(sf.read(io.BytesIO(r["audio"]["bytes"]), dtype="float32")[0]) for r in t])).to(m.device)
with torch.no_grad():
    feat, ours = m.net(x)
    theirs = ref(feat)
print("ours   ", np.round(ours[:, 1].cpu().numpy(), 4)); print("authors", np.round(theirs[:, 1].cpu().numpy(), 4))
print(f"max|logits diff| = {(ours - theirs).abs().max().item():.3g}")