"""WST-Graph TorchScript loader + preprocessing (HF kwokho1/wst-graph @ 080000c1685451c7759f43d2ae7d0f08477cfdf8, MIT).

There is no Python model code to vendor: the authors release only traced TorchScript programs (waveform -> logits,
Kymatio WST front-end + GraphAASIST back-end, all parameters inside). The program was reviewed before first use
(read from its zip `code/` dir and from `model.code`/`model.graph` of every submodule): only aten:: / prim:: nodes,
classes limited to the model's own modules (__torch__.models.*, kymatio ScatteringTorch1D, torch.nn modules), pickled
globals only torch storages / _rebuild_tensor_v2 / OrderedDict. `load_ts` re-checks the node kinds on every load.

Preprocessing = the model card's quick start: first 64000 samples, right zero-pad to 64000, input_length = retained
sample count (int64). Class 0 = bona fide, 1 = spoof. The `_cuda` file was traced on cuda:0 (device constants are
baked in), so it runs on CUDA only.
"""
import os

import numpy as np
import torch

TARGET_SAMPLES = 64000  # model card: keep the first 64,000 samples (4 s)


def prep_np(wav, max_len=TARGET_SAMPLES):
    """-> (float32 [max_len] right zero-padded, input_length). Empty clip -> zeros, length 1 (the card raises)."""
    wav = np.asarray(wav, dtype=np.float32).reshape(-1)
    n = min(wav.shape[0], max_len)
    out = np.zeros(max_len, dtype=np.float32)
    out[:n] = wav[:n]
    return out, max(n, 1)


def check_graph(module):
    """Assert the TorchScript program only uses aten::/prim:: ops (no custom/IO ops)."""
    bad = set()

    def walk(block):
        for node in block.nodes():
            k = node.kind()
            if not k.startswith(("aten::", "prim::")):
                bad.add(k)
            for b in node.blocks():
                walk(b)

    for _, sub in module.named_modules():
        try:
            g = sub.graph
        except (AttributeError, RuntimeError):
            continue
        walk(g)
    assert not bad, f"unexpected TorchScript ops: {sorted(bad)}"


def load_ts(path, device="cuda"):
    m = torch.jit.load(path, map_location=device).eval()
    check_graph(m)
    return m