def compute_A(x, u, w, i):
    r = x[:,0]
    alpha = x[:,1]
    beta = x[:,2]
    Vx = x[:,3]
    Vy = x[:,4]
    Vz = x[:,5]
    m = x[:,6]
    phi = x[:,7]
    psi = x[:,8]

    T = u[:,0]
    omega_phi = u[:,1]
    omega_psi = u[:,2]
    A = np.array([  [0,0,0,1,0,0,0,0,0],
                    [-Vy[i]/(r[i]**2*drake_math.cos(beta[i])), 				0,	Vy[i]*drake_math.sin(beta[i])/(r[i]*drake_math.cos(beta[i])**2), 	0, 		1/(r[i]*drake_math.cos(beta[i])), 		0, 				0, 							0, 						0],
                    [-Vz[i]/r[i]**2, 							0, 	0, 								0, 		0, 						1/r[i], 			0, 							0, 						0],
                    [-(Vy[i]**2+Vz[i]**2)/r[i]**2, 				0, 	0, 								0, 		2*Vy[i]/r[i], 				2*Vz[i]/r[i], 		-T[i]*drake_math.sin(phi[i])/(m[i]+w)**2, 			T[i]*drake_math.cos(phi[i])/(m[i]+w), 			0],
                    [Vx[i]*Vy[i]/r[i]**2-Vy[i]*Vz[i]*drake_math.tan(beta[i])/r[i]**2,	0,	Vy[i]*Vz[i]*(1+drake_math.tan(beta[i])**2)/r[i], 		-Vy[i]/r[i],	-Vx[i]/r[i]+Vz[i]*drake_math.tan(beta[i])/r[i],	Vy[i]*drake_math.tan(beta[i])/r[i],	-T[i]*drake_math.cos(phi[i])*drake_math.cos(psi[i])/(m[i]+w)**2,	-T[i]*drake_math.sin(phi[i])*drake_math.cos(psi[i])/(m[i]+w),	-T[i]*drake_math.cos(phi[i])*drake_math.sin(psi[i])/(m[i]+w)],
                    [Vx[i]*Vz[i]/r[i]**2+Vy[i]**2*drake_math.tan(beta[i])/r[i]**2, 	0, 	-Vy[i]**2*(1+drake_math.tan(beta[i])**2)/r[i], 		-Vz[i]/r[i], 	-2*Vy[i]*drake_math.tan(beta[i])/r[i], 		-Vx[i]/r[i], 			-T[i]*drake_math.cos(phi[i])*drake_math.sin(psi[i])/(m[i]+w)**2,	-T[i]*drake_math.sin(phi[i])*drake_math.sin(psi[i])/(m[i]+w),	T[i]*drake_math.cos(phi[i])*drake_math.cos(psi[i])/(m[i]+w)],
                    [0, 0, 0, 0, 0, 0, 0, 0,0],
                    [0, 0, 0, 0, 0, 0, 0, 0,0],
                    [0, 0, 0, 0, 0, 0, 0, 0,0]])

    return A

def compute_B(x, u, w, i):

    m = x[:,6]
    phi = x[:,7]
    psi = x[:,8]

    B = np.array([  [0, 0, 0],
                    [0, 0, 0],
                    [0, 0, 0],
                    [drake_math.sin(phi[i])/(m[i]+w), 0, 0],
                    [drake_math.cos(phi[i])*drake_math.cos(psi[i])/(m[i]+w), 0, 0],
                    [drake_math.cos(phi[i])*drake_math.sin(psi[i])/(m[i]+w), 0, 0],
                    [-1/(Isp*g), 0, 0],
                    [0, 1, 0],
                    [0, 0, 1]])
    return B

def compute_G(x, u, w, i):
    m = x[:,6]
    phi = x[:,7]
    psi = x[:,8]
    T = u[:,0]
    G = np.array([ [0,0,0,0,0,0,0,0,0],
                   [0,0,0,0,0,0,0,0,0],
                   [0,0,0,0,0,0,0,0,0],
                   [0,0,0,0,0,0,-T[i]*drake_math.sin(phi[i])/(m[i]+w)**2,0,0],
                   [0,0,0,0,0,0,-T[i]*drake_math.cos(phi[i])*drake_math.cos(psi[i])/(m[i]+w)**2,0,0],
                   [0,0,0,0,0,0,-T[i]*drake_math.cos(phi[i])*drake_math.sin(psi[i])/(m[i]+w)**2,0,0],
                   [0,0,0,0,0,0,0,0,0],
                   [0,0,0,0,0,0,0,0,0],
                   [0,0,0,0,0,0,0,0,0]])
    return G

def ell_w(x, u, D, E_1, Q_l, R_l, Ql_n, Q, R, N):
    
    ell = 0
    w = 20*np.random.rand() # [0, 1)
    A = np.zeros((N,n_x,n_x), dtype=AutoDiffXd)
    B = np.zeros((N,n_x,n_u), dtype=AutoDiffXd)
    G = np.zeros((N,n_x,n_x), dtype=AutoDiffXd)

    K = [0]*(N-1)

    K = np.zeros((N-1,n_u,n_x), dtype=AutoDiffXd)    
    E = np.zeros((N,n_x,n_x), dtype=AutoDiffXd)    
    H = np.zeros((N,n_x,n_x), dtype=AutoDiffXd)    
    for i in range(N):
        A[i] = compute_A(x, u, w, i)
        B[i] = compute_B(x, u, w, i)
        G[i] = compute_G(x, u, 0, i)
    P = [0]*N
    P[-1] = Ql_n
    
    # compute Riccati difference equation for i in 1:N
    for i in reversed(range(1,N)):
        P[i-1] = Q + A[i].T@P[i]@A[i]-A[i].T@P[i]@B[i]@(R + B[i].T@P[i]@B[i]).T@(B[i].T@P[i]@A[i])
    for i in range(N-1):
        K[i] = (R + B[i].T@P[i+1]@B[i]).T@(B[i].T@P[i+1]@A[i])
    
    H[0] = np.zeros_like(Q)
    E[0] = E_1
    for i in range(N-1):
        ell += np.trace(Q_l + K[i].T @ R_l @ K[i] @ E[i])
        E[i+1] = (A[i] - B[i] @ K[i]) @ E[i] @ (A[i] - B[i] @ K[i]).T + \
                 (A[i] - B[i] @ K[i]) @ H[i] @ G[i].T + \
                  G[i] @ H[i].T @  (A[i] - B[i] @ K[i]).T + G[i] @ D @ G[i].T
        H[i+1] =  (A[i] - B[i] @ K[i]) @ H[i] + G[i] @ D
    ell += np.trace(Ql_n @ E[-1])
    return ell

def AddCost(x, u, use_dirtrel=False):
    TT = 0
    for i in range(N-1):
        TT += (u[i,0]**2 + u[i+1,0]**2)
    if use_dirtrel == False:
        prog.AddQuadraticCost(TT)
    else:
        D = 20*np.eye(9)
        E_1 = np.eye(9)
        TT += ell_w(x, u, D, E_1, Q, R,Q,Q,R,N,0,6)
        TT += ell_w(x, u, D, E_1, Q, R,Q,Q,R,N,6,N-1)
        prog.AddCost(TT)
