# Path: deluca/igpc/rollout.py
# @partial(jax.jit, static_argnums=(1,))
# def rollout(env,
#             cost_func,
#             U_old,
#             k=None,
#             K=None,
#             X_old=None,
#             alpha=1.0,
#             H=None,
#             start_state=None):
#   """
#     Arg List
#     env: The environment to do the rollout on. This is treated as a
#     derstructible copy.
#     U_old: A base open loop control sequence
#     k: open loop gain (iLQR)
#     K: closed loop gain (iLQR)
#     X_old: Previous trajectory to compute gain (iLQR)
#     alpha: Optional multplier to the open loop gain (iLQR)
#     D: Optional noise vectors for the rollout (GPC)
#     F: Linearization shift ??? (GPC)
#     H: The horizon length to perform a rollout.
#     """
#   ## problematic
#   # if H == None:
#   #     H = env.H
#   H = env.H
#   X, U = [None] * (H + 1), [None] * (H)
#   if start_state is None:
#     start_state = env.init()
#   X[0], cost = start_state, 0.0
#   for h in range(H):
#     if k is None:
#       U[h] = U_old[h]
#     else:
#       X_flat = X[h].flatten()
#       X_old_flat = X_old[h].flatten()
#       U[h] = U_old[h] + alpha * k[h] + K[h] @ (X_flat - X_old_flat)
# 
#     X[h + 1], _ = env(X[h], U[h])
#     cost += cost_func(X[h], U[h], env)
#   return X, U, cost
# 
# def compute_ders(env, cost, X, U, H=None):
#   D, F, C = compute_ders_inner(env, cost, X, U, H=None)
#   F = list(zip(*F))
#   C = list(zip(*C))
#   return D, F, C
# 
# Path: deluca/igpc/lqr_solver.py
# def LQR(F, C):
#   H, d = len(C), C[0][0].shape[0]
#   K, k = (
#       [None for _ in range(H)],
#       [None for _ in range(H)],
#   )
# 
#   V_x, V_xx = np.zeros(d), np.zeros((d, d))
#   # V_x, V_xx = jnp.zeros(d), jnp.zeros((d, d))
# 
#   # F_x, F_u = F
#   # C_x, C_u, C_xx, C_uu = C
# 
#   # def loop(carry, args):
#   # v_x, v_xx = carry
#   # f_x, f_u, c_x, c_u, c_xx, c_uu = args
# 
#   # q_x, q_u = c_x + f_x.T @ v_x, c_u + f_u.T @ v_x
#   # q_xx, q_ux, q_uu = (
#   # c_xx + f_x.T @ v_xx @ f_x,
#   # f_u.T @ v_xx @ f_x,
#   # c_uu + f_u.T @ v_xx @ f_u,
#   # )
# 
#   # K, k = -jnp.linalg.inv(q_uu) @ q_ux, -jnp.linalg.inv(q_uu) @ q_u
#   # v_x = q_x - K.T @ q_uu @ k
#   # v_xx = q_xx - K.T @ q_uu @ K
#   # v_xx = (v_xx + v_xx.T) / 2
# 
#   # return (v_x, v_xx), (K, k)
# 
#   # _, (K, k) = jax.lax.scan(loop, (V_x, V_xx), (F_x, F_u, C_x, C_u, C_xx, C_uu))
#   for h in range(H - 1, -1, -1):
#     f_x, f_u = F[h]
#     c_x, c_u, c_xx, c_uu = C[h]
#     Q_x, Q_u = c_x + f_x.T @ V_x, c_u + f_u.T @ V_x
# 
#     Q_xx, Q_ux, Q_uu = (
#         c_xx + f_x.T @ V_xx @ f_x,
#         f_u.T @ V_xx @ f_x,
#         c_uu + f_u.T @ V_xx @ f_u,
#     )
#     K[h], k[h] = -np.linalg.inv(Q_uu) @ Q_ux, -np.linalg.inv(Q_uu) @ Q_u
#     V_x = Q_x - K[h].T @ Q_uu @ k[h]
#     V_xx = Q_xx - K[h].T @ Q_uu @ K[h]
#     V_xx = (V_xx + V_xx.T) / 2
#   return k, K


import os
import time
import jax.numpy as jnp
import jax
from deluca.igpc.rollout import rollout, compute_ders
from deluca.igpc.lqr_solver import LQR
    verbose=True,
    backtracking=True,
):
  alpha, r = ref_alpha, 1
  X, U, c = rollout(env_true, cost_func, U, k, K, X)
  print(f"iLC: t = {-1}, r = {r}, c = {c}")
  prev_loop_fail = False
  for t in range(T):
    _, F, C = compute_ders(env_sim, cost_func, X, U)
    k, K = LQR(F, C)
    if backtracking:
      if prev_loop_fail:
        if verbose:
          print("Backtracking failed - Quitting")
          break

      prev_loop_fail = True
      for alphaC in alpha * 1.1 * 1.1**(-jnp.arange(10)**2):
        r += 1
        XC, UC, cC = rollout(env_true, cost_func, U, k, K, X, alphaC)
        if cC <= c:
          X, U, c, alpha = XC, UC, cC, alphaC
          if verbose:
            print(
                f"iLC (closed+alpha={ref_alpha}): t = {t}, r = {r}, c = {c}, alpha = {alpha}"
            )
            prev_loop_fail = False
          break
    else:
      r += 1

# Path: deluca/igpc/rollout.py
# @partial(jax.jit, static_argnums=(1,))
# def rollout(env,
#             cost_func,
#             U_old,
#             k=None,
#             K=None,
#             X_old=None,
#             alpha=1.0,
#             H=None,
#             start_state=None):
#   """
#     Arg List
#     env: The environment to do the rollout on. This is treated as a
#     derstructible copy.
#     U_old: A base open loop control sequence
#     k: open loop gain (iLQR)
#     K: closed loop gain (iLQR)
#     X_old: Previous trajectory to compute gain (iLQR)
#     alpha: Optional multplier to the open loop gain (iLQR)
#     D: Optional noise vectors for the rollout (GPC)
#     F: Linearization shift ??? (GPC)
#     H: The horizon length to perform a rollout.
#     """
#   ## problematic
#   # if H == None:
#   #     H = env.H
#   H = env.H
#   X, U = [None] * (H + 1), [None] * (H)
#   if start_state is None:
#     start_state = env.init()
#   X[0], cost = start_state, 0.0
#   for h in range(H):
#     if k is None:
#       U[h] = U_old[h]
#     else:
#       X_flat = X[h].flatten()
#       X_old_flat = X_old[h].flatten()
#       U[h] = U_old[h] + alpha * k[h] + K[h] @ (X_flat - X_old_flat)
# 
#     X[h + 1], _ = env(X[h], U[h])
#     cost += cost_func(X[h], U[h], env)
#   return X, U, cost
# 
# def compute_ders(env, cost, X, U, H=None):
#   D, F, C = compute_ders_inner(env, cost, X, U, H=None)
#   F = list(zip(*F))
#   C = list(zip(*C))
#   return D, F, C
# 
# Path: deluca/igpc/lqr_solver.py
# def LQR(F, C):
#   H, d = len(C), C[0][0].shape[0]
#   K, k = (
#       [None for _ in range(H)],
#       [None for _ in range(H)],
#   )
# 
#   V_x, V_xx = np.zeros(d), np.zeros((d, d))
#   # V_x, V_xx = jnp.zeros(d), jnp.zeros((d, d))
# 
#   # F_x, F_u = F
#   # C_x, C_u, C_xx, C_uu = C
# 
#   # def loop(carry, args):
#   # v_x, v_xx = carry
#   # f_x, f_u, c_x, c_u, c_xx, c_uu = args
# 
#   # q_x, q_u = c_x + f_x.T @ v_x, c_u + f_u.T @ v_x
#   # q_xx, q_ux, q_uu = (
#   # c_xx + f_x.T @ v_xx @ f_x,
#   # f_u.T @ v_xx @ f_x,
#   # c_uu + f_u.T @ v_xx @ f_u,
#   # )
# 
#   # K, k = -jnp.linalg.inv(q_uu) @ q_ux, -jnp.linalg.inv(q_uu) @ q_u
#   # v_x = q_x - K.T @ q_uu @ k
#   # v_xx = q_xx - K.T @ q_uu @ K
#   # v_xx = (v_xx + v_xx.T) / 2
# 
#   # return (v_x, v_xx), (K, k)
# 
#   # _, (K, k) = jax.lax.scan(loop, (V_x, V_xx), (F_x, F_u, C_x, C_u, C_xx, C_uu))
#   for h in range(H - 1, -1, -1):
#     f_x, f_u = F[h]
#     c_x, c_u, c_xx, c_uu = C[h]
#     Q_x, Q_u = c_x + f_x.T @ V_x, c_u + f_u.T @ V_x
# 
#     Q_xx, Q_ux, Q_uu = (
#         c_xx + f_x.T @ V_xx @ f_x,
#         f_u.T @ V_xx @ f_x,
#         c_uu + f_u.T @ V_xx @ f_u,
#     )
#     K[h], k[h] = -np.linalg.inv(Q_uu) @ Q_ux, -np.linalg.inv(Q_uu) @ Q_u
#     V_x = Q_x - K[h].T @ Q_uu @ k[h]
#     V_xx = Q_xx - K[h].T @ Q_uu @ K[h]
#     V_xx = (V_xx + V_xx.T) / 2
#   return k, K


# Path: deluca/igpc/ilc.py
import os
import time
import jax.numpy as jnp
import jax
from deluca.igpc.rollout import rollout, compute_ders
from deluca.igpc.lqr_solver import LQR

    verbose=True,
    backtracking=True,
):
  alpha, r = ref_alpha, 1
  X, U, c = rollout(env_true, cost_func, U, k, K, X)
  print(f"iLC: t = {-1}, r = {r}, c = {c}")
  prev_loop_fail = False
  for t in range(T):
    _, F, C = compute_ders(env_sim, cost_func, X, U)
    k, K = LQR(F, C)
    if backtracking:
      if prev_loop_fail:
        if verbose:
          print("Backtracking failed - Quitting")
          break

      prev_loop_fail = True
      for alphaC in alpha * 1.1 * 1.1**(-jnp.arange(10)**2):
        r += 1
        XC, UC, cC = rollout(env_true, cost_func, U, k, K, X, alphaC)
        if cC <= c:
          X, U, c, alpha = XC, UC, cC, alphaC
          if verbose:
            print(
                f"iLC (closed+alpha={ref_alpha}): t = {t}, r = {r}, c = {c}, alpha = {alpha}"
            )
            prev_loop_fail = False
          break
    else:
      r += 1

      XC, UC, cC = rollout(env_true, cost_func, U, k, K, X, alpha)