###
# Create some plots of the robot example
# needs a result file generated by robot_nl_run_example.py
###
import matplotlib.pylab as plt
import numpy as np

from ellipsoid_utils import splitSigmaIntoBlocksList
from latexify import latexify
from robot_nl_plotutils import *

latexify(fig_width=4)

saveplots = True
showplots = False
gamma_conf = 3          # confidence level for ellipsoids representing normal distribution

outfolder = 'robot_nl_results/'
resfile = 'res_2022-11-22-15-00-08.npy'

results_dict = np.load(outfolder + resfile, allow_pickle=True)
results_dict = results_dict[()]
params = results_dict["params"]
N_cl = results_dict["N_cl"]

#%% open loop /  OCP solutions
# for each controller, plot the solution of the very first OCP

Xopt_nom, Uopt_nom = results_dict["traj_ocp"]["nommpc"]
plotTrajectoryInTime(Xopt_nom, Uopt_nom, params, gamma=gamma_conf, title=r'OCP nominal')
plt.tight_layout()
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_nom_time.pdf')
plotTrajectoryInSpace(Xopt_nom, params, gamma=gamma_conf, title=r'OCP nominal')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_nom_space.pdf')

Xopt_ol, Uopt_ol, P_opt_ol = results_dict["traj_ocp"]["olsmpc"]
plotTrajectoryInTime(Xopt_ol, Uopt_ol, params, P=P_opt_ol, gamma=gamma_conf, title=r'OCP open loop stochastic')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_ol_time.pdf')
plotTrajectoryInSpace(Xopt_ol, params, P=P_opt_ol, gamma=gamma_conf, title=r'OCP open loop stochastic')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_ol_space.pdf')

Xopt_cl, Uopt_cl, P_opt_cl, K_fb_opt_cl = results_dict["traj_ocp"]["sfsmpc"]
plotTrajectoryInTime(Xopt_cl, Uopt_cl, params, P=P_opt_cl, K_fb=K_fb_opt_cl, gamma=gamma_conf, title=r'OCP state feedback stochastic')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_cl_time.pdf')
plotTrajectoryInSpace(Xopt_cl, params, P=P_opt_cl, gamma=gamma_conf,  title=r'OCP state feedback stochastic')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_cl_space.pdf')

Xopt_dual, Uopt_dual, Sigma_opt_dual, K_fb_opt_dual, K_ob_opt_dual = results_dict["traj_ocp"]["ofsmpc"]
P_opt_dual, Phat_opt_dual, Pbrev_opt_dual = splitSigmaIntoBlocksList(Sigma_opt_dual)
plotTrajectoryInTime(Xopt_dual, Uopt_dual, params, P=P_opt_dual, Phat=Phat_opt_dual, Pbrev=Pbrev_opt_dual, K_fb=K_fb_opt_dual, gamma=gamma_conf, title=r'OCP dual')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_dual_time.pdf')
plotTrajectoryInSpace(Xopt_dual, params, P=P_opt_dual, gamma=gamma_conf, title=r'OCP state feedback stochastic, with $P$ plotted for uncertainty')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_dual_space_P.pdf')
plotTrajectoryInSpace(Xopt_dual, params, P=Phat_opt_dual, gamma=gamma_conf,  title=r'OCP state feedback stochastic, with $\hat P$ plotted for uncertainty')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'OCP_dual_space_Phat.pdf')

#%%  MPC closed loop trajectories
# example trajectories
X_true_nom , U_true_nom , X_hat_nom , P_hat_nom , _ = results_dict["traj_closedloop"][1][0]
X_true_ol  , U_true_ol  , X_hat_ol  , P_hat_ol  , _ = results_dict["traj_closedloop"][1][1]
X_true_cl  , U_true_cl  , X_hat_cl  , P_hat_cl  , _ = results_dict["traj_closedloop"][1][2]
X_true_dual, U_true_dual, X_hat_dual, P_hat_dual, _ = results_dict["traj_closedloop"][1][3]

plotCompareTrueAndEstimatedTraj(X_true_nom, X_hat_nom, params, Phat=P_hat_nom, gamma=gamma_conf, title='Nominal MPC')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'MPC_nom_space.pdf')
plotCompareTrueAndEstimatedTraj(X_true_ol, X_hat_ol, params, Phat=P_hat_ol, gamma=gamma_conf, title='Open loop SMPC')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'MPC_ol_space.pdf')
plotCompareTrueAndEstimatedTraj(X_true_cl, X_hat_cl, params, Phat=P_hat_cl, gamma=gamma_conf, title='closed loop SMPC')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'MPC_cl_space.pdf')
plotCompareTrueAndEstimatedTraj(X_true_dual, X_hat_dual, params, Phat=P_hat_dual, gamma=gamma_conf, title='dual MPC')
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'MPC_dual_space.pdf')

plotCompareMPCTrajInSpace(results_dict["traj_closedloop"][1], params, gamma=gamma_conf)
if saveplots: plt.tight_layout(pad=.5,h_pad=.5, w_pad=0); plt.savefig(outfolder + 'MPC_comp_traj.pdf')

# compare constraint viol for sample traj
traj_list = []
traj_list += [(X_true_nom, 'Nominal MPC')]
traj_list += [(X_true_ol, 'open loop SMPC')]
traj_list += [(X_true_cl, 'closed loop SMPC')]
traj_list += [(X_true_dual, 'Dual MPC')]
plotCompareConstraintViolation(traj_list, params)
if saveplots: plt.tight_layout(); plt.savefig(outfolder + 'MPC_comp_constr.pdf')

plotCompareConstraintViolationMany(results_dict["traj_closedloop"], params, N_cl)
if saveplots: plt.tight_layout(pad=.5,h_pad=.5, w_pad=0); plt.savefig(outfolder + 'MPC_comp_constr_many.pdf')

plotCompareMPCtrajInSpaceAndConstraintViolationMany(results_dict["traj_closedloop"], params, N_cl, gamma=gamma_conf )
if saveplots: plt.tight_layout(pad=.5); plt.savefig(outfolder + 'MPC_comp_traj_and_constr_many.pdf')

plotMPCSuccessRatio(results_dict["traj_closedloop"], params, None)
if saveplots: plt.tight_layout(pad=0.3); plt.savefig(outfolder + 'MPC_comp_success.pdf')

if showplots: plt.show()
