"""This is the script run by nox session "test-data-generate".

It makes all the 'generated' test data, i.e. files created by vak,
It's called 'generated' test data to distinguish it from the
'source' test data, i.e., files *not* created by vak:
the input data such as audio and annotation files that are used
when vak *does* create files (csv files, logs,
neural network checkpoints, etc.).
This script generates those files for use by unit and integration tests.
Specifically, it generates:
* temporary config.toml files used when generating results
* `prep`d (prepared) datasets, and results created with those datasets,
  both of which were generated using the temporary config.toml files

All the setup configs send output to one of two places:
- for any prep command, the output goes to some child directory of ./tests/data_for_tests/generated/prep
- for any command run with a `prep`d dataset, the output goes to some child directory
  of ./tests/data_for_tests/generated/results

Examples:
    when we run `vak prep tests/data_for_tests/generated/configs/test_train_audio_wav_annot_birdsongrec.toml`
    the `prep`d dataset will be in a new directory created in
    `./tests/data_for_tests/generated/prep/train/audio_wav_annot_birdsongrec`

    when we run `vak train tests/data_for_tests/genereated/configs/test_train_audio_wav_annot_birdsongrec.toml`
    it will use the `prep`d dataset csv that's now in
    `./tests/data_for_tests/generated/prep/train/audio_wav_annot_birdsongrec`, that the temporary config.toml points to,
    and the results will go to a new directory created in
    `./tests/data_for_tests/generated/results/train/audio_wav_annot_birdsongrec`

To set up this directory structure, we iterate through the constants defined below.

The constants are:
- TOP_LEVEL_DIRS
    name of sub-directories in .tests/data_for_tests/generated that correspond to
    either tempory config files, 'prep'd datasets, or `results` generated from those `prep`d datasets
- COMMAND_DIRS
    names of sub-sub-directories in ./tests/data_for_tests/generated/prep &
    ./tests/data_for_tests/generated/results that correspond to cli commands;
    e.g., dataset from running `vak prep $A_TRAIN_CONFIG.TOML` will be in
    ./tests/data_for_tests/generated/prep/train/audio_{}_annot_{}
- DATA_DIRS
    names of "subsub"directories in ./tests/data_for_tests/$COMMAND that correspond to toy test data sets;
    those sub-directories contain test data generated by $COMMAND using the specified toy test data set

In other words, the parent directory for all the generated directories we need to remove
will have a path of the form: `./tests/data_for_tests/$TOP_LEVEL_DIRS/COMMAND_DIRS/DATA_DIRS`.
For example:
The prep`d dataset from running `vak prep $AUDIO_CBIN_ANNOT_NOTMAT_CONFIG.TOML` will be in
`./tests/data_for_tests/prep/train/audio_cbin_annot_notmat`
and corresponding results will be in
`./tests/data_for_tests/results/train/audio_cbin_annot_notmat`
The directories will have names with timestamps like `prep_20201015_1115`.
Those are the generated directories we want to remove.
"""
import logging
import sys
import warnings

# Do this here to suppress warnings before we import vak
from numba.core.errors import NumbaDeprecationWarning
warnings.simplefilter('ignore', category=NumbaDeprecationWarning)
import vak

import vaktestdata


logger = logging.getLogger('vaktestdata')  # 'base' logger
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)
logger.setLevel('INFO')


def generate_test_data(
        step: str = 'all',
        commands=vaktestdata.constants.COMMANDS,
        single_train_result:bool = True,
):
    """Main function that generates all the test data.

    Parameters
    ----------
    step : str
        The step of generating test data to carry out.
        One of ('prep', 'results', or 'all').
        This option is provided through the command-line
        because generating all the data sets is time-consuming
        and may not be needed every time test data is re-generated.
        Running the 'prep' step will run just `vak prep`
        for every configuration file.
        Running the 'results' step assumes that the 'prep' step
        has already been run, and it will run the appropriate command
        for each configuration file to generate results with the prep'd
        dataset, e.g. `vak train`.
        Running 'both' will first run 'prep' and then 'results'--this is
        the default.
    commands : list
        List of commands to run for 'results' step.
    """
    # need to run `prep` before we run other commands
    if step in ('prep', 'all'):
        # first we generate outputs of processing steps
        # leading up to a dataset that speed up tests
        vaktestdata.dirs.make_spect_output_dir_in_generated()  # for any prepared spectrograms
        vaktestdata.dirs.make_source_files_csv_dir_in_generated()  # for csvs of source files
        vaktestdata.dirs.make_source_files_with_splits_csv_dir_in_generated()  # same csvs, with splits added

        # -- now actually run prep for all the configs
        config_paths = vaktestdata.configs.copy_config_files()

        vaktestdata.source_files.set_up_source_files_and_csv_files_for_frame_classification_models()

        vaktestdata.dirs.make_subdirs_in_generated(config_paths)
        # run prep for some models
        vaktestdata.prep.run_prep()
        # re-use some of the prepped datasets for other models
        # this makes time to prep all datasets shorter
        vaktestdata.configs.add_dataset_path_from_prepped_configs()

    if step in ('results', 'all'):
        # Note we need to run `train` first to get results needed for `eval', 'predict' and continuing 'train'
        for command in commands:
            if command == "prep":
                continue  # we don't run prep in this code block
            logger.info(f"running configs for command: {command}")
            command_config_metadata = [
                config_metadata
                for config_metadata in vaktestdata.constants.CONFIG_METADATA
                if config_metadata.config_type == command
            ]

            if command in ("predict", "eval", "train_continue"):
                # Fix values for required options in predict / eval / train_continue configs
                # using results from running the corresponding train configs.
                # this only works if we ran the train configs already,
                # which we should have because of ordering of COMMANDS constant above
                vaktestdata.configs.fix_options_in_configs(command_config_metadata, command, single_train_result)

            if command == "train_continue":
                # so we don't get 'command not recognized' error in next code block
                command = "train"

            for config_metadata in command_config_metadata:
                config_path = vaktestdata.constants.GENERATED_TEST_CONFIGS_ROOT / config_metadata.filename
                logger.info(
                    f"n\Running 'vak {command}', using config: {config_path.name}"
                )
                vak.cli.cli.cli(command, config_path)


if __name__ == "__main__":
    parser = vaktestdata.parser.get_parser()
    args = parser.parse_args()
    generate_test_data(
        step=args.step,
        commands=args.commands,
        single_train_result=args.single_train_result,
    )
