# %%
import functools as ft

import tensorflow as tf
import tensorflow_probability as tfp

import src.bnn.dropout.functions as do_fn
import src.bnn.map.functions as map_fn
import src.data as data
import src.functions as src_fn
import src.gp.functions as gp_fn
import src.rf.functions as rf_fn
import src.utils.evaluate as evl
import src.utils.ipython as ipython
import src.utils.logging as logging
import src.utils.plots as plot

# import os

# # %%
# # Check os.getenv("save_to_disk") == "True" to see if
# # logs, plots and other results should be saved to disk.
# os.environ["save_to_disk"] = "True"

# %%
log_dir = logging.get_log_dir("leaderboard/single_run")
logger = logging.get_logger(path=log_dir)

# %%
data_sets, (X_scaler, y_scaler), label_names = data.load_gaultois_train_test()
X_train, y_train, X_test, y_test = data_sets
label_names, log_label_names, scd_label_names = label_names

data_sets = (X_train, y_train[scd_label_names].T, X_test, y_test[scd_label_names].T)

scd_label_names_tex = [
    r"$\log\ \tilde\rho$",
    r"$\log\ \tilde S^2$",
    r"$\log\ \tilde\kappa$",
    r"$\log\ \tilde{zT}$",
]
log_label_names_tex = [r"$\log\ \rho$", r"$\log\ S^2$", r"$\log\ \kappa$", r"$\log\ zT$"]


# %% [markdown]
# ---
# # Random Forest Regression
# ---

# %%
rf_y_preds_scd, rf_y_vars_scd, rf_models = src_fn.predict_multiple_labels(
    rf_fn.rf_predict, *data_sets
)

# %%
rf_y_preds, rf_y_vars = src_fn.back_transform_labels(
    y_scaler, rf_y_preds_scd, rf_y_vars_scd
)

# %%
rf_mses = evl.get_mses_as_df(
    [y_test[scd_label_names], rf_y_preds_scd], [y_test[label_names], rf_y_preds]
)
rf_mses

# %%
for label, y_true, y_pred, y_var in zip(
    scd_label_names_tex,
    y_test[scd_label_names].values.T,
    rf_y_preds_scd.values.T,
    rf_y_vars_scd.values.T,
):
    ipython.display_md(f"# {label}", center=True)
    evl.gen_plots(y_true, y_pred, y_var ** 0.5, path=log_dir, title=label)

# %%
rf_y_preds["zT_log_computed"] = evl.compute_log_zT(
    *rf_y_preds[log_label_names[:3]].values.T, y_test["T_log"]
).reset_index(drop=True)
zT_mse = ((y_test["zT_log"] - rf_y_preds["zT_log_computed"]) ** 2).mean()
logger.info(f"zT_log_computed mse: {'{:.4g}'.format(zT_mse)}")


# %% [markdown]
# ---
# # Maximum A Posteriori Neural Network
# ---

# %%
# Per-label weight and bias priors in the order rho, seebeck, kappa, zT.
weight_priors = [tfp.distributions.Normal(0, std) for std in [0.1, 0.1, 0.1, 0.1]]
bias_priors = [tfp.distributions.Normal(0, std) for std in [1.0, 1.0, 1.0, 1.0]]
map_predictors = [
    ft.partial(map_fn.map_predict, *priors) for priors in zip(weight_priors, bias_priors)
]

# %%
map_results = src_fn.predict_multiple_labels(map_predictors, *data_sets)
map_y_preds_scd, map_y_vars_scd, map_log_probs, map_initial_states = map_results

# %%
map_y_preds, map_y_vars = src_fn.back_transform_labels(
    y_scaler, map_y_preds_scd, map_y_vars_scd
)

# %%
map_mses = evl.get_mses_as_df(
    [y_test[scd_label_names], map_y_preds_scd], [y_test[label_names], map_y_preds]
)
map_mses


# %%
for label, map_log_prob in zip(scd_label_names_tex, map_log_probs.values.T):
    ipython.display_md(f"# {label}", center=True)
    plot.log_probs(
        zip(["train", "test"], map_log_prob),
        title=label,
        path=log_dir + f"{label}_map_log_probs.pdf",
    )

# %%
for label, y_true, y_pred, y_var in zip(
    scd_label_names_tex,
    y_test[scd_label_names].values.T,
    map_y_preds_scd.values.T,
    map_y_vars_scd.values.T,
):
    ipython.display_md(f"# {label}", center=True)
    evl.gen_plots(y_true, y_pred, y_var ** 0.5, path=log_dir, title=label)

# %%
map_y_preds["zT_log_computed"] = evl.compute_log_zT(
    *map_y_preds[log_label_names[:3]].values.T, y_test["T_log"]
).reset_index(drop=True)
zT_mse = ((y_test["zT_log"] - map_y_preds["zT_log_computed"]) ** 2).mean()
logger.info(f"zT_log_computed mse: {'{:.4g}'.format(zT_mse)}")


# %% [markdown]
# ---
# # Gaussian Process Regression
# ---

# %%
# GPR might benefit from PCA
# from sklearn.decomposition import PCA
# n_features = X_train.shape[1]
# pca = PCA(n_components=n_features // 4)
# X_train = pca.fit(X_train).transform(X_train)
# X_test = pca.transform(X_test)

gp_y_preds_scd, gp_y_vars_scd, gp_models = src_fn.predict_multiple_labels(
    gp_fn.gp_predict, *data_sets
)

# %%
gp_y_preds, gp_y_vars = src_fn.back_transform_labels(
    y_scaler, gp_y_preds_scd, gp_y_vars_scd
)

# %%
gp_mses = evl.get_mses_as_df(
    [y_test[scd_label_names], gp_y_preds_scd], [y_test[label_names], gp_y_preds]
)
gp_mses

# %%
for label, y_true, y_pred, y_std in zip(
    scd_label_names_tex,
    y_test[scd_label_names].values.T,
    gp_y_preds_scd.values.T,
    gp_y_vars_scd.values.T ** 0.5,
):
    ipython.display_md(f"# {label}", center=True)
    evl.gen_plots(y_true, y_pred, y_std, path=log_dir, title=label)


# %% [markdown]
# ---
# # Dropout Neural Network
# ---

# %%
# one of ["aleatoric", "epistemic", "aleatoric+epistemic"]
do_uncertainty = "aleatoric+epistemic"
tb_cb = tf.keras.callbacks.TensorBoard(log_dir=log_dir)

do_y_preds_scd, do_y_vars_scd, do_histories, do_models = src_fn.predict_multiple_labels(
    ft.partial(do_fn.do_predict, uncert=do_uncertainty, cbs=[tb_cb], epochs=500),
    *data_sets,
)

# %%
do_y_preds, do_y_vars = src_fn.back_transform_labels(
    y_scaler, do_y_preds_scd, do_y_vars_scd
)

# %%
do_mses = evl.get_mses_as_df(
    [y_test[scd_label_names], do_y_preds_scd], [y_test[label_names], do_y_preds]
)
do_mses

# %%
for label, y_true, y_pred, y_std, hist in zip(
    scd_label_names_tex,
    y_test[scd_label_names].values.T,
    do_y_preds_scd.values.T,
    do_y_vars_scd.values.T ** 0.5,
    do_histories.values(),
    # models,
):
    ipython.display_md(f"# {label}", center=True)
    do_fn.save_loss_history(hist, path=log_dir)
    evl.gen_plots(y_true, y_pred, y_std, path=log_dir, title=label)
    # do_fn.plot_model(model, log_dir)
