diff --git a/AHhw1_2.py b/AHhw1_2.py
index 8464073..ecab9ba 100644
--- a/AHhw1_2.py
+++ b/AHhw1_2.py
@@ -1,29 +1,43 @@
 import numpy as np
 import matplotlib.pyplot as plt
+from scipy.integrate import solve_bvp
 
 
-# define function
-def f(P):
+# define ricatti equation
+def riccati(P):
     return -P @ A - A.T @ P + P @ B / R @ B.T @ P - Q
 
 
+def TBVBP(t, x, p):
+    return np.vstack((A@x-B/R@B.T@p, -Q@x-A.T@p))
+
+
+def bc(ya, yb, p):
+    return np.array([X_initial, np.zeros((2, 1)), np.zeros((2, 1))])
+
+
 # Finds value of y for a given x using step size h
 # and initial value y0 at x0.
-def rk4(P_final, t_final, dt, f):
+def rk4(sol_final, t, dt, f):
     # Iterate for number of iterations
-    P = np.zeros((2, 2 * int(t_final / dt)))
-    P[:, (2 * int(t_final / dt) - 2):(2 * int(t_final / dt))] = P_final
-    for i in range(int(t_final / dt) - 1, 0, -1):
-        print(P[:, 2 * i:2 * i + 2])
-        k1 = dt * f(P[:, 2 * i:2 * i + 2])
-        k2 = dt * f(P[:, 2 * i:2 * i + 2] - k1 / 2.)
-        k3 = dt * f(P[:, 2 * i:2 * i + 2] - k2 / 2.)
-        k4 = dt * f(P[:, 2 * i:2 * i + 2] - k3)
+    row = len(sol_final)
+    columns = len(sol_final[0])
+    sol = np.zeros((row, columns * len(t)))
+    sol[:, (columns * len(t) - columns):(columns * len(t))] = sol_final
+    for i in range(len(t) - 1, 0, -1):
+        k1 = dt * f(sol[:, columns * i:(columns * i + columns)])
+        k2 = dt * f(sol[:, columns * i:(columns * i + columns)] - k1 / 2.)
+        k3 = dt * f(sol[:, columns * i:(columns * i + columns)] - k2 / 2.)
+        k4 = dt * f(sol[:, columns * i:(columns * i + columns)] - k3)
         # Update next value of p
-        P[:, 2 * i - 2:2 * i] = P[:, 2 * i:2 * i + 2] - (1 / 6.) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
-    return P
+        sol[:, (columns * i - columns):columns * i] = sol[:, columns * i:(columns * i + columns)] - (1 / 6.) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
+    return sol
 
 
+# time step
+t = np.arange(0, 10, 1e-3)
+dt = t[1]
+
 # Define variable
 Q = np.array([[2, 0], [0, 0.01]])
 R = 0.1
@@ -31,27 +45,45 @@ def rk4(P_final, t_final, dt, f):
 A = np.array([[0, 1], [-1.6, -0.4]])
 B = np.array([[0], [1]])
 P_10 = P1
-t_final = 10
-dt = 0.01
-
+x_10 = np.zeros((4, 1))
+Arow1 = np.append(A, -B/R@B.T, axis=1)
+Arow2 = np.append(-Q, -A.T, axis=1)
+A_1 = np.append(Arow1, Arow2, axis=0)
+y = np.zeros((2, len(t)))
 # find P
-P = rk4(P_10, t_final, dt, f)
+P = rk4(P_10, t, dt, riccati)
+
+# find TPBVP solution
+sol = solve_bvp(TBVBP, bc, t, y, 0)
 
 # simulation
-X = np.zeros((2, int(t_final / dt)))
-u = np.zeros((1, int(t_final / dt)))
+X = np.zeros((2, len(t)))
+u = np.zeros((1, len(t)))
+X_initial = np.array([[10], [0]])
 X[0:2, 0:1] = np.array([[10], [0]])
 
