# source_parity.py SRC_DIR [CKPT] -- authors' ProSDDStage2 + main_eval.inference_forward + data_utils_eval.pad
# (from_pretrained swapped for the vendored config; checkpoint loaded strict=False exactly as main_eval.py)
# vs this wrapper, on 16 J-SPAW_LA clips. Prints max |diff| of logits[:,1].
import io, os, sys
import numpy as np, pyarrow.parquet as pq, soundfile as sf, torch
from transformers import Wav2Vec2Config, Wav2Vec2Model

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import _net
W = __import__([m for m in ("prosdd_asv19", "prosdd_asv5") if os.path.exists(os.path.join(HERE, m + ".py"))][0])
src = sys.argv[1]
ckpt = sys.argv[2] if len(sys.argv) > 2 else W._CKPT
sys.path.append(src)
Wav2Vec2Model.from_pretrained = classmethod(lambda cls, *a, **k: cls(Wav2Vec2Config.from_json_file(_net.XLSR_CONFIG)))
import model_stage2realfake, main_eval, data_utils_eval  # noqa: E402

torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
dev = "cuda"
ref = model_stage2realfake.ProSDDStage2(classifier_pool="mean", T_target=200)
state = torch.load(ckpt, map_location="cpu", weights_only=True)
r = ref.load_state_dict({k.replace("module.", ""): v for k, v in state.items()}, strict=False)
print("source strict=False: missing", r.missing_keys, "unexpected", r.unexpected_keys)
ref = ref.to(dev).eval()
m = W.__dict__[[n for n in dir(W) if n.startswith("ProSDDASV")][0]]()
m.load()

t = pq.read_table(os.path.join(HERE, "..", "J-SPAW_LA", "data", sorted(os.listdir(os.path.join(HERE, "..", "J-SPAW_LA", "data")))[-1]),
                  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))  # short + long cases
ours = np.array(m.score_batch(auds, [16000] * len(auds)))
theirs = []
with torch.no_grad():
    for a in auds:
        w = data_utils_eval.pad(torch.from_numpy(a).float())[None].to(dev)
        theirs.append(float(main_eval.inference_forward(ref, w)[:, 1]))
theirs = np.array(theirs)
print("lens", [len(a) for a in auds])
print("ours  ", np.round(ours, 4))
print("source", np.round(theirs, 4))
print("max_abs_diff", float(np.max(np.abs(ours - theirs))))