# Usage: CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 python source_parity.py
# Authors' model.py + inference.py vs our wrapper on 20 In-the-Wild clips (10 bona fide + 10 spoof).
# The authors' DeepfakeDetector calls WavLMModel.from_pretrained("microsoft/wavlm-large"); with no Hub access we
# point that call at a local dir holding the vendored config + the checkpoint's own wavlm.* tensors (the real
# from_pretrained code path), then load the full checkpoint with the authors' strict=False, as inference.py does.
import os, sys, io, importlib.util, tempfile, numpy as np, torch, soundfile as sf, pyarrow.parquet as pq
os.environ["HF_HUB_OFFLINE"] = "1"
print("device", torch.cuda.get_device_name())
MD = os.path.dirname(os.path.abspath(__file__))
MOD, CLS, SRC, CKPT = {
    "Forensics-WavLM-Base": ("forensics_wavlm_base", "ForensicsWavLMBase",
                             "eliya__forensics_0.3B_base_deepfake_classifier", "checkpoint_epoch_5.safetensors"),
    "Forensics-WavLM-OCSoftmax": ("forensics_wavlm_ocsoftmax", "ForensicsWavLMOCSoftmax",
                                  "eliya__forensics_0.3B_wavlm_oc_softmax_deepfake_classifier",
                                  "2nd_oc_softmax_ce_from16_last.safetensors"),
}[os.path.basename(MD)]
SRC = os.path.join(MD, "..", "_ckpt_dl", SRC)
from safetensors.torch import load_file, save_file
import transformers
def imp(name, path):
    spec = importlib.util.spec_from_file_location(name, path); m = importlib.util.module_from_spec(spec)
    sys.modules[name] = m; spec.loader.exec_module(m); return m
sd = load_file(f"{SRC}/{CKPT}")
tmp = tempfile.mkdtemp()
import shutil; shutil.copy(f"{MD}/wavlm_large_config.json", f"{tmp}/config.json")
save_file({k[len("wavlm."):]: v.contiguous() for k, v in sd.items() if k.startswith("wavlm.")}, f"{tmp}/model.safetensors")
_orig = transformers.WavLMModel.from_pretrained.__func__
transformers.WavLMModel.from_pretrained = classmethod(lambda cls, name, *a, **k: _orig(cls, tmp, *a, **k))
srcmodel = imp("model", f"{SRC}/model.py")
srcinf = imp("inference", f"{SRC}/inference.py")
ref = srcmodel.DeepfakeDetector().cuda().eval()
r = ref.load_state_dict(sd, strict=False)
print("source strict=False:", "missing", r.missing_keys, "unexpected", r.unexpected_keys)
sys.path.insert(0, MD)
F = importlib.import_module(MOD)
m = getattr(F, CLS)(); m.load()
print("params", sum(p.numel() for p in m.net.parameters()))
t = pq.read_table(f"{MD}/../InTheWild/data/test-00000-of-00008.parquet", columns=["audio", "label"]).slice(0, 400).to_pylist()
by = {0: [], 1: []}
for row in t:
    if len(by[row["label"]]) < 10:
        by[row["label"]].append(row["audio"]["bytes"])
a, b, lab = [], [], []
for y in (0, 1):
    for i, byts in enumerate(by[y]):
        f = f"{tmp}/clip_{y}_{i}.wav"; open(f, "wb").write(byts)
        with torch.no_grad():
            a.append(ref(srcinf.load_audio(f).unsqueeze(0).cuda()).float().item())
        w, sr = sf.read(f, dtype="float32"); assert sr == 16000, sr
        if w.ndim > 1: w = w.mean(1)
        b.append(m.score_batch([w], [16000])[0]); lab.append(y)
a, b = np.array(a), np.array(b)
print("src ", np.round(a, 4)); print("ours", np.round(b, 4)); print("labels", lab)
print("n", len(a), "max abs diff", np.abs(a - b).max())
shutil.rmtree(tmp)