-for i in range(0, int(t_final / dt) - 1):
-    u[0:1, i:i + 1] = -1 / R * B.T @ P[0:2, 2 * i:2 * i + 2] @ X[0:2, i:i + 1]
-    X[0:2, i + 1:i + 2] = X[0:2, i:i + 1] + dt * (A @ X[0:2, i:i + 1] + B @ u[0:1, i:i + 1])
-
-# plot
-plt.plot(X[0, :], X[1, :])
-plt.scatter(X[0, 0], X[1, 0], label = 'start point')
-plt.scatter(X[0, int(t_final / dt) - 1], X[1, int(t_final / dt) - 1], label = 'final point')
-plt.xlabel("x1")
-plt.ylabel("x2")
-plt.legend()
-plt.grid()
+
+for i in range(0, len(t) - 1):
+    u[:, i:i + 1] = -1 / R * B.T @ P[0:2, 2 * i:2 * i + 2] @ X[0:2, i:i + 1]
+    X[:, i + 1:i + 2] = X[0:2, i:i + 1] + dt * (A @ X[0:2, i:i + 1] + B @ u[0:1, i:i + 1])
+
+
+# two plot
+fig, (ax1, ax2) = plt.subplots(1, 2)
+
+# plot riccati
+ax1.plot(X[0, :], X[1, :], label = 'riccati')
+ax1.scatter(X[0, 0], X[1, 0], label = 'start point')
+ax1.scatter(X[0, len(t) - 1], X[1, len(t) - 1], label = 'final point')
+ax1.legend()
+ax1.grid()
+
+# plot TPBVP
+ax2.plot(X2[0, :], X2[1, :], label='TPBVP')
+ax2.scatter(X2[0, 0], X2[1, 0], label='start point')
+ax2.scatter(X2[0, len(t) - 1], X2[1, len(t) - 1], label='final point')
+ax2.legend()
+ax2.grid()
+
+# show fig
 plt.show()
diff --git a/AHhw1_3.py b/AHhw1_3.py
index a172a06..67868ba 100644
--- a/AHhw1_3.py
+++ b/AHhw1_3.py
@@ -1,215 +1,283 @@
 import numpy as np
-import matplotlib.pyplot as plt
 from math import pi, cos, sin, pow
-
-# define system
-def f(x, u):
-    dx_dt = np.zeros((3, 1))
-    dx_dt[0] = cos(x[2]) * u[0]
-    dx_dt[1] = sin(x[2]) * u[0]
-    dx_dt[2] = u[1]
-    return dx_dt
-
-
-# simulation
-def sim(u, t, X_initial):
-    X = np.zeros((3, len(t)))
-    X[:, 0:1] = X_initial
-    dt = t[1]
-    for i in range(0, len(t)-1):
-        X[:, i+1:i+2] = X[:, i:i+1] + dt * f(X[:, i:i+1], u[:, i])
-    return X
-
-
-# define Riccati equation
-def ricatti(P, A, B, a, b, P_know, r_know):
-    return -P @ A - A.T @ P + P @ B @ np.linalg.inv(R) @ B.T @ P - Q
-
-
-# define r equation
-def rdot(r, A, B, a, b, P_know, r_know):
-    return -(A - B @ np.linalg.inv(R) @ B.T @ P_know).T @ r - a + P_know @ B @ np.linalg.inv(R) @ b
+import matplotlib.pyplot as plt
 
 
-# define z equation
-def zdot(z, A, B, a, b, P_know, r_know):
-    return A @ z + B @ (-np.linalg.inv(R) @ B.T @ P_know @ z - np.linalg.inv(R) @ B.T @ r_know - np.linalg.inv(R) @ b)
+def integrate(f, xt, dt, tt):
+    """
+    This function takes in an initial condition x(t) and a timestep dt,
+    as well as a dynamical system f(x) that outputs a vector of the
+    same dimension as x(t). It outputs a vector x(t+dt) at the future
+    time step.
+
+    Parameters
+    ============
+    dyn: Python function
+        derivate of the system at a given step x(t),
+        it can considered as \dot{x}(t) = func(x(t))
+    xt: NumPy array
+        current step x(t)
+    dt:
+        step size for integration
+    tt:
+        current time
+
+    Return
+    ============
+    new_xt:
+        value of x(t+dt) integrated from x(t)
+    """
+    k1 = dt * f(xt, tt)
+    k2 = dt * f(xt + k1 / 2., tt + dt / 2.)
+    k3 = dt * f(xt + k2 / 2., tt + dt / 2.)
+    k4 = dt * f(xt + k3, tt + dt)
+    new_xt = xt + (1 / 6.) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
+    return new_xt
+
+
+def simulate(f, x0, tspan, dt, integrate):
+    """
+    This function takes in an initial condition x0, a timestep dt,
+    a time span tspan consisting of a list [min_time, max_time],
+    as well as a dynamical system f(x) that outputs a vector of the
+    same dimension as x0. It outputs a full trajectory simulated
+    over the time span of dimensions (xvec_size, time_vec_size).
+
+    Parameters
+    ============
+    f: Python function
+        derivate of the system at a given step x(t),
+        it can considered as \dot{x}(t) = func(x(t))
+    x0: NumPy array
+        initial conditions
+    tspan: Python list
+        tspan = [min_time, max_time], it defines the start and end
+        time of simulation
+    dt:
+        time step for numerical integration
+    integrate: Python function
+        numerical integration method used in this simulation
+
+    Return
+    ============
+    x_traj:
+        simulated trajectory of x(t) from t=0 to tf
+    """
+    N = int((max(tspan) - min(tspan)) / dt)
+    x = np.copy(x0)
+    tvec = np.linspace(min(tspan), max(tspan), N)
+    xtraj = np.zeros((len(x0), N))
+    for i in range(N):
+        xtraj[:, i] = integrate(f, [*x], dt, tvec[i])
+        x = np.copy(xtraj[:, i])
+    return xtraj
 
 
-# define v equation
-def v1(z, A, B, a, b, P_know, r_know):
-    return -np.linalg.inv(R) @ B.T @ P_know @ z - np.linalg.inv(R) @ B.T @ r_know - np.linalg.inv(R) @ b
+# define system
+def f(xu, tt):
+    try:
+        dx_dt = np.array([cos(xu[2]) * u[0, int(tt/dt)], sin(xu[2]) * u[0, int(tt/dt)], u[1, int(tt/dt)]])
+    except:
+        dx_dt = np.array([cos(xu[2]) * u[0, len(t)-1], sin(xu[2]) * u[0, len(t)-1], u[1, len(t)-1]])
+    return dx_dt
 
 
-# find cost function f
-def J(x, desire, input):
-    cost = 0
-    for i in range(0, len(t)):
-        cost = cost + 1/2*(x[:, i] - desire[:, i]).T @ Q @ (x[:, i] - desire[:, i]) + 1/2*input[:, i].T @ R @ input[:, i]
-    cost = cost + 1/2*(x[:, len(t)-1] - desire[:, len(t)-1]).T @ P1 @ (x[:, len(t)-1] - desire[:, len(t)-1])
+# define cost function l
+def l(J, tt):
+    try:
+        cost = (kxi[:, int(tt/dt)]-desire[:, int(tt/dt)]).T @ Q @ (kxi[:, int(tt/dt)] - desire[:, int(tt/dt)]) + u[:, int(tt/dt)].T @ R @ u[:, int(tt/dt)]
+    except:
+        cost = 0
     return cost
 
 
