import re


def check_dictionaries(dataset_name, model_dict = None, train_dict = None):
    """
    Checks that dictionaries are filled in correctly, meeting models' specifications.
    
    : dataset_name (str): name assigned to the model for saving purposes
    : model_dict (dict, optional): contains hyperparameters for SNN model
    : train_dict (dict, optional): contains hyperparameters for training the SNN model
    
    """
    
    if model_dict is not None:
        print(model_dict)
        if 'is_mlp' in list(model_dict.keys()):
            assert train_dict is not None, 'Please, introduce training dictionary.'
            assert 'is_mlp' in list(train_dict.keys()), 'Please, introduce is_mlp key in the training dictionary.'
            if model_dict['is_mlp'] is True:  
                assert train_dict['is_mlp'] is True, 'Please, introduce is_mlp = True in the training dictionary, too.'
                dataset_name += '_mlp_'
            else:
                # Adding hidden states processor in the dataset name
                assert model_dict['hidden_states_processor'] is not None, 'Please introduce hidden states processor'
                if model_dict['hidden_states_processor']=='internal_processing':
                    dataset_name += '_internal_processing_'
                elif model_dict['hidden_states_processor']=='attention':
                    dataset_name += '_attention_'
                else:
                    raise NotImplementedError('hidden states processor not recognised.')
                
        if 'is_regression' in list(model_dict.keys()) and model_dict['is_regression'] is True:
            assert train_dict is not None, 'Please, introduce training dictionary.'
            assert 'is_regression' in list(train_dict.keys()), 'Please, introduce is_regression key in the training dictionary.'
            assert train_dict['is_regression'] is True, 'Please, introduce is_regression = True in the training dictionary, too.'
            dataset_name += '_regression_'
    else:
        # default hidden states processor not covered in previous cases
        dataset_name += 'attention'

    if train_dict is not None:
        # Update dataset name according to conditions
        if 'augmentation_factor_train' in list(train_dict.keys()) and train_dict['augmentation_factor_train'] > 1:
            dataset_name += '_train_augmentation_' + str(train_dict['augmentation_factor_train'])

        if 'randomising_times' in list(train_dict.keys()) and train_dict['randomising_times'] is not None:
            dataset_name += '_random_' + str(train_dict['randomising_times'])
        else:
            dataset_name += '_canonical'
   
    # Remove potential double underscores
    dataset_name = re.sub('__','_', dataset_name)
    
    return dataset_name