import os
import joblib
from sklearn.metrics import mean_absolute_error, mean_squared_error


class Model:
    def __init__(self, tree=None, feature_names=None, train_test_split=None):
        self._model = tree
        self._feature_names = feature_names
        self._x_train = self._x_test = None
        self._y_train = self._y_test = None

        if train_test_split is not None:
            self._x_train, self._x_test, \
            self._y_train, self._y_test = train_test_split

    @classmethod
    def load(cls, path):
        data = joblib.load(path)

        model = cls()
        model._model = data["model"]
        model._feature_names = data["feature_names"]

        print(f"Loaded model from {path}")

        return model

    def get_feature_names(self):
        return self._feature_names

    def get_feature_importance(self):
        importances = self._model.feature_importances_

        return sorted(
            zip(self._feature_names, importances),
            key=lambda x: x[1],
            reverse=True
        )

    def train(self):
        print("Start training...")
        self._model.fit(self._x_train, self._y_train)
        print("Training complete.")

    def evaluate(self):
        print("Evaluating model...")

        train_predictions = self._model.predict(self._x_train)
        test_predictions = self._model.predict(self._x_test)

        train_mae = mean_absolute_error(self._y_train, train_predictions)
        train_mse = mean_squared_error(self._y_train, train_predictions)

        test_mae = mean_absolute_error(self._y_test, test_predictions)
        test_mse = mean_squared_error(self._y_test, test_predictions)

        print()
        print("Training:")
        print("MAE:", train_mae)
        print("MSE:", train_mse)

        print()
        print("Test:")
        print("MAE:", test_mae)
        print("MSE:", test_mse)
        
        return [train_mae, train_mse, test_mae, test_mse]

    def release_training_data(self):
        self._x_train = None
        self._x_test = None
        self._y_train = None
        self._y_test = None


    def save_model(self, path):
        joblib.dump({
            "model": self._model,
            "feature_names": self._feature_names
        }, path)

        print(f"Saved model to {path}")