add a check for loudness as well

import numpy as np

sys.path.append(str(Path(__file__).parent.parent))

from ivas_processing_scripts.audiotools.audio import fromfile

from ivas_processing_scripts.audiotools.audioarray import getdelay

from ivas_processing_scripts.audiotools.audiofile import read

from ivas_processing_scripts.audiotools.wrappers.bs1770 import get_loudness

from ivas_processing_scripts.utils import progressbar_update, spinner

TEST_TO_OUTPUT_FORMAT = {

    "P800-1": "STEREO",

    "P800-2": "STEREO",

    "P800-3": "STEREO",

    "P800-4": "BINAURAL",

    "P800-5": "BINAURAL",

    "P800-6": "BINAURAL",

    "P800-7": "BINAURAL",

    "P800-8": "BINAURAL",

    "P800-9": "BINAURAL",

    "BS1534-1a": "STEREO",

    "BS1534-1b": "STEREO",

    "BS1534-2a": "5_1",

    "BS1534-2b": "5_1",

    "BS1534-3a": "7_1_4",

    "BS1534-3b": "7_1_4",

    "BS1534-4a": "BINAURAL",

    "BS1534-4b": "BINAURAL",

    "BS1534-5a": "BINAURAL",

    "BS1534-5b": "7_1_4",

    "BS1534-6a": "BINAURAL",

    "BS1534-6b": "BINAURAL",

    "BS1534-7a": "BINAURAL",

    "BS1534-7b": "BINAURAL",

}

def compare_audio_arrays(

    left: np.ndarray, left_fs: int, right: np.ndarray, right_fs: int

    return snr, gain_b, max_diff

def determine_out_format(file_path):

    try:

        fmt = TEST_TO_OUTPUT_FORMAT[

            re.search(r"(BS1534-\d\w)|(P800-\d)", str(file_path)).group(0)

        ]

    except KeyError:

        raise ValueError(

            "Unable to determine output format of test from given directory! Please ensure the test name is in the path (BS1534-XX or P800-X)!"

        )

    return fmt

def compare_loudness(ref, cut):

    ref_loudness, _, _ = get_loudness(ref)

    cut_loudness, _, _ = get_loudness(cut)

    if abs(ref_loudness - cut_loudness) > 0.5:

        warn(

            f"Loudness differs by more than 0.5 LFKS! File A {ref_loudness}, File B {cut_loudness}"

        )

def compare_audio_arrays_wrap(ref_file: Path, cut_file: Path):

    ref, ref_fs = read(ref_file)

    cut, cut_fs = read(cut_file)

    output_format = determine_out_format(ref_file)

    ref = fromfile(output_format, ref_file)

    cut = fromfile(output_format, cut_file)

    with catch_warnings(record=True) as warnings_list:

        snr, gain_b, max_diff = compare_audio_arrays(ref, ref_fs, cut, cut_fs)

        snr, gain_b, max_diff = compare_audio_arrays(

            ref.audio, ref.fs, cut.audio, cut.fs

        )

        if np.isnan(snr) or gain_b == 0:

            raise ValueError(f"Invalid signals! Check {ref_file} and {cut_file}!")

        compare_loudness(ref, cut)

        for w in warnings_list:

            print(f"\r{cut_file.stem} : {w.message}", flush=True)

    parser = argparse.ArgumentParser(

        description="Script perform sanity checks on a listening test output directory"

    )

    parser.add_argument("test_dir", help="Test directory with cXX directories", type=Path)

    parser.add_argument(

        "test_dir", help="Test directory with cXX directories", type=Path

    )

    args = parser.parse_args()