"""Checks quran_asr.features against torchaudio's Kaldi-compliance fbank (the reference recipe)."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from quran_asr.features import fbank  # noqa: E402


def main():
    import torchaudio.compliance.kaldi as kaldi

    torch.manual_seed(0)
    worst = 0.0
    for n in (16000, 16000 * 3 + 123, 48000 + 77, 5000):
        wave = (torch.randn(n) * 0.1).clamp(-1, 1)
        ours = fbank(wave)
        ref = kaldi.fbank(
            wave[None, :], num_mel_bins=80, frame_length=25.0, frame_shift=10.0, dither=0.0,
            preemphasis_coefficient=0.97, remove_dc_offset=True, window_type="povey",
            round_to_power_of_two=True, low_freq=20.0, high_freq=-400.0, snip_edges=False,
            use_energy=False, use_log_fbank=True, use_power=True, sample_frequency=16000.0,
        )
        assert ours.shape == ref.shape, (ours.shape, ref.shape)
        diff = (ours - ref).abs().max().item()
        worst = max(worst, diff)
        print(f"{n} samples: {ours.shape[0]} frames, max |diff| = {diff:.2e}")
    assert worst < 1e-3, worst
    print("OK: features match the Kaldi recipe")


if __name__ == "__main__":
    main()