import copy

from sklearn.metrics import f1_score, accuracy_score
from sklearn.model_selection import cross_validate
import numpy as np

from core.src.metrics.metrics import Metrics
from core.src.repository.DataRepository import DataRepository


class WithoutHyperParamTuningTrainer:

    @staticmethod
    def run_training():
        x_train, y_train = DataRepository.get_train_datasets()
        x_test, y_test = DataRepository.get_test_datasets()

        training_results = {}
        for model_name in DataRepository.get_initialized_models():
            training_results[model_name] =\
                WithoutHyperParamTuningTrainer.train_algorithm(
                    x_train=x_train,
                    y_train=y_train,
                    x_test=x_test,
                    y_test=y_test,
                    model_name=model_name
                )

        return training_results

    @staticmethod
    def get_classification_metrics(y_test, y_pred):
        return {
            'accuracy': Metrics.accuracy(y_test=y_test, y_pred=y_pred),
            'f1': Metrics.f1(y_test=y_test, y_pred=y_pred),
            'precision': Metrics.precision(y_test=y_test, y_pred=y_pred),
            'recall': Metrics.recall(y_test=y_test, y_pred=y_pred),
            'matthews': Metrics.matthews_corr(y_test=y_test, y_pred=y_pred)
        }

    @staticmethod
    def get_best_estimator(model, x_train, y_train):
        estimators_cv = cross_validate(estimator=model, X=x_train, y=y_train, return_estimator=True)
        max_accuracy_idx = np.argmax(estimators_cv['test_score'])
        best_estimator = estimators_cv['estimator'][max_accuracy_idx]

        return best_estimator

    @staticmethod
    def get_estimator_results(model_name, x_train, y_train, x_test, y_test):
        model = DataRepository.get_algorithm_by_name(model_name)

        best_estimator = WithoutHyperParamTuningTrainer.get_best_estimator(model=model, x_train=x_train, y_train=y_train)
        y_pred = best_estimator.predict(x_test)

        f1_list, accuracy_list = WithoutHyperParamTuningTrainer.get_elbow_curve_data(
            estimator=model,
            x_train=x_train,
            y_train=y_train,
            x_test=x_test,
            y_test=y_test
        )

        metrics_result = WithoutHyperParamTuningTrainer.get_classification_metrics(y_test=y_test, y_pred=y_pred)

        return {
            'best_estimator': best_estimator,
            'metrics': metrics_result,
            'f1_list': f1_list,
            'accuracy_list': accuracy_list
        }

    @staticmethod
    def get_elbow_curve_data(estimator, x_train, y_train, x_test, y_test):
        f1_list = []
        accuracy_list = []

        tmp_cols = []
        for col_name in x_train.columns:
            tmp_cols.append(col_name)
            tmp_x_train = x_train[tmp_cols]
            tmp_x_test = x_test[tmp_cols]

            model = copy.deepcopy(estimator)
            model.fit(tmp_x_train, y_train)

            y_pred = model.predict(tmp_x_test)

            f1_list.append(f1_score(y_test, y_pred))
            accuracy_list.append(accuracy_score(y_test, y_pred))

        return f1_list, accuracy_list

    @staticmethod
    def train_algorithm(x_train, y_train, x_test, y_test, model_name: str):
        return WithoutHyperParamTuningTrainer.get_estimator_results(
            model_name=model_name,
            x_train=x_train,
            y_train=y_train,
            x_test=x_test,
            y_test=y_test,
        )



