import random
import string
from typing import Optional
import os
import pandas as pd
import numpy as np
import autokeras as ak
from sklearn import preprocessing
from mindsdb.integrations.libs.base import BaseMLEngine
from tensorflow.keras.models import load_model


# Makes this run on systems where this arg isn't specified, like Windows Subsystem for Linux
# Doesn't break things on Mac
os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/lib/cuda"

trainer_dict = {"regression": ak.StructuredDataRegressor, "classification": ak.StructuredDataClassifier}

DEFAULT_TRIALS = 100


def train_model(df, target, max_trials=DEFAULT_TRIALS):
    """Helper function to trains an AutoKeras model with an input df.

    Automatically decides on classification vs. regression depending on
    the type of the target.

    Will auto-encode categorical variables as dummies.

    Returns both trained model and the names of categoric dummy columns, which
    are passed later to the prediction method.
    """
    # Choose regressor of classifier based on target data type
    if np.issubdtype(df[target].dtype, np.number):
        mode = "regression"
        y_train = df[target]
    else:
        mode = "classification"
        lb = preprocessing.LabelBinarizer()
        y_train = lb.fit_transform(df[target])

    training_df = df.drop(target, axis=1)
    trainer = trainer_dict[mode](overwrite=True, max_trials=max_trials)

    # Save the column names of all numeric columns before transforming any categorical columns into dummies
    numeric_column_names = training_df.select_dtypes(include=[np.number]).columns.values.tolist()
    training_df = pd.get_dummies(training_df)
    categorical_dummy_column_names = [
        col for col in training_df.columns.values.tolist() if col not in numeric_column_names
    ]
    trainer.fit(training_df, y_train, verbose=2)
    return trainer.export_model(), categorical_dummy_column_names


def get_preds_from_model(df, model, target, column_count, categorical_dummy_column_names):
    """Gets predictions from the stored AutoKeras model."""
    df_to_predict = df.copy()
    for col in ["__mindsdb_row_id", target]:
        if col in df_to_predict.columns.values.tolist():
            df_to_predict = df_to_predict.drop(col, axis=1)

    if len(df_to_predict.columns) != column_count:
        raise Exception("All feature columns must be defined in the WHERE clause when making predictions")
    # Get dummies for any categorical columns and then populate the missing ones with zeros
    df_with_dummies = pd.get_dummies(df_to_predict)
    for col in categorical_dummy_column_names:  # exception handler for empty columns
        if col not in df_with_dummies.columns.values.tolist():
            df_with_dummies[col] = 0

    return model.predict(df_with_dummies, verbose=2)


def format_categorical_preds(predictions, original_y, df_to_predict, target_col):
    """Transforms class predictions back to their original class.

    Categoric predictions come out the AutoKeras model in a binary
    format e.g. (0, 1, 0). This function maps them back to their
    original class e.g. 'Blue', and adds a DF column for the
    model confidence score.
    """
    # Turn prediction back into categorical value
    lb = preprocessing.LabelBinarizer()
    lb.fit(original_y)
    preds = lb.inverse_transform(predictions)

    # Add the confidence score next to the prediction
    df_to_predict[target_col] = pd.Series(preds).astype(original_y.dtype)
    df_to_predict["confidence"] = [max(row) for _, row in enumerate(predictions)]
    return df_to_predict


class AutokerasHandler(BaseMLEngine):
    """
    Integration with the AutoKeras ML library.
    """  # noqa

    name = "autokeras"

    def create(self, target: str, df: Optional[pd.DataFrame] = None, args: Optional[dict] = None) -> None:
        """Create and tune AutoKeras model using the input df.

        Saves the AutoKeras best model params to model storage.
        """
        args = args["using"]  # ignore the rest of the problem definition
        args["target"] = target
        max_trials = int(args["train_time"] * DEFAULT_TRIALS) if "train_time" in args else DEFAULT_TRIALS
        # Save the training df in order to filter the training data based on the predict df
        args["training_df"] = df.to_json()
        args["training_data_column_count"] = len(df.columns) - 1  # subtract 1 for target

        random_string = "".join(random.choices(string.ascii_uppercase + string.digits, k=24))
        args["folder_path"] = os.path.join("autokeras", random_string)

        model, args["data_column_names"] = train_model(df, target, max_trials)
        model.save(args["folder_path"])
        self.model_storage.json_set("predict_args", args)

    def predict(self, df, args=None):
        """Predicts with best saved AutoKeras model."""
        args = self.model_storage.json_get("predict_args")
        training_df = pd.read_json(args["training_df"])
        model = load_model(args["folder_path"], custom_objects=ak.CUSTOM_OBJECTS)

        df_to_predict = df.copy()
        predictions = get_preds_from_model(df_to_predict, model, args["target"], args["training_data_column_count"], args["data_column_names"])

        # If we used the classifier we need to pre-process the predictions before returning them
        original_y = training_df[args["target"]]
        if not np.issubdtype(original_y.dtype, np.number):
            return format_categorical_preds(predictions, original_y, df_to_predict, args["target"])

        df_to_predict[args["target"]] = predictions
        return df_to_predict
