# %%
import src.bnn.dropout.functions as do_fn
import src.data as data
import src.utils.evaluate as evl
import src.utils.logging as logging
import src.utils.plots as plot

# %%
log_dir = logging.get_log_dir("dropout")
logger = logging.get_logger(path=log_dir)

# %%
data_sets, (X_scaler, y_scaler), label_names = data.load_gaultois_train_test(dropna=False)
X_train, y_train, X_test, y_test = data_sets
scd_label_names = label_names[-1]
rho_train, seebeck_train, kappa_train, zT_train = y_train[scd_label_names].values.T
rho_test, seebeck_test, kappa_test, zT_test = y_test[scd_label_names].values.T
X_train, X_test = X_train.values, X_test.values

# %% [markdown]
# ## Resistivity model with epistemic uncertainty

# %%
rho_ep_pred, rho_ep_var, rho_ep_history, rho_ep_model = do_fn.do_predict(
    X_train, seebeck_train, X_test, seebeck_test, uncert="epistemic"
)

# %%
logger.info(f"rho epistemic mse: {evl.np_mse(rho_ep_pred, seebeck_test)}")
plot.loss_history(rho_ep_history, path=log_dir + "rho_ep_loss.pdf")
evl.gen_plots(seebeck_test, rho_ep_pred, rho_ep_var ** 0.5, path=log_dir, title="rho")

# %% [markdown]
# ## Resistivity model with aleatoric and epistemic uncertainty

# %%
rho_al_ep_pred, rho_al_ep_var, rho_al_ep_history, rho_al_ep_model = do_fn.do_predict(
    X_train, seebeck_train, X_test, seebeck_test
)

# %%
logger.info(f"rho aleatoric+epistemic mse: {evl.np_mse(rho_al_ep_pred, seebeck_test)}")
plot.loss_history(rho_al_ep_history, path=log_dir + "rho_al_ep_loss.pdf")
evl.gen_plots(
    seebeck_test, rho_al_ep_pred, rho_al_ep_var ** 0.5, path=log_dir, title="rho"
)

# %% [markdown]
# ## Seebeck model with epistemic uncertainty

# %%
see_ep_pred, see_ep_var, see_ep_history, see_ep_model = do_fn.do_predict(
    X_train, seebeck_train, X_test, seebeck_test, uncert="epistemic"
)

# %%
logger.info(f"see epistemic mse: {evl.np_mse(see_ep_pred, seebeck_test)}")
plot.loss_history(see_ep_history, path=log_dir + "see_ep_loss.pdf")
evl.gen_plots(seebeck_test, see_ep_pred, see_ep_var ** 0.5, path=log_dir, title="seebeck")

# %% [markdown]
# ## Seebeck model with aleatoric and epistemic uncertainty

# %%
see_al_ep_pred, see_al_ep_var, see_al_ep_history, see_al_ep_model = do_fn.do_predict(
    X_train, seebeck_train, X_test, seebeck_test
)

# %%
logger.info(f"see aleatoric+epistemic mse: {evl.np_mse(see_al_ep_pred, seebeck_test)}")
plot.loss_history(see_al_ep_history, path=log_dir + "see_al_ep_loss.pdf")
evl.gen_plots(
    seebeck_test, see_al_ep_pred, see_al_ep_var ** 0.5, path=log_dir, title="seebeck"
)

# %% [markdown]
# ## Thermal conducitivity model with epistemic uncertainty

# %%
kappa_ep_pred, kappa_ep_var, kappa_ep_history, kappa_ep_model = do_fn.do_predict(
    X_train, kappa_train, X_test, kappa_test, uncert="epistemic"
)

# %%
logger.info(f"kappa epistemic mse: {evl.np_mse(kappa_ep_pred, kappa_test)}")
plot.loss_history(kappa_ep_history, path=log_dir + "kappa_ep_loss.pdf")
evl.gen_plots(kappa_test, kappa_ep_pred, kappa_ep_var ** 0.5, path=log_dir, title="kappa")

# %% [markdown]
# ## Thermal conducitivity model with aleatoric and epistemic uncertainty

# %%
kappa_al_ep_pred, kappa_al_ep_var, kappa_al_ep_history, _ = do_fn.do_predict(
    X_train, kappa_train, X_test, kappa_test
)

# %%
logger.info(f"kappa aleatoric+epistemic mse: {evl.np_mse(kappa_al_ep_pred, kappa_test)}")
plot.loss_history(kappa_al_ep_history, path=log_dir + "kappa_al_ep_loss.pdf")
evl.gen_plots(
    kappa_test, kappa_al_ep_pred, kappa_al_ep_var ** 0.5, path=log_dir, title="kappa"
)
