"""Source parity: the authors' LaST-Net code (model.py LaSTNet + load_dataset.read_audio/crop_or_repeat + the infer.py
scoring loop, one clip at a time) vs this wrapper (batched), on J-SPAW_LA clips, at 4 s (Arena) and 6 s (authors' default).

python source_parity.py <LaST-Net clone @ 3838f0c> [n_clips]
Deviations from running infer.py itself (both are loading-only, not math):
  * eval.load_model does torch.load(best.pt, weights_only=False) -> not allowed here; the authors' LaSTNet instead
    loads the same tensors from last_net_model.safetensors (= best.pt["model"], extract_state.py) with strict=True.
  * fairseq.checkpoint_utils.load_model_ensemble_and_task([xlsr2_300m.pt]) fails on fairseq 0.12.2 (stale task cfg
    keys) -> patched to build the same Wav2Vec2Model from the vendored xlsr2_300m model cfg; every tensor is then
    overwritten by the strict load anyway.
TF32 off on both sides (FP32).
"""
import io
import os
import sys
import tempfile

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)
src = sys.argv[1]
n = int(sys.argv[2]) if len(sys.argv) > 2 else 16
import _net  # noqa: E402
import last_net as W  # noqa: E402
import fairseq  # noqa: E402

fairseq.checkpoint_utils.load_model_ensemble_and_task = lambda paths, **k: ([_net.build_xlsr()], None, None)
sys.path.append(src)
import model as src_model  # noqa: E402
from load_dataset import SAMPLE_RATE, crop_or_repeat, read_audio  # noqa: E402

torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
ref = src_model.LaSTNet(ssl_path="<patched>", temperature=0.8, uniform_floor=0.05, layer_dropout=0.10)
r = ref.load_state_dict(W.load_state(), strict=True)
print("source strict=True: missing", r.missing_keys, "unexpected", r.unexpected_keys)
ref = ref.cuda().eval()
ours = W.LaSTNetModel()
ours.load()

t = pq.read_table(os.path.join(HERE, "..", "J-SPAW_LA", "data", "test-00000-of-00001.parquet"),
                  columns=["audio", "label"]).to_pylist()
rows = [x for x in t if x["label"] == 0][: n // 2] + [x for x in t if x["label"] == 1][: n - n // 2]
auds = [sf.read(io.BytesIO(x["audio"]["bytes"]), dtype="float32")[0] for x in rows]
auds += [auds[0][:20000], auds[1][:5000], np.concatenate([auds[2]] * 3)]  # short (repeat) + long (crop) cases
tmp = tempfile.mkdtemp()
paths = []
for k, a in enumerate(auds):
    p = os.path.join(tmp, f"{k}.wav")
    sf.write(p, a, 16000, subtype="FLOAT")  # lossless float container for the authors' path-based reader
    paths.append(p)
print("lens", [len(a) for a in auds])
for seconds in (4, 6):
    samples = seconds * SAMPLE_RATE
    theirs = []
    with torch.inference_mode():
        for p in paths:  # infer.py main() loop, verbatim
            wav = read_audio(p)
            available = min(len(wav), samples) / SAMPLE_RATE
            batch = torch.from_numpy(crop_or_repeat(wav, samples)).unsqueeze(0).to("cuda")
            duration = torch.tensor([available], device="cuda")
            logits = ref(batch, duration).float()
            theirs.append(logits.log_softmax(dim=1)[0, 1].item())
    theirs = np.array(theirs)
    if seconds == 4:
        mine = np.array(ours.score_batch(auds, [16000] * len(auds)))
    else:  # same wrapper code path with the 6 s window
        with torch.no_grad():
            xs, av = zip(*(_net.crop_np(a, samples) for a in auds))
            lg = ours.net(torch.from_numpy(np.stack(xs)).cuda(), torch.tensor(av, dtype=torch.float64).cuda()).float()
            mine = lg.log_softmax(dim=1)[:, 1].cpu().numpy()
    print(f"--- {seconds} s ({samples} samples), n={len(auds)}")
    print("source", np.round(theirs, 4).tolist())
    print("ours  ", np.round(np.asarray(mine, dtype=np.float64), 4).tolist())
    print(f"max_abs_diff_{seconds}s", float(np.max(np.abs(mine - theirs))))