# double integrator ( G(s)=1/s^2 )

import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
from scipy import signal
np.set_printoptions(precision=4)

A = np.array([[ 0.0, 1.0 ],
              [ 0.0, 0.0 ]])
B = np.array([[ 0.0 ],
              [ 1.0 ]])
C = np.array([[ 1.0, 0.0 ]])
D = np.zeros((1,1))
m = C.shape[0]
p = D.shape[1]
sysC = signal.lti(A,B,C,D)

# Q = C.T@C = np.diag([1.0, 0.0])

Q = C.T@C
rho_list = [100,10,1,0.1,0.01]
K_list = []
for rho in rho_list:
  R = rho*np.eye(1)
  P = sp.linalg.solve_continuous_are(A,B,Q,R)
  K = -np.linalg.inv(R)@B.T@P
  K_list.append(K)

K_list_ = np.array([x[0]  for x in K_list])

plt.figure(figsize=(8,8), dpi=100)
for i in range(2):
  plt.subplot(2,1,i+1)
  plt.semilogx(rho_list, K_list_[:,i], 'x-', label=f'$K_{i+1}$')
  plt.legend()
  plt.grid()
plt.xlabel(r'$\rho$')
plt.show()
Tf = 10
t, yo = signal.impulse(sysC, T=np.arange(0,Tf,0.01))

plt.figure(figsize=(8,4), dpi=100)
plt.plot(t,yo, label='open loop')
i = 0
for K in K_list:
  sysCc = signal.lti(A+B@K, B, C, D)
  t, yc = signal.impulse(sysCc, T=np.arange(0,Tf,0.01))
  plt.plot(t,yc, label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel(r'$t$')
plt.ylabel(r'$x$')

plt.figure(figsize=(8,8), dpi=100)
plt.plot(0,0, 'x', label='open loop')
i = 0
for K in K_list:
  eval, evec = np.linalg.eig(A+B@K)
  plt.plot(eval.real, eval.imag, 'x', label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel('real')
plt.ylabel('imag')
plt.axis('equal')
plt.show()

# Q = np.diag([1.0, 0.2])

Q = np.diag([1.0, 0.2])
rho_list = [100,10,1,0.1,0.05,0.04,0.03,0.02,0.01,0.005]
K_list = []
for rho in rho_list:
  R = rho*np.eye(1)
  P = sp.linalg.solve_continuous_are(A,B,Q,R)
  K = -np.linalg.inv(R)@B.T@P
  K_list.append(K)

K_list_ = np.array([x[0]  for x in K_list])

plt.figure(figsize=(8,8), dpi=100)
for i in range(2):
  plt.subplot(2,1,i+1)
  plt.semilogx(rho_list, K_list_[:,i], 'x-', label=f'$K_{i+1}$')
  plt.legend()
  plt.grid()
plt.xlabel(r'$\rho$')
plt.show()
Tf = 10
t, yo = signal.impulse(sysC, T=np.arange(0,Tf,0.01))

plt.figure(figsize=(8,4), dpi=100)
plt.plot(t,yo, label='open loop')
i = 0
for K in K_list:
  sysCc = signal.lti(A+B@K, B, C, D)
  t, yc = signal.impulse(sysCc, T=np.arange(0,Tf,0.01))
  plt.plot(t,yc, label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel(r'$t$')
plt.ylabel(r'$x$')

plt.figure(figsize=(8,8), dpi=100)
plt.plot(0,0, 'x', label='open loop')
i = 0
for K in K_list:
  eval, evec = np.linalg.eig(A+B@K)
  plt.plot(eval.real, eval.imag, 'x', label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel('real')
plt.ylabel('imag')
plt.axis('equal')
plt.show()

# Q = np.diag([1.0, 0.5])

Q = np.diag([1.0, 0.5])
rho_list = [100,10,1,0.1,0.05,0.04,0.03,0.02,0.01]
K_list = []
for rho in rho_list:
  R = rho*np.eye(1)
  P = sp.linalg.solve_continuous_are(A,B,Q,R)
  K = -np.linalg.inv(R)@B.T@P
  K_list.append(K)

K_list_ = np.array([x[0]  for x in K_list])

plt.figure(figsize=(8,8), dpi=100)
for i in range(2):
  plt.subplot(2,1,i+1)
  plt.semilogx(rho_list, K_list_[:,i], 'x-', label=f'$K_{i+1}$')
  plt.legend()
  plt.grid()
plt.xlabel(r'$\rho$')
plt.show()
Tf = 10
t, yo = signal.impulse(sysC, T=np.arange(0,Tf,0.01))

plt.figure(figsize=(8,4), dpi=100)
plt.plot(t,yo, label='open loop')
i = 0
for K in K_list:
  sysCc = signal.lti(A+B@K, B, C, D)
  t, yc = signal.impulse(sysCc, T=np.arange(0,Tf,0.01))
  plt.plot(t,yc, label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel(r'$t$')
plt.ylabel(r'$x$')

plt.figure(figsize=(8,8), dpi=100)
plt.plot(0,0, 'x', label='open loop')
i = 0
for K in K_list:
  eval, evec = np.linalg.eig(A+B@K)
  plt.plot(eval.real, eval.imag, 'x', label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel('real')
plt.ylabel('imag')
plt.axis('equal')
plt.show()

# Q = np.diag([1.0, 2.0])

Q = np.diag([1.0, 2.0])
rho_list = [100,10,1,0.1,0.05,0.04,0.03,0.02,0.01]
K_list = []
for rho in rho_list:
  R = rho*np.eye(1)
  P = sp.linalg.solve_continuous_are(A,B,Q,R)
  K = -np.linalg.inv(R)@B.T@P
  K_list.append(K)

K_list_ = np.array([x[0]  for x in K_list])

plt.figure(figsize=(8,8), dpi=100)
for i in range(2):
  plt.subplot(2,1,i+1)
  plt.semilogx(rho_list, K_list_[:,i], 'x-', label=f'$K_{i+1}$')
  plt.legend()
  plt.grid()
plt.xlabel(r'$\rho$')
plt.show()
Tf = 10
t, yo = signal.impulse(sysC, T=np.arange(0,Tf,0.01))

plt.figure(figsize=(8,4), dpi=100)
plt.plot(t,yo, label='open loop')
i = 0
for K in K_list:
  sysCc = signal.lti(A+B@K, B, C, D)
  t, yc = signal.impulse(sysCc, T=np.arange(0,Tf,0.01))
  plt.plot(t,yc, label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel(r'$t$')
plt.ylabel(r'$x$')

plt.figure(figsize=(8,8), dpi=100)
plt.plot(0,0, 'x', label='open loop')
i = 0
for K in K_list:
  eval, evec = np.linalg.eig(A+B@K)
  plt.plot(eval.real, eval.imag, 'x', label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel('real')
plt.ylabel('imag')
plt.axis('equal')
plt.show()

# double integrator with a zero ( G(s) = (2s+1)/s^2 )

import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
from scipy import signal

h = 0.01

A = np.array([[ 0.0, 1.0 ],
              [ 0.0, 0.0 ]])
B = np.array([[ 0.0 ],
              [ 1.0 ]])
C = np.array([[ 1.0, 2.0 ]])
D = np.zeros((1,1))
m = C.shape[0]
p = D.shape[1]

sysC = signal.lti(A,B,C,D)


# Q = C.T@C = np.array([[1.0, 2.0], [2.0, 4.0]])

Q = C.T@C
rho_list = [100,10,5,4,3,2,1]
K_list = []
for rho in rho_list:
  R = rho*np.eye(1)
  P = sp.linalg.solve_continuous_are(A,B,Q,R)
  K = -np.linalg.inv(R)@B.T@P
  K_list.append(K)

K_list_ = np.array([x[0]  for x in K_list])

plt.figure(figsize=(8,8), dpi=100)
for i in range(2):
  plt.subplot(2,1,i+1)
  plt.semilogx(rho_list, K_list_[:,i], 'x-', label=f'$K_{i+1}$')
  plt.legend()
  plt.grid()
plt.xlabel(r'$\rho$')
plt.show()
Tf = 10
t, yo = signal.impulse(sysC, T=np.arange(0,Tf,0.01))

plt.figure(figsize=(8,4), dpi=100)
plt.plot(t,yo, label='open loop')
i = 0
for K in K_list:
  sysCc = signal.lti(A+B@K, B, C, D)
  t, yc = signal.impulse(sysCc, T=np.arange(0,Tf,0.01))
  plt.plot(t,yc, label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel(r'$t$')
plt.ylabel(r'$x$')

plt.figure(figsize=(8,8), dpi=100)
plt.plot(0,0, 'x', label='open loop')
i = 0
for K in K_list:
  eval, evec = np.linalg.eig(A+B@K)
  plt.plot(eval.real, eval.imag, 'x', label=rf'Infinite horizon LQR, $\rho$={rho_list[i]}')
  i = i+1
plt.grid()
plt.legend()
plt.xlabel('real')
plt.ylabel('imag')
plt.axis('equal')
plt.show()