"""Source parity: the authors' DeepFense code (ModularDetector + fairseq_local Wav2Vec2 loader + StandardDataset
transforms load_audio/pad) vs this wrapper, on real clips.

python source_parity.py <deepfense-framework checkout @ 3f2fe11> <xlsr2_300m.pt> [n_clips]
Imports only the needed source files (package __init__s that pull mamba/qwen deps are bypassed).
"""
import io
import os
import sys
import tempfile
import types

import numpy as np
import pyarrow.parquet as pq
import soundfile as sf
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
repo, xlsr = sys.argv[1], sys.argv[2]
n = int(sys.argv[3]) if len(sys.argv) > 3 else 12
for pkg in ("deepfense", "deepfense.models", "deepfense.models.modules", "deepfense.utils",
            "deepfense.models.frontends", "deepfense.models.backends", "deepfense.models.losses",
            "deepfense.data", "deepfense.data.transforms"):
    m = types.ModuleType(pkg)
    m.__path__ = [os.path.join(repo, *pkg.split("."))]
    sys.modules[pkg] = m
import deepfense.models.frontends.wav2vec2  # noqa: E402,F401  (registers)
import deepfense.models.backends.mlp  # noqa: E402,F401
import deepfense.models.losses.cross_entropy  # noqa: E402,F401
from deepfense.models.detector import ModularDetector  # noqa: E402
from deepfense.data.transforms.transforms import load_audio, pad_combined  # noqa: E402

import deepfense_codecfake_w2v2_mlp as W  # noqa: E402

torch.backends.cudnn.allow_tf32 = False
cfg = {"frontend": {"type": "wav2vec2", "args": {"ckpt_path": xlsr, "freeze": False, "source": "fairseq"}},
       "backend": {"type": "MLP", "args": {"input_dim": 1024, "projection": [128], "pooling_type": "mean"}},
       "loss": [{"type": "CrossEntropy", "embedding_dim": 128, "n_classes": 2, "weight": 1.0}]}
src = ModularDetector(cfg)
r = src.load_state_dict(W.load_state(), strict=False)
print("source strict=False: missing", r.missing_keys, "unexpected", r.unexpected_keys)
src = src.cuda().eval()

ours = W.DeepFenseCodecFakeW2V2MLP()
ours.load()

clips = []
for ds in ("J-SPAW_LA", "InTheWild"):
    f = sorted(x for x in os.listdir(os.path.join(HERE, "..", ds, "data")) if x.startswith("test-"))[0]
    t = pq.read_table(os.path.join(HERE, "..", ds, "data", f), columns=["audio"]).slice(0, n // 2).to_pylist()
    clips += [r["audio"]["bytes"] for r in t]
tmp = tempfile.mkdtemp()
d_logit, d_score, s_src, s_ours = [], [], [], []
with torch.no_grad():
    for k, b in enumerate(clips):
        a, sr = sf.read(io.BytesIO(b), dtype="float32")
        p = os.path.join(tmp, f"{k}.wav")
        sf.write(p, a, sr, subtype="FLOAT")  # lossless float container for the source path-based loader
        x = pad_combined(load_audio(p), max_len=64000, random_pad=False, pad_type="repeat")
        out = src(torch.tensor(x, dtype=torch.float32)[None].cuda())
        ref = out["logits"][0].cpu().numpy()
        mine = ours.net(torch.from_numpy(W.pad_np(a))[None].cuda())[0].cpu().numpy()
        d_logit.append(np.abs(ref - mine).max())
        s_src.append(float(out["scores"][0]))  # source get_score = logits[:, bonafide_label=1]
        s_ours.append(ours.score_batch([a], [16000])[0])
        d_score.append(abs((ref[1] - ref[0]) - s_ours[-1]))
print(f"n={len(clips)} max|dlogits|={max(d_logit):.3g} max|d(logit1-logit0)|={max(d_score):.3g}")
print("source get_score (logit1):", np.round(s_src, 3).tolist())
print("ours (logit1-logit0):     ", np.round(s_ours, 3).tolist())