all_models = ["LinearModel", "KNN", "DecisionTree", "RandomForest", "XGBoost", "CatBoost", "LightGBM", "ModelTree",
               "MLP", "TabNet", "VIME", "TabTransformer", "NODE", "DeepGBM", "RLN", "DNFNet", "STG", "NAM", "DeepFM",
               "SAINT", "DANet"]


def str2model(model):

    if model == "LinearModel":
        from tabicl.models.baseline_models import LinearModel
        return LinearModel

    elif model == "KNN":
        from tabicl.models.baseline_models import KNN
        return KNN

    elif model == "SVM":
        from tabicl.models.baseline_models import SVM
        return SVM

    elif model == "DecisionTree":
        from tabicl.models.baseline_models import DecisionTree
        return DecisionTree

    elif model == "RandomForest":
        from tabicl.models.baseline_models import RandomForest
        return RandomForest

    elif model == "XGBoost":
        from tabicl.models.tree_models import XGBoost
        return XGBoost

    elif model == "CatBoost":
        from tabicl.models.tree_models import CatBoost
        return CatBoost

    elif model == "LightGBM":
        from tabicl.models.tree_models import LightGBM
        return LightGBM

    elif model == "MLP":
        from tabicl.models.mlp import MLP
        return MLP

    elif model == "ModelTree":
        from tabicl.models.modeltree import ModelTree
        return ModelTree

    elif model == "TabNet":
        from tabicl.models.tabnet import TabNet
        return TabNet

    elif model == "VIME":
        from tabicl.models.vime import VIME
        return VIME

    elif model == "TabTransformer":
        from tabicl.models.tabtransformer import TabTransformer
        return TabTransformer

    elif model == "NODE":
        from tabicl.models.node import NODE
        return NODE

    elif model == "DeepGBM":
        from tabicl.models.deepgbm import DeepGBM
        return DeepGBM

    elif model == "RLN":
        from tabicl.models.rln import RLN
        return RLN

    elif model == "DNFNet":
        from tabicl.models.dnf import DNFNet
        return DNFNet

    elif model == "STG":
        from tabicl.models.stochastic_gates import STG
        return STG

    elif model == "NAM":
        from tabicl.models.neural_additive_models import NAM
        return NAM

    elif model == "DeepFM":
        from tabicl.models.deepfm import DeepFM
        return DeepFM

    elif model == "SAINT":
        from tabicl.models.saint import SAINT
        return SAINT

    elif model == "DANet":
        from tabicl.models.danet import DANet
        return DANet

    else:
        raise NotImplementedError("Model \"" + model + "\" not yet implemented")
