from predictors import base_predictors, ev_predictors, hmm_predictors, onehot_predictors
from predictors import unirep_predictors, esm_predictors, vae_predictors
from predictors.base_predictors import BoostingPredictor, JointPredictor


BASIC_PREDICTOR_MAP = {
    'mutation': base_predictors.MutationRadiusPredictor,
    'ev': ev_predictors.EVPredictor,
    'onehot': onehot_predictors.OnehotRidgePredictor,
    'georgiev': onehot_predictors.GeorgievRidgePredictor,
    'eunirep_reg': unirep_predictors.EUniRepRegressionPredictor,
    'gunirep_reg': unirep_predictors.GUniRepRegressionPredictor,
    'eunirep_ll': unirep_predictors.EUniRepLLPredictor,
    'gunirep_ll': unirep_predictors.GUniRepLLPredictor,
    'hmm': hmm_predictors.HMMPredictor,
    'blosum': base_predictors.BLOSUM62Predictor,
    'gesm': esm_predictors.GlobalESMPredictor,
    'gesm_reg': esm_predictors.GlobalESMRegressionPredictor,
    'vae': vae_predictors.VaePredictor,
}

CORE_PREDICTORS = [
    'eunirep_reg', 'ev+onehot', 'gesm+onehot', 'eunirep_ll+onehot',
    'vae+onehot',
]

BASELINE_PREDICTORS = [
    'georgiev', 'onehot', 'hmm+onehot', 'blosum+onehot', 'mutation+onehot',
]

ADDITIONAL_PREDICTORS = [
    'gunirep_ll+onehot',
    'gesm_reg',
]

UNSUPERVISED_PREDICTORS = [
    'ev', 'vae', 'hmm', 'blosum', 'mutation', 'eunirep_ll', 'gunirep_ll', 'gesm', 
]

def get_predictor_cls(predictor_name):
    names = predictor_name.split('+')
    return [BASIC_PREDICTOR_MAP[n] for n in names]

def get_predictor_names(key):
    if key == 'core':
        return CORE_PREDICTORS
    elif key == 'baselines':
        return BASELINE_PREDICTORS
    elif key == 'additional':
        return ADDITIONAL_PREDICTORS
    elif key == 'unsupervised':
        return UNSUPERVISED_PREDICTORS
    elif key == 'all':
        return CORE_PREDICTORS + BASELINE_PREDICTORS + ADDITIONAL_PREDICTORS + UNSUPERVISED_PREDICTORS
    else:
        return [key]