import sys
import os
file_path = os.getcwd()
sys.path.append(file_path)

import root_gnn_base.utils as utils
import argparse
from root_gnn_base.batched_dataset import PreBatchedDataset
from root_gnn_base.batched_dataset import LazyPreBatchedDataset

def main():
    parser = argparse.ArgumentParser()
    add_arg = parser.add_argument
    add_arg('--config', type=str, required=True)
    add_arg('--dataset', type=str, required=True)
    add_arg('--chunk', type=int, default=0)
    add_arg('--shuffle_mode', action='store_true', help='Shuffle the dataset before training.')
    add_arg('--drop_last', action='store_false', help='Set drop_last to False if the flag is provided. Defaults to True.')
    args = parser.parse_args()

    config = utils.load_config(args.config)
    dset_config = config['Datasets'][args.dataset]
    batch_size = config['Training']['batch_size']
    if not args.shuffle_mode:
        dset = utils.buildFromConfig(dset_config, {'process_chunks': [args.chunk,]})
    else:
        dset = utils.buildFromConfig(dset_config)
        if 'batch_size' in dset_config:
            batch_size = dset_config['batch_size']

        shuffle_chunks = dset_config.get('shuffle_chunks', 10)
        padding_mode = dset_config.get('padding_mode', 'STEPS')
        fold_conf = dset_config["folding"]
        print(f"shuffle_chunks = {shuffle_chunks}, args.chunk = {args.chunk}, padding_mode = {padding_mode}")
        if dset_config["class"] == "LazyMultiLabelDataset":
            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'] )
            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'])

        else:
            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'])
            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'] )

if __name__ == "__main__":
    main()