#!/usr/bin/env python

# Usage:
#   python example_dynamic_time_indexed [problem_id]
# 
from __future__ import print_function, division

import pyexotica as exo
from pyexotica.publish_trajectory import *
import signal, sys
from time import time, sleep
import numpy as np

exo.Setup.init_ros()
sleep(0.2)

configs = [
    '{exotica_examples}/resources/configs/dynamic_time_indexed/01_ilqr_cartpole.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/02_ilqr_lwr.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/03_ilqr_valkyrie.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/04_analytic_ddp_cartpole.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/05_analytic_ddp_lwr.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/06_analytic_ddp_valkyrie.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/07_control_limited_ddp_cartpole.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/08_control_limited_ddp_lwr.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/09_control_limited_ddp_valkyrie.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/10_ilqg_cartpole.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/11_ilqg_lwr.xml',
    '{exotica_examples}/resources/configs/dynamic_time_indexed/12_ilqg_valkyrie.xml',
]


solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/04_analytic_ddp_cartpole.xml')
problem = solver.get_problem()
solution = solver.solve()
costs_ddp_second = problem.get_cost_evolution()

solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/13_ddp_cartpole_no_2nd_order.xml')
problem = solver.get_problem()
solution = solver.solve()
costs_ddp_first = problem.get_cost_evolution()

# solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/07_control_limited_ddp_cartpole.xml')
# problem = solver.get_problem()
# solution = solver.solve()
# cost_cl_ddp = problem.get_cost_evolution()

# solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/14_analytic_ddp_cartpole_1st_order.xml')
# problem = solver.get_problem()
# solution = solver.solve()
# cost_cl_1st_ddp = problem.get_cost_evolution()

# solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/01_ilqr_cartpole.xml')
# problem = solver.get_problem()
# solution = solver.solve()
# costs_ilqr = problem.get_cost_evolution()

# solver = exo.Setup.load_solver( '{exotica_examples}/resources/configs/dynamic_time_indexed/10_ilqg_cartpole.xml')
# problem = solver.get_problem()
# solution = solver.solve()
# costs_ilqg = problem.get_cost_evolution()


import matplotlib.pyplot as plt
plt.figure()
plt.plot(range(len(costs_ddp_first[1])), costs_ddp_first[1], label='analytic ddp 1st order')
plt.plot(range(len(costs_ddp_second[1])), costs_ddp_second[1], label='analytic ddp 2nd order')
# plt.plot(range(len(cost_cl_ddp[1])), cost_cl_ddp[1], label='control-limited ddp (1st order)')
# plt.plot(range(len(cost_cl_ddp[1])), cost_cl_ddp[1], label='control-limited ddp (1st order)')
# plt.plot(range(len(costs_ilqr[1])), costs_ilqr[1], label='ilqr')
# plt.plot(range(len(costs_ilqg[1])), costs_ilqg[1], label='ilqg')
plt.xlabel('Iteration')
plt.ylabel('Cost')
plt.legend()
plt.show()
