"""The fused attention used in training and the plain formula used for ONNX must agree."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from quran_asr import model as model_mod, text  # noqa: E402
from quran_asr.config import _wrap  # noqa: E402


def main():
    torch.manual_seed(0)
    cfg = _wrap(dict(d_model=96, num_layers=3, num_heads=4, ff_mult=2, conv_kernel=15, dropout=0.0,
                     subsampling_channels=64))
    m = model_mod.build(cfg, len(text.VOCAB)).eval()
    x = torch.randn(3, 700, 80)
    lens = torch.tensor([700, 450, 123])
    with torch.no_grad():
        fused, l1 = m(x, lens)
        plain, l2 = model_mod.use_plain_attention(m)(x, lens)
    keep = torch.arange(fused.shape[1])[None, :] < l1[:, None]
    diff = (fused - plain).abs()[keep].max().item()
    print(f"max |fused - plain| over valid frames: {diff:.2e}")
    assert torch.equal(l1, l2) and diff < 1e-4
    print("OK: training and export attention agree")


if __name__ == "__main__":
    main()