import os
import numpy as np
import pandas as pd

from time import time
import matplotlib.pyplot as plt

from ast import literal_eval

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error

from joblib import dump, load

from crabnet.utils.utils import get_cbfv
from crabnet.utils.utils import CONSTANTS

plt.rcParams.update({"font.size": 16})


# %%
cons = CONSTANTS()

bm_names = cons.benchmark_props
bm_pretty = cons.benchmark_names_dict
mb_names = cons.matbench_props
mb_pretty = cons.matbench_names_dict
elem_props = cons.eps
elem_props = ["magpie"]

mat_props = mb_names + bm_names

model_names = {
    "RandomForestRegressor": RandomForestRegressor,
}


# %%
def concatenate_best_results(directory):
    files = os.listdir(directory)

    csv_list = [file for file in files if ".csv" in file if "best_" in file]
    csv_paths = [os.path.join(directory, file) for file in csv_list]

    df = pd.DataFrame()

    for file in csv_paths:
        df_file = pd.read_csv(file)
        df = pd.concat([df, df_file], axis=0, ignore_index=True)

    df = df[
        [
            "estimator",
            "elem_prop",
            "mat_prop",
            "params",
            "mean_test_neg_MAE",
            "mean_test_r2",
        ]
    ]

    return df


def get_best_model(df, mat_prop, sortby="mean_test_neg_MAE"):
    df_mp = df[df["mat_prop"] == mat_prop]
    df_best_row = df_mp.sort_values(sortby, axis=0, ascending=False).iloc[0]

    params = df_best_row["params"]
    params = literal_eval(params)

    est_name = df_best_row["estimator"]
    est_class = model_names[est_name]
    estimator = est_class(**params)

    elem_prop = df_best_row["elem_prop"]

    return est_name, estimator, params, elem_prop