-# find a(t) and b(t)
-def Dl(x, desire, input):
-    a = np.zeros((3, len(t)))
-    b = np.zeros((2, len(t)))
-    for i in range(0, len(t)):
-        a[:, i] = Q.T @ (x[:, i] - desire[:, i])
-        b[:, i] = R.T @ input[:, i]
-    return a, b
+# define A matrix
+def A(tt):
+    try:
+        return np.array([[0, 0, -sin(kxi[2, int(tt/dt)]) * u[0, int(tt/dt)]], [0, 0, cos(kxi[2, int(tt/dt)]) * u[0, int(tt/dt)]], [0, 0, 0]])
+    except:
+        return np.array([[0, 0, -sin(kxi[2, len(kxi[0])-1]) * u[0, len(kxi[0])-1]], [0, 0, cos(kxi[2, len(kxi[0])-1]) * u[0, len(kxi[0])-1]], [0, 0, 0]])
 
 
-# define zeta norm
-def zetanorm(zeta):
-    norm = 0
-    for i in range(0, len(t)):
-        norm = norm + np.linalg.norm(zeta[:, i])
-    return norm
+# define B matrix
+def B(tt):
+    try:
+        return np.array([[cos(kxi[2, int(tt/dt)]), 0], [sin(kxi[2, int(tt/dt)]), 0], [0, 1]])
+    except:
+        return np.array([[cos(kxi[2, len(kxi[0])-1]), 0], [sin(kxi[2, len(kxi[0])-1]), 0], [0, 1]])
 
 
-# define A matrix
-def Amatrix(zeta, u):
-    A = np.zeros((3, 3 * len(t)))
-    for i in range(0, len(t)):
-        A[:, 3 * i:3 * i + 3] = np.array(
-            [[0, 0, cos(zeta[2, i]) * u[0, i]], [0, 0, sin(zeta[2, i]) * u[0, i]], [0, 0, 0]])
-    return A
+def a(tt):
+    try:
+        return (kxi[:, int(tt/dt)]-desire[:, int(tt/dt)]).T @ Q
+    except:
+        return (kxi[:, len(kxi[0])-1]-desire[:, len(kxi[0])-1]).T @ Q
 
 
-# define B matrix
-def Bmatrix(zeta, u):
-    B = np.zeros((3, 2 * len(t)))
-    for i in range(0, len(t)):
-        B[:, 2 * i:2 * i + 2] = np.array([[cos(zeta[2, i]), 0], [sin(zeta[2, i]), 0], [0, 1]])
-    return B
-
-
-# Finds value of y for a given x using step size dt
-# and initial value y0 at x0.
-def rk4(sol_final, t, dt, f, A, B, a, b, P, r):
-    # Iterate for number of iterations
-    row = len(sol_final)
-    columns = len(sol_final[0])
-    sol = np.zeros((row, columns * len(t)))
-    sol[:, (columns * len(t) - columns):(columns * len(t))] = sol_final
-    for i in range(len(t) - 1, 0, -1):
-        k1 = dt * f(sol[:, columns * i:(columns * i + columns)], A[:, 3 * i:3 * i + 3], B[:, 2 * i:2 * i + 2],
-                    a[:, i:i + 1], b[:, i:i + 1], P[:, 3 * i:3 * i + 3], r[:, i:i + 1])
-        k2 = dt * f(sol[:, columns * i:(columns * i + columns)] - k1 / 2., A[:, 3 * i:3 * i + 3], B[:, 2 * i:2 * i + 2],
-                    a[:, i:i + 1], b[:, i:i + 1], P[:, 3 * i:3 * i + 3], r[:, i:i + 1])
-        k3 = dt * f(sol[:, columns * i:(columns * i + columns)] - k2 / 2., A[:, 3 * i:3 * i + 3], B[:, 2 * i:2 * i + 2],
-                    a[:, i:i + 1], b[:, i:i + 1], P[:, 3 * i:3 * i + 3], r[:, i:i + 1])
-        k4 = dt * f(sol[:, columns * i:(columns * i + columns)] - k3, A[:, 3 * i:3 * i + 3], B[:, 2 * i:2 * i + 2],
-                    a[:, i:i + 1], b[:, i:i + 1], P[:, 3 * i:3 * i + 3], r[:, i:i + 1])
-        # Update next value of sol
-        sol[:, (columns * i - columns):columns * i] = sol[:, columns * i:(columns * i + columns)] - (1 / 6.) * (
-                    k1 + 2.0 * k2 + 2.0 * k3 + k4)
-    return sol
+def b(tt):
+    try:
+        return u[:, int(tt/dt)].T @ R
+    except:
+        return u[:, len(kxi[0])-1].T @ R
 
 
-# time step
-t = np.arange(0, 2 * pi, 1e-3)
+# define Riccati equation
+def riccati(P, tt):
+    return -P @ A(tt) - A(tt).T @ P + P @ B(tt) @ np.linalg.inv(R) @ B(tt).T @ P - Q
 
