from ml_lib.utils.config import Config
from models.conv_classifier import ConvClassifier
from models.bandwise_conv_classifier import BandwiseConvClassifier
from models.bandwise_conv_multihead_classifier import BandwiseConvMultiheadClassifier
from models.ensemble import Ensemble
from models.residual_conv_classifier import ResidualConvClassifier


class MConfig(Config):
    # TODO: There should be a way to automate this.

    @property
    def _model_map(self):
        return dict(ConvClassifier=ConvClassifier,
                    CC=ConvClassifier,
                    BandwiseConvClassifier=BandwiseConvClassifier,
                    BCC=BandwiseConvClassifier,
                    BandwiseConvMultiheadClassifier=BandwiseConvMultiheadClassifier,
                    BCMC=BandwiseConvMultiheadClassifier,
                    Ensemble=Ensemble,
                    E=Ensemble,
                    ResidualConvClassifier=ResidualConvClassifier,
                    RCC=ResidualConvClassifier
                    )