from argparse import Namespace

import warnings

from datasets.ccs_librosa_datamodule import CCSLibrosaDatamodule
from datasets.primates_librosa_datamodule import PrimatesLibrosaDatamodule
from datasets.mask_librosa_datamodule import MaskLibrosaDatamodule
from ml_lib.utils.config import parse_comandline_args_add_defaults

warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)


def rebuild_dataset(h_params, data_class):

    # START
    # =============================================================================
    # Let Datamodule pull what it wants
    datamodule = data_class.from_argparse_args(h_params)
    assert datamodule.purge()
    datasets = datamodule.manual_setup()
    print(f'Dataset length is: {len(datasets)}')


if __name__ == '__main__':
    for dataset in [CCSLibrosaDatamodule]:
        # Parse comandline args, read config and get model
        cmd_args, _, _, _ = parse_comandline_args_add_defaults('_parameters.ini')

        # To NameSpace
        hparams = Namespace(**cmd_args)

        # Start
        # -----------------
        rebuild_dataset(hparams, dataset)
        print('done')
        pass