-# initial condition
-X_initial = np.array([[0], [0], [pi / 2]])
 
-# desire trajectory
-x_d = 2 / pi * t
-y_d = np.zeros(t.shape)
-theta_d = pi / 2 * np.ones(t.shape)
-desire = np.vstack((x_d, y_d, theta_d))
+# define r equation
+def rdot(r, tt):
+    return -(A(tt) - B(tt) @ np.linalg.inv(R) @ B(tt).T @ P[:, 3*int(tt/dt):3*int(tt/dt)+3]).T @ r - a(tt).T + P[:, 3*int(tt/dt):3*int(tt/dt)+3] @ B(tt) @ np.linalg.inv(R) @ b(tt).T
 
-# initial input
-u = np.vstack((1 * np.ones(t.shape), -1/2 * np.ones(t.shape)))
-u_old = u
 
-# initial simulation
-x_0 = sim(u, t, X_initial)
-x_old = x_0
+# define z equation
+def zdot(z, tt):
+    return A(tt) @ z + B(tt) @ v1(z, tt)
 
-# initial kxi
-kxi = np.append(x_0, u, axis=0)
 
-# Define LQR variable
-Q = np.array([[2, 0, 0], [0, 0.01, 0], [0, 0, 0.01]])
-R = np.array([[0.1, 0], [0, 0.1]])
-P1 = np.array([[1, 0, 0], [0, 0.01, 0], [0, 0, 0.01]])
+# define v equation
+def v1(z, tt):
+    try:
+        return -np.linalg.inv(R) @ B(tt).T @ P[:, 3*int(tt/dt):3*int(tt/dt)+3] @ z - np.linalg.inv(R) @ B(tt).T @ r[:, int(tt/dt)] - np.linalg.inv(R) @ b(tt).T
+    except:
+        return -np.linalg.inv(R) @ B(tt).T @ P[:, 3 * (len(kxi[0])-1):3 * (len(kxi[0])-1) + 3] @ z - np.linalg.inv(R) @ B(tt).T @ r[:, len(kxi[0])-1] - np.linalg.inv(R) @ b(tt).T
 
