from argparse import Namespace

import warnings

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.prepare_data()
    datasets = datamodule.setup()
    print(f'Dataset length is: {len(datasets)}')


if __name__ == '__main__':
    # Parse comandline args, read config and get model
    cmd_args, found_data_class, _ = parse_comandline_args_add_defaults('_parameters.ini')

    # To NameSpace
    hparams = Namespace(**cmd_args)

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