import importlib
import re

MODElNAMES = {
    'lenet5'  : 'LeNet5',
    'alexnetcifar' : 'AlexNetCifar',
    'cnnham' : 'CNNHAM',
    }
dataset_num_classes = {'fashionmnist': 10, 'cifar10': 10, 'cinic10': 10, 'ham': 7}

def check(model, dataset):
    if 'resnet' in model:
        model_pattern1 = r'logistic|mlp|lenet5|cnn|vgg|wvgg|resnetgnii|resnetii|wrngn|wrnnsgn|wrn|resnetgn\d+$|resnetgn\d+x\d+$'
        model_result1 = re.findall(model_pattern1, model)
        assert model_result1 != []
        model = re.sub(r'\d+', '', model)
    model_pattern2 = r'lenet5|cnn|alexnet|resnetii'    
    model_result = re.search(model_pattern2, model)
    model_result = model_result.group(0) if model_result != None else model_result
    
    if dataset == 'cinic10':
        dataset_result = 'cifar'
    else:
        dataset_pattern = r'mnist|cifar|ham'
        dataset_result = re.search(dataset_pattern, dataset)
        dataset_result = dataset_result.group(0) if dataset_result != None else dataset_result
    return model_result, dataset_result

def build_model(model='lenet5', dataset='mnist'):
    check_result = check(model, dataset)
    assert check_result[0] != None and check_result[1] != None

    # module
    module = importlib.import_module('.{}_{}'.format(check_result[0], check_result[1]), 'sim.models')
    # class or function 
    if model in ['lenet5', 'alexnetcifar', 'cnnham']:
        model_class = getattr(module, MODElNAMES[model])
    else:
        model_class = getattr(module, model)
    # model
    num_classes = dataset_num_classes[dataset]
    if model =='lenet5' and dataset in ['mnist', 'fashionmnist']:
        model = model_class()
    elif model == 'alexnetcifar' and dataset in ['cifar10', 'cinic10']:
        model = model_class(num_classes=num_classes)
    elif model == 'cnnham' and dataset == 'ham':
        model = model_class()
    elif check_result[0] in ['vgg', 'wvgg', 'resnetii', 'resnetgnii', 'resnetgn', 'wrn', 'wrngn', 'wrnnsgn'] and dataset in ['cifar10', 'cifar100', 'cinic10']:
        model = model_class(num_classes=num_classes)
    return model