-# Define cost function
-cost = J(x_0, desire, u)
+# define directional diff.
+def DiffJdotzeta(Dz, tt):
+    try:
+        return a(tt) @ z[:, int(tt/dt)] + b(tt) @ v[:, int(tt/dt)]
+    except:
+        return a(tt) @ z[:, len(kxi[0])-1] + b(tt) @ v[:, len(kxi[0])-1]
 
 
-def iLQR(x_i, desire, u_i):
+# define zeta norm
+def zetanorm(zn, tt):
+    try:
+        return 0.5*z[:, int(tt/dt)].T @ Q @ z[:, int(tt/dt)] + 0.5*v[:, int(tt/dt)].T @ R @ v[:, int(tt/dt)]
+    except:
+        return 0
 
-    # find a(t),b(t)
-    a, b = Dl(x_i, desire, u_i)
 
-    # Define system variable
-    P_final = P1
-    r_final = (x_i[:, len(t.shape) - 1:len(t.shape)] - desire[:, len(t.shape) - 1:len(t.shape)]).T @ P1
-    dt = t[1]
-    A = Amatrix(x_i, u_i)
-    B = Bmatrix(x_i, u_i)
-    P = np.zeros((3, 3 * len(t)))
-    r = np.zeros((3, 1 * len(t)))
-    z_final = np.zeros((3, 1))
+# time step
+t = np.arange(0, 2 * pi, 1e-3)
+dt = t[1]
+tspan = np.array([t[0], t[len(t)-1]])
 
-    # find P
-    P = rk4(P_final, t, dt, ricatti, A, B, a, b, P, r)
+# initial value
+kxi_0 = np.array([0, 0, pi / 2])
+Q = np.array([[2, 0, 0], [0, 2, 0], [0, 0, 0.01]])
+R = np.array([[0.01, 0], [0, 0.01]])
+P1 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 0]])
 
-    # find r
-    r = rk4(r_final.T, t, dt, rdot, A, B, a, b, P, r)
+# input
+u = np.vstack((1 * np.ones(len(t)), -1/2 * np.ones(len(t))))
+u_initial = u
 
-    # find z
-    z = rk4(z_final, t, dt, zdot, A, B, a, b, P, r)
 
-    # find v
-    v = np.zeros((2, len(t)))
-    for i in range(0, len(t)):
-        v[:, i:i + 1] = v1(z[:, i: i + 1], A[:, 3 * i: 3 * i + 3], B[:, 2 * i:2 * i + 2], a[:, i:i + 1], b[:, i:i + 1],
-                           P[:, 3 * i:3 * i + 3], r[:, i: i + 1])
+# initial simulation
+kxi = simulate(f, kxi_0, tspan, dt, integrate)
+kxi = np.append(kxi_0.reshape((3, 1)), kxi, axis=1)
 
-    # find gradient descent
-    zeta = np.append(z, v, axis=0)
+# store initial simulation
+kxi_initial = kxi
 
-    # find DJ(kxi) dot zeta
-    DJdotzeta = 0
-    for i in range(0, len(t)):
-        DJdotzeta = DJdotzeta + np.dot(a[:, i], z[:, i]) + np.dot(b[:, i], v[:, i])
+# desire trajectory
+x_d = 2 / pi * t
+y_d = np.zeros(t.shape)
+theta_d = pi / 2 * np.ones(t.shape)
+desire = np.vstack((x_d, y_d, theta_d))
 
-    # initial  Armijo line search coefficient
-    alpha = 0.3
-    beta = 0.5
-    N = 1
+for j in range(0, 5):
+
+    # calculate cost function
+    J = 0
+    for i in range(0, len(kxi[0])):
+        J = integrate(l, J, dt, t[i])
+    print('J')
+    print(J)
+
+    # solve P
+    P = np.zeros((3, 3 * len(kxi[0])))
+    P[:, 3*len(kxi[0])-3:3*len(kxi[0])] = P1
+    for i in range(len(kxi[0])-1, 0, -1):
+        P[:, 3*i-3: 3*i] = integrate(riccati, P[:, 3*i: 3*i+3], -dt, t[i])
+
+    # solve r
+    r = np.zeros((3, len(kxi[0])))
+    r[:, len(kxi[0])-1] = P1.T@(kxi[:, len(kxi[0])-1]-desire[:, len(kxi[0])-1])
+    for i in range(len(kxi[0])-1, 0, -1):
+        r[:, i-1] = integrate(rdot, r[:, i], -dt, t[i])
+
+    # simulation z
+    z = np.zeros((3, len(kxi[0])))
+    z_0 = np.array([0, 0, 0])
+    z = simulate(zdot, z_0, tspan, dt, integrate)
+    z = np.append(z_0.reshape((3, 1)), z, axis=1)
+
+    # solve v
+    v = np.zeros((2, len(kxi[0])))
+    for i in range(0, len(kxi[0])):
+        v[:, i] = v1(z[:, i], t[i])
+
+    # DJ(kxi) dot zeta
+    DJdotz = 0
+    for i in range(0, len(kxi[0])):
+        DJdotz = integrate(DiffJdotzeta, DJdotz, dt, t[i])
+    '''print('DJdotz')
+    print(DJdotz)'''
+    # initial Armijo line search coefficient
+    alpha = 0.45
+    beta = 0.2
+    N = 0
+
+    # store current input
+    u_old = u
 
     # Armijo line search
     while True:
         gamma = pow(beta, N)
-        u_new = u_i + gamma * v
-        x_new = sim(u_new, t, X_initial)
+        u = u + gamma * v
+        kxi = simulate(f, kxi_0, tspan, dt, integrate)
+        kxi = np.append(kxi_0.reshape((3, 1)), kxi, axis=1)
+        J_new = 0
+        for i in range(0, len(kxi[0])):
+            J_new = integrate(l, J_new, dt, t[i])
+        '''print('J_new')
+        print(J_new)'''
+
         # compare current cost function with linear
-        if J(x_new, desire, u_new) <= J(x_0, desire, u)+alpha*gamma*DJdotzeta:
+        if J_new <= J + alpha*gamma*DJdotz:
             break
+
         N = N + 1
 
-    print(zetanorm(zeta))
-    return x_new, u_new, z, v
+        # reset input
+        u = u_old
 
+    # calculate zeta norm
+    zetaN = 0
+    for i in range(0, len(kxi[0])):
+        zetaN = integrate(zetanorm, zetaN, dt, t[i])
+    print('zetaN')
+    print(zetaN)
+
+    print(str(j+1) + 'loop complete')
 
-# iterative
-for i in range(0, 5):
-    x_new, u_new, z, v = iLQR(x_old, desire, u_old)
-    x_old = x_new
-    u_old = u_new
 
 # plot initial guess and reference
-plt.plot(x_0[0, :], x_0[1, :], label = 'initial')
-plt.plot(x_new[0, :], x_new[1, :], label = 'first iteration')
-plt.xlabel("x")
-plt.ylabel("y")
-plt.legend()
+fig, (ax1, ax2) = plt.subplots(2)
+ax1.plot(kxi_initial[0, :], kxi_initial[1, :], label='initial')
+ax1.plot(kxi[0, :], kxi[1, :], label='iterative')
+ax1.plot(desire[0, :], desire[1, :], label='desire')
+ax1.legend()
+
+ax2.plot(t, u_initial[0, :], label='initial input 1')
+ax2.plot(t, u[0, :], label='new input 1')
+ax2.plot(t, u_initial[1, :], label='initial input 2')
+ax2.plot(t, u[1, :], label='new input 2')
+ax2.legend()
+
+plt.grid
 plt.show()
 
-
-
diff --git a/AHhw1_5.py b/AHhw1_5.py
new file mode 100644
index 0000000..34057b7
--- /dev/null
+++ b/AHhw1_5.py
@@ -0,0 +1,306 @@
+import numpy as np
+from math import pi, cos, sin, pow
+import matplotlib.pyplot as plt
+
+
+def integrate(f, xt, dt, tt):
+    """
+    This function takes in an initial condition x(t) and a timestep dt,
+    as well as a dynamical system f(x) that outputs a vector of the
+    same dimension as x(t). It outputs a vector x(t+dt) at the future
+    time step.
+
+    Parameters
+    ============
+    dyn: Python function
+        derivate of the system at a given step x(t),
+        it can considered as \dot{x}(t) = func(x(t))
+    xt: NumPy array
+        current step x(t)
+    dt:
+        step size for integration
+    tt:
+        current time
+
+    Return
+    ============
+    new_xt:
+        value of x(t+dt) integrated from x(t)
+    """
+    # print(tt+dt)
+    k1 = dt * f(xt, tt)
+    k2 = dt * f(xt + k1 / 2., tt + dt / 2.)
+    k3 = dt * f(xt + k2 / 2., tt + dt / 2.)
+    k4 = dt * f(xt + k3, tt + dt)
+    new_xt = xt + (1 / 6.) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
+    return new_xt
+
+
+def simulate(f, x0, tspan, dt, integrate):
+    """
+    This function takes in an initial condition x0, a timestep dt,
+    a time span tspan consisting of a list [min_time, max_time],
+    as well as a dynamical system f(x) that outputs a vector of the
+    same dimension as x0. It outputs a full trajectory simulated
+    over the time span of dimensions (xvec_size, time_vec_size).
+
+    Parameters
+    ============
+    f: Python function
+        derivate of the system at a given step x(t),
+        it can considered as \dot{x}(t) = func(x(t))
+    x0: NumPy array
+        initial conditions
+    tspan: Python list
+        tspan = [min_time, max_time], it defines the start and end
+        time of simulation
+    dt:
+        time step for numerical integration
+    integrate: Python function
+        numerical integration method used in this simulation
+
+    Return
+    ============
+    x_traj:
+        simulated trajectory of x(t) from t=0 to tf
+    """
+    N = int((max(tspan) - min(tspan)) / dt)
+    x = np.copy(x0)
+    tvec = np.linspace(min(tspan), max(tspan), N)
+    xtraj = np.zeros((len(x0), N))
+    for i in range(N):
+        xtraj[:, i] = integrate(f, [*x], dt, tvec[i])
+        x = np.copy(xtraj[:, i])
+    return xtraj
+
+
+# define system
+def f1(x, tt):
+    dx_dt = np.array([cos(x[2]) * u1[0, int(tt / dt)], sin(x[2]) * u1[0, int(tt / dt)], u1[1, int(tt / dt)]])
+    return dx_dt
+
+
+# define Dl1(tt)
+def Dl1(tt):
+    return (x[:, int(tt / dt)] - desire[:, int(tt / dt)]).T @ Q
+
+
+# define A matrix
+def A(tt):
+    return np.array([[0, 0, -sin(x[2, int(tt / dt)]) * u1[0, int(tt / dt)]],
+                     [0, 0, cos(x[2, int(tt / dt)]) * u1[0, int(tt / dt)]], [0, 0, 0]])
+
+
+# define h matrix
+def h(tt):
+    return np.array([[cos(x[2, int(tt/dt)]), 0], [sin(x[2, int(tt/dt)]), 0], [0, 1]])
+
+
+# define rhodot
+def rhodot(rho, tt):
+    return -A(tt).T @ rho - Dl1(tt)
+
+
+# define cost function
+def l(J1, tt):
+    cost = (x[:, int(tt / dt)] - desire[:, int(tt / dt)]).T @ Q @ \
+           (x[:, int(tt / dt)] - desire[:, int(tt / dt)]) + u1[:, int(tt / dt)].T @ R @ u1[:, int(tt / dt)]
+    return cost
+
+
+# define lambda
+def Lambda(tt):
+    return h(tt).T @ rho[:, int(tt / dt):(int(tt / dt)+1)] @ rho[:, int(tt / dt):(int(tt / dt)+1)].T @ h(tt)
+
+
+# define u2*
+def u2star(tt):
+    return np.linalg.inv(Lambda(tt) + R.T) @ (Lambda(tt) @ u1[:, int(tt / dt)] + h(tt).T @ rho[:, int(tt / dt)] * alpha_d)
+
+
+# calculate dJ1_dlambda
+def dJ1_dlambda(tt):
+    return rho[:, int(tt / dt)].T @ (h(tt) @ u2star(tt) - h(tt) @ u1[:, int(tt / dt)])
+
+
+# second cost function J_tao
+def J_tao(tt):
+    return np.linalg.norm(u2star(tt)) + dJ1_dlambda(tt) + pow(tt-t_0, beta)
+
+
+# define all constant
+
+J_min = 0.0001  # minimum change in cost
+t_curr = 0  # current time
+t_init = 0.1  # default control duration
+omega = 0.5  # scale factor
+T = 2.  # predictive horizon
+ts = 0.01  # sampling time
+t_calc = 0.  # the max time for iterative control calculation
+k_max = 10  # the max backtracking iterations
+gamma = -9.  # the coefficient for first order sensitivity alpha
+beta = 0.7
+Q = np.array([[2, 0, 0], [0, 2, 0], [0, 0, 0.01]])  # Q in cost function
+R = np.array([[0.01, 0], [0, 0.01]])  # R in cost function
+
+# the end time t_end
+t_end = 2*pi
+
+# define x_init
+x_init = np.array([0, 0, pi / 2])
+
+while t_curr < t_end:
+
+    # define t_0 and t_final
+    t_0 = 0.
+    t_final = T
+    t = np.arange(t_0, t_final + ts, ts)
+    tspan = np.array([t_0, t_final])
+    dt = ts
+
+    # define td for desire trajectory
+    td = np.array([i+t_curr for i in t])
+
+    # desire trajectory
+    x_d = 2 / pi * td
+    y_d = np.zeros(td.shape)
+    theta_d = pi / 2 * np.ones(td.shape)
+    desire = np.vstack((x_d, y_d, theta_d))
+    #print('desire.shape')
+    #print(desire.shape)
+
+    # define nominal input
+    if t_curr == 0:
+        u1 = np.vstack((0.1 * np.ones(len(t)), -0.05 * np.ones(len(t))))
+    else:
+        u1_new = np.array([[1], [-0.5]])
+        u1 = np.append(u1[:, 1:len(t)], u1_new, axis=1)
+    #print('u1.shape')
+    #print(u1.shape)
+
+    # solve x
+    x = simulate(f1, x_init, tspan, ts, integrate)
+    x = np.append(x_init.reshape((3, 1)), x, axis=1)
+    #print('x.shape')
+    #print(x.shape)
+
+    # solve rho
+    rho = np.zeros((3, len(t)))
+    rho[:, len(t) - 1] = np.zeros(3)
+    for i in range(len(t) - 1, 0, -1):
+        rho[:, i - 1] = integrate(rhodot, rho[:, i], -dt, t[i])
+    #print('rho.shape')
+    #print(rho.shape)
+
+    # calculate cost function
+    J_init = 0
+    for i in range(0, len(t)):
+        J_init = integrate(l, J_init, dt, t[i])
+    print('J')
+    print(J_init)
+
+    # calculate sensitivity alpha
+    alpha_d = gamma * J_init
+
+    # find tao when u2star(tao) is min
+    tao = t_0+t_calc
+    minJtao = J_tao(t_0+t_calc)
+    for i in range(0, len(x[0])):
+        if t[i] > t_0+t_calc:
+            if J_tao(t[i]) < minJtao:
+                minJtao = J_tao(t[i])
+                tao = t[i]
+    print('tao')
+    print(tao)
+
+    # initial k and J1_new
+    k = 0
+    J_new = 10e8
+
+    # store old u1
+    u1_old = u1
+
+    # search for lambda
+    while J_new - J_init > J_min and k <= k_max:
+
+        # restore the data
+        u1 = u1_old
+
+        # calculate lambda, tao_0 & tao_f
+        lambda1 = (omega**k)*t_init
+        tao_0 = tao - lambda1/2
+        if tao_0 < 0:
+            tao_0 = 0
+        tao_f = tao + lambda1/2
+        if tao_f > t_final:
+            tao_f = t_final
+        #print('lambda1, tao_0, tao_f')
+        #print(lambda1, tao_0, tao_f)
+
+        # add u2star to u1 and saturate
+        for i in range(int(tao_0/dt), int(tao_f/dt)):
+
+            # saturate u2star
+            u2_new = u2star(i*dt)
+            if u2_new[0] > 5:
+                u1[0, i] = 5
+            elif u2_new[0] < -5:
+                u1[0, i] = -5
+            else:
+                u1[0, i] = u2_new[0]
+
+            if u2_new[1] > 5:
+                u1[1, i] = 5
+            elif u2_new[1] < -5:
+                u1[1, i] = -5
+            else:
+                u1[1, i] = u2_new[1]
+
+        # re_simulation
+        x = simulate(f1, x_init, tspan, dt, integrate)
+        x = np.append(x_init.reshape((3, 1)), x, axis=1)
+
+        # calculate new cost
+        J_new = 0
+        for i in range(0, len(x[0])):
+            J_new = integrate(l, J_new, dt, t[i])
+        k = k+1
+        #print('J_new')
+        #print(J_new)
+
+    # store desire trajectory and real trajectory
+    if t_curr == 0:
+        X_d = desire[:, 0:1]
+        X = x[:, 0:1]
+        U = u1[:, 0:1]
+    else:
+        X_d = np.append(X_d, desire[:, 0:1], axis=1)
+        X = np.append(X, x[:, 0:1], axis=1)
+        U = np.append(U, u1[:, 0:1], axis=1)
+
+    # update time and x_init
+    t_curr = t_curr + ts
+    x_init = x[:, int(ts/ts)]
+
+    print('current time ' + str(t_curr))
+
+# simulation with update input
+x_init1 = np.array([0, 0, pi / 2])
+u1 = U
+tspan1 = np.array([0, int(2*pi)])
+X1 = simulate(f1, x_init1, tspan1, dt, integrate)
+X1 = np.append(x_init1.reshape((3, 1)), X1, axis=1)
+
+# plot initial
+fig, (ax1, ax2, ax3) = plt.subplots(3)
+ax1.plot(X1[0, :], X1[1, :], label='update')
+ax1.plot(X_d[0, :], X_d[1, :], label='desire')
+ax1.legend()
+
+ax2.plot(X[0, :], X[1, :])
+
+ax3.plot(U[0, :], label='u1')
+ax3.plot(U[1, :], label='u2')
+ax3.legend()
+
+plt.show()