"""This is the main script run by the nox session that generates data found in
``./tests/data-for-tests/generated``.

Some data is generated by other scripts in this directory,
that need to be run first:
* tests/scripts/generate_ava_segment_test_data/generate_ava_segment_text_files_from_jourjine_et_al_2023.py
  that can be run with the nox session test-data-generate-ava-segments
"""
import json
import pathlib
import shutil

import vocalpy as voc


# Annoyingly, we have to repeat ourselves here and declare all these paths
# that are *also* declared as constants in tests/fixtures.
# That's because tests/fixtures expect the files that this script generates to exist.
# So if we try to import test/fixtures, and the files get loaded, e.g. to parse json for a fixture,
# but they don't exist yet, we get FileNotFoundErrors.
HERE = pathlib.Path(__file__).parent
PROJ_ROOT = (HERE / '..' / '..').resolve()
# next line, use `resolve` + `relative_to` so all paths we compute with this are relative to project root,
# which matters e.g. when we have fixtures that use paths to files, we don't want them to be
# absolute for a specific system (such as, my computer)
TEST_DATA_ROOT = HERE.joinpath('..', 'data-for-tests').resolve().relative_to(PROJ_ROOT)
SOURCE_TEST_DATA_ROOT = TEST_DATA_ROOT / 'source'
GENERATED_TEST_DATA_ROOT = TEST_DATA_ROOT / 'generated'

# ---- sound
AUDIO_DIR_CBIN = SOURCE_TEST_DATA_ROOT / 'audio_cbin_annot_notmat' / 'gy6or6' / '032312'
AUDIO_LIST_CBIN = sorted(AUDIO_DIR_CBIN.glob('*.cbin'))
BIRDSONGREC_WAV_DIR =  SOURCE_TEST_DATA_ROOT / 'audio_wav_annot_birdsongrec' / 'Bird0' / 'Wave'
BIRDSONGREC_WAV_LIST = sorted(BIRDSONGREC_WAV_DIR.glob('*.wav'))

# ---- "oracle" segments for 'vocalpy.segment.meansquared'
SEGMENT_MAT_DIR = SOURCE_TEST_DATA_ROOT / 'segment' / 'evsonganaly'
GENERATED_SEGMENT_ROOT = GENERATED_TEST_DATA_ROOT / 'segment'
# ---- "oracle" segments for 'vocalpy.segment.ava'
AVA_SEGMENT_TXT_DIR = GENERATED_SEGMENT_ROOT / 'ava-segment-txt'
AVA_SEGMENT_TXT_LIST = sorted(AVA_SEGMENT_TXT_DIR.glob('*.txt'))

# ---- these are the dirs and files that this script actually makes
PATHS_DIR = GENERATED_TEST_DATA_ROOT / 'paths'
DATA_ROOTS_WITH_SUBDIRS = PATHS_DIR / 'root_with_subdirs_to_test_paths'
SPECTROGRAM_DIR = GENERATED_TEST_DATA_ROOT / 'spectrogram'
SPECT_DIR_NPZ = SPECTROGRAM_DIR / 'spect_npz'

EVSONGANALY_SEGMENT_DIR = GENERATED_SEGMENT_ROOT / 'evsonganaly'
EVSONGANALY_SEGMENT_JSON_PATH = EVSONGANALY_SEGMENT_DIR / 'evsonganaly-cbin-notmat-segment-mat-paths.json'


def mkdirs():
    print(
        f"Making data roots with subdirs for testing ``vocalpy.paths``: {DATA_ROOTS_WITH_SUBDIRS}"
    )
    DATA_ROOTS_WITH_SUBDIRS.mkdir()
    print(
        f"Making spectrogram directory for testing ``vocalpy.Spectrogram``: {SPECT_DIR_NPZ}"
    )
    SPECT_DIR_NPZ.mkdir()
    print(
        f"Making directory for testing ``vocalpy.segment.meansquared``: {EVSONGANALY_SEGMENT_DIR}"
    )
    EVSONGANALY_SEGMENT_DIR.mkdir()


def generate_npz_spect_files():
    """This generates .npz files with spectrograms in them,
    so we can test that we can load these files
    """
    print(
        f"Making spectrogram files in spectrogram directory, for testing ``vocalpy.Spectrogram``: {SPECT_DIR_NPZ}"
    )
    for wav_path in BIRDSONGREC_WAV_LIST:
        sound = voc.Sound.read(wav_path)
        spect = voc.spectrogram(sound)
        dst_path = SPECT_DIR_NPZ / f'{wav_path.name}.npz'
        print(
            f"Writing spectrogram to path: {dst_path}"
        )
        spect.write(path=dst_path)


def generate_data_roots_with_subdirs():
    print(
        "Copying sound files into sub-dirs for testing ``vocalpy.paths``"
    )
    for ext, path_list in (
        ('wav', BIRDSONGREC_WAV_LIST),
        ('cbin', AUDIO_LIST_CBIN),
    ):
        print(
            f"Copying files with extension: {ext}"
        )
        # ---- set-up
        data_root = DATA_ROOTS_WITH_SUBDIRS / f'{ext}'
        print(
            f"Making dir for sound files with extension {ext}: {data_root}"
        )
        data_root.mkdir()

        dst1 = data_root / f'{ext}-first-half'
        print(
            f"Making sub-dir for sound files with extension {ext}: {dst1}"
        )
        dst1.mkdir()
        dst2 = data_root / f'{ext}-second-half'
        print(
            f"Making sub-dir for sound files with extension {ext}: {dst2}"
        )
        dst2.mkdir()

        half_ind = len(path_list) // 2
        paths1 = path_list[:half_ind]
        paths2 = path_list[half_ind:]

        for dst, paths in zip(
            (dst1, dst2),
            (paths1, paths2)
        ):
            for path in paths:
                path_dst = dst / path.name
                print(
                    f"Copying sound file to path: {path_dst}"
                )
                shutil.copy(path, path_dst)


def check_ava_segment_data_exists():
    if len(AVA_SEGMENT_TXT_LIST) < 1:
        raise ValueError(
            f"No segment .txt files from ava segmentation found in directory: {AVA_SEGMENT_TXT_DIR}."
            f"Run the nox session 'test-data-generate-ava-segment' before running the session that runs this script."
        )


def generate_evsonganaly_segment_json():
    """Generate .json with paths needed to test that
    vocalpy.segment.meansquared replicates the evsonganaly segment function
    """
    cbin_paths = AUDIO_LIST_CBIN  # same list we use above, cbin files from AUDIO_DIR_CBIN
    notmat_paths = sorted(
        # .not.mat annotation files are in this same directory
        AUDIO_DIR_CBIN.glob('*not.mat')
    )
    segment_mat_paths = sorted(SEGMENT_MAT_DIR.glob('*.mat'))

    for_json = []
    for segment_mat_path in segment_mat_paths:
        cbin_ind = segment_mat_path.name.find('cbin')
        cbin_name = segment_mat_path.name[:cbin_ind + 4]
        cbin_path = [
            cbin_path
            for cbin_path in cbin_paths
            if cbin_path.name.startswith(cbin_name)
        ]
        assert len(cbin_path) == 1
        cbin_path = cbin_path[0]
        notmat_path = [
            notmat_path
            for notmat_path in notmat_paths
            if notmat_path.name.startswith(cbin_name)
        ]
        assert len(notmat_path) == 1
        notmat_path = notmat_path[0]
        for_json.append(
            {
                'cbin_path': str(cbin_path),
                'notmat_path': str(notmat_path),
                'segment_mat_path': str(segment_mat_path),
            }
        )
        with (
            EVSONGANALY_SEGMENT_JSON_PATH
        ).open('w') as fp:
            json.dump(
                for_json, fp
            )


def main():
    """Makes directories in `./tests/data-for-tests/generated
    (after those have been removed by running ``nox -s test-data-clean-generated``).
    Then runs helper functions that generate test data from source.
    """
    print(
        f"Making sub-directories in generated test data dir: {GENERATED_TEST_DATA_ROOT}"
    )
    mkdirs()
    generate_npz_spect_files()
    generate_data_roots_with_subdirs()
    check_ava_segment_data_exists()
    generate_evsonganaly_segment_json()


main()