# %%
if __name__ == "__main__":
    metrics_path = r"metrics/rf_gridsearch/"

    bm_predictions_path = r"metrics/rf_benchmark__predictions/"
    mb_predictions_path = r"metrics/rf_matbench__predictions/"
    clf_save_path = f"models/trained_rf/"
    os.makedirs(clf_save_path, exist_ok=True)
    os.makedirs(metrics_path, exist_ok=True)
    os.makedirs(bm_predictions_path, exist_ok=True)
    os.makedirs(mb_predictions_path, exist_ok=True)

    bm_props_dir = r"data/benchmark_data/"
    mb_props_dir = r"data/matbench_cv/"

    df_best_classics = concatenate_best_results(metrics_path)

    retrain_columns = ["estimator_name", "estimator", "params", "mat_prop", "elem_prop"]
    df_models_to_retrain = pd.DataFrame(columns=retrain_columns)

    for mp in mat_props:
        output = get_best_model(df_best_classics, mp)
        est_name, estimator, params, elem_prop = output

        row_dict = {
            "estimator_name": est_name,
            "estimator": estimator,
            "params": params,
            "mat_prop": mp,
            "elem_prop": elem_prop,
        }
        df_models_to_retrain = df_models_to_retrain.append(row_dict, ignore_index=True)

    df_retrain_columns = [
        "estimator_name",
        "fit_time",
        "mae_train",
        "mse_train",
        "r2_train",
        "mae_test",
        "mse_test",
        "r2_test",
        "params",
        "mat_prop",
        "elem_prop",
    ]
    df_retrain_results = pd.DataFrame(columns=df_retrain_columns)
    ti_retrain_classics = time()

    mb_cvs = 5
    for i in range(len(df_models_to_retrain)):
        output = df_models_to_retrain.iloc[i].values
        est_name, estimator, param, mat_prop, elem_prop = output

        if mat_prop in bm_names:
            df_predictions = pd.DataFrame()
            trainpath = os.path.join(bm_props_dir, mat_prop, "train.csv")
            valpath = os.path.join(bm_props_dir, mat_prop, "val.csv")
            testpath = os.path.join(bm_props_dir, mat_prop, "test.csv")

            X, y, form, skipped = get_cbfv(trainpath, elem_prop=elem_prop)
            X_val, y_val, form_val, skipped_val = get_cbfv(valpath, elem_prop=elem_prop)
            X_test, y_test, form_test, skipped_test = get_cbfv(
                testpath, elem_prop=elem_prop
            )

            X = pd.concat([X, X_val], axis=0, ignore_index=True)
            y = pd.concat([y, y_val], axis=0, ignore_index=True)
            form = pd.concat([form, form_val], axis=0, ignore_index=True)

            print(f"fitting {mat_prop} with {elem_prop} using {est_name}")
            if X.shape[0] > 2000:
                print(
                    f"there are {X.shape[0]} data samples, " f"this may take a while..."
                )
            ti_model = time()
            estimator.fit(X, y)
            dt_model = time() - ti_model
            print(f"finished fitting {mat_prop} with {elem_prop} using {est_name}")
            print(f"fitting time: {dt_model:0.4f} s")

            savepath = f"models/trained_rf/{mat_prop}.joblib"
            print(f"Saving rf model {savepath}")
            dump(estimator, savepath, compress=3, protocol=5)

            print(f"evaluating fitted model on whole trainset and testset")
            ti_eval = time()
            target_test = y_test
            pred_test = estimator.predict(X_test)

            target_train = y
            pred_train = estimator.predict(X)

            r2_test = r2_score(target_test, pred_test)
            mae_test = mean_absolute_error(target_test, pred_test)
            mse_test = mean_squared_error(target_test, pred_test)

            r2_train = r2_score(target_train, pred_train)
            mae_train = mean_absolute_error(target_train, pred_train)
            mse_train = mean_squared_error(target_train, pred_train)

            dt_eval = time() - ti_eval
            print(f"evaluation time: {dt_eval:0.4f} s")

            df_retrain_row = {
                "estimator_name": est_name,
                "fit_time": dt_model,
                "mae_train": mae_train,
                "mse_train": mse_train,
                "r2_train": r2_train,
                "mae_test": mae_test,
                "mse_test": mse_test,
                "r2_test": r2_test,
                "params": param,
                "mat_prop": mat_prop,
                "cv": "NaN",
                "elem_prop": elem_prop,
            }

            df_retrain_results = df_retrain_results.append(
                df_retrain_row, ignore_index=True
            )

            df_predictions["composition"] = form_test
            df_predictions["target"] = target_test
            df_predictions["pred"] = pred_test
            outfile = os.path.join(bm_predictions_path, f"{mat_prop}_test_output.csv")
            df_predictions.to_csv(outfile)
            print(f"mae_test for {mat_prop}: {mae_test}")

        elif mat_prop in mb_names:
            trainpath_pre = os.path.join(mb_props_dir, mat_prop, "train")
            valpath_pre = os.path.join(mb_props_dir, mat_prop, "val")
            testpath_pre = os.path.join(mb_props_dir, mat_prop, "test")

            maes = []
            for cv in range(mb_cvs):
                df_predictions = pd.DataFrame()
                trainpath = f"{trainpath_pre}{cv}.csv"
                valpath = f"{valpath_pre}{cv}.csv"
                testpath = f"{testpath_pre}{cv}.csv"

                X, y, form, skipped = get_cbfv(trainpath, elem_prop=elem_prop)
                X_val, y_val, form_val, skipped_val = get_cbfv(
                    valpath, elem_prop=elem_prop
                )
                X_test, y_test, form_test, skipped_test = get_cbfv(
                    testpath, elem_prop=elem_prop
                )

                X = pd.concat([X, X_val], axis=0, ignore_index=True)
                y = pd.concat([y, y_val], axis=0, ignore_index=True)
                form = pd.concat([form, form_val], axis=0, ignore_index=True)

                print(f"fitting {mat_prop} with {elem_prop} using {est_name}")
                print(f"cv {cv}")
                if X.shape[0] > 2000:
                    print(
                        f"there are {X.shape[0]} data samples, "
                        f"this may take a while..."
                    )
                ti_model = time()
                estimator.fit(X, y)
                dt_model = time() - ti_model
                print(f"finished fitting {mat_prop} with {elem_prop} using {est_name}")
                print(f"fitting time: {dt_model:0.4f} s")

                savepath = f"models/trained_rf/{mat_prop}_cv{cv}.joblib"
                print(f"Saving rf model {savepath}")
                dump(estimator, savepath, compress=3, protocol=5)

                print(f"evaluating fitted model on whole trainset and testset")
                ti_eval = time()
                target_test = y_test
                pred_test = estimator.predict(X_test)

                target_train = y
                pred_train = estimator.predict(X)

                r2_test = r2_score(target_test, pred_test)
                mae_test = mean_absolute_error(target_test, pred_test)
                mse_test = mean_squared_error(target_test, pred_test)

                r2_train = r2_score(target_train, pred_train)
                mae_train = mean_absolute_error(target_train, pred_train)
                mse_train = mean_squared_error(target_train, pred_train)

                dt_eval = time() - ti_eval
                print(f"evaluation time: {dt_eval:0.4f} s")

                df_retrain_row = {
                    "estimator_name": est_name,
                    "fit_time": dt_model,
                    "mae_train": mae_train,
                    "mse_train": mse_train,
                    "r2_train": r2_train,
                    "mae_test": mae_test,
                    "mse_test": mse_test,
                    "r2_test": r2_test,
                    "params": param,
                    "mat_prop": mat_prop,
                    "cv": cv,
                    "elem_prop": elem_prop,
                }

                df_retrain_results = df_retrain_results.append(
                    df_retrain_row, ignore_index=True
                )

                df_predictions["composition"] = form_test
                df_predictions["target"] = target_test
                df_predictions["pred"] = pred_test
                outfile = os.path.join(
                    mb_predictions_path, f"{mat_prop}_test_output_cv{cv}.csv"
                )
                df_predictions.to_csv(outfile)
                print(f"mae_test: {mae_test}")
                maes.append(mae_test)
            print(f"average mae_test for {mat_prop}: {np.mean(maes)}")

    dt_retrain_classics = time() - ti_retrain_classics
    print("*********** retrain_classics finished ***********")
    print(f"retrain_classics finished, elapsed time: " f"{dt_retrain_classics:0.4g} s")
    print("*********** retrain_classics finished ***********")

    outfile = os.path.join(metrics_path, "retrained_test_scores.csv")
    df_retrain_results.to_csv(outfile, index=False)
