# modelrunner.py

from nkululeko.util import Util
from nkululeko import glob_conf


class Modelrunner:
    """
    Class to model one run
    """

    def __init__(self, df_train, df_test, feats_train, feats_test, run):
        """Constructor setting up the dataframes
        Args:
            df_train: train dataframe
            df_test: test dataframe
            feats_train: train features
            feats_train: test features
        """
        self.df_train, self.df_test, self.feats_train, self.feats_test = (
            df_train,
            df_test,
            feats_train,
            feats_test,
        )
        self.util = Util("modelrunner")
        self.run = run
        self.target = glob_conf.config["DATA"]["target"]
        # intialize a new model
        model_type = glob_conf.config["MODEL"]["type"]
        self._select_model(model_type)

    def do_epochs(self):
        # initialze results
        reports = []
        plot_epochs = self.util.config_val("PLOT", "epochs", False)
        only_test = self.util.config_val("MODEL", "only_test", False)
        # for all epochs
        for epoch in range(int(self.util.config_val("EXP", "epochs", 1))):
            if only_test:
                self.model.load(self.run, epoch)
                self.util.debug(f"reusing model: {self.model.store_path}")
                self.model.reset_test(self.df_test, self.feats_test)
            else:
                self.model.set_id(self.run, epoch)
                self.model.train()
            report = self.model.predict()
            report.set_id(self.run, epoch)
            plot_name = self.util.get_plot_name() + f"_{self.run}_{epoch:03d}_cnf"
            reports.append(report)
            self.util.debug(
                f"run: {self.run} epoch: {epoch}: result: "
                f"{reports[-1].get_result().get_test_result()}"
            )
            if plot_epochs:
                self.util.debug(f"plotting conf matrix to {plot_name}")
                report.plot_confmatrix(plot_name, epoch)
            store_models = self.util.config_val("MODEL", "save", False)
            plot_best_model = self.util.config_val("PLOT", "best_model", False)
            if (store_models or plot_best_model) and (
                not only_test
            ):  # in any case the model needs to be stored to disk.
                self.model.store()
        if not plot_epochs:
            # Do at least one confusion matrix plot
            self.util.debug(f"plotting confusion matrix to {plot_name}")
            reports[-1].plot_confmatrix(plot_name, epoch)
        return reports

    def _select_model(self, model_type):
        if model_type == "svm":
            from nkululeko.models.model_svm import SVM_model

            self.model = SVM_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "svr":
            from nkululeko.models.model_svr import SVR_model

            self.model = SVR_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "xgb":
            from nkululeko.models.model_xgb import XGB_model

            self.model = XGB_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "xgr":
            from nkululeko.models.model_xgr import XGR_model

            self.model = XGR_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "bayes":
            from nkululeko.models.model_bayes import Bayes_model

            self.model = Bayes_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "gmm":
            from nkululeko.models.model_gmm import GMM_model

            self.model = GMM_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "knn":
            from nkululeko.models.model_knn import KNN_model

            self.model = KNN_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "knn_reg":
            from nkululeko.models.model_knn_reg import KNN_reg_model

            self.model = KNN_reg_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "tree":
            from nkululeko.models.model_tree import Tree_model

            self.model = Tree_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "tree_reg":
            from nkululeko.models.model_tree_reg import Tree_reg_model

            self.model = Tree_reg_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "cnn":
            from nkululeko.models.model_cnn import CNN_model
            from nkululeko.models.model_cnn import CNN_model

            self.model = CNN_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "mlp":
            from nkululeko.models.model_mlp import MLP_model

            self.model = MLP_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        elif model_type == "mlp_reg":
            from nkululeko.models.model_mlp_regression import MLP_Reg_model

            self.model = MLP_Reg_model(
                self.df_train, self.df_test, self.feats_train, self.feats_test
            )
        else:
            self.util.error(f"unknown model type: '{model_type}'")
        if self.util.exp_is_classification() and not self.model.is_classifier:
            self.util.error(
                "Experiment type set to classification but model type is not a classifier"
            )
        return self.model
