import numpy as np
import matplotlib.pyplot as plt
a2 = -3.35
a1 = 3.15
a0 = -0.9


x = []
u = 1

def odpSkok(x0, u, A, B, C, D):
    x = []
    y = []
    for i in range(20):
        xk = A*x0 + B*u
        yk = C*xk + D*u
        x0 = xk

        #x.append(xk)
        y.append(yk.item(0))
    return y



x0 = np.matrix([0,0,0])
x0 = np.transpose(x0)
Aa = np.matrix([[-a2, -a1, -a0], [1, 0, 0], [0, 1, 1]])
Ba = np.transpose(np.matrix([1, 0, 0]))
Ca = np.matrix([1, 1, 1])
Da = [0]
#x0 = np.zeros(3)
y = odpSkok(x0, 1, Aa, Ba, Ca, Da)

def Riccatiego(A, B, Q, R):
    Pp = 0*np.identity(3)
    P = Pp
    for i in range(10):
        P = Q + np.transpose(A)*(Pp - Pp*B*(R + np.transpose(B)*P*B)**(-1)*
                                 np.transpose(B)*P)*A
        Pp = P

    return P


def sterownik(c1, c2, A, B):
    Q = c1*np.identity(3)
    R = c2
    P = Riccatiego(A, B, Q, R)

    F = (R + np.transpose(B)*P*B)**(-1)*np.transpose(B)*P*A
    return F
    pass

plt.figure(1)

plt.subplot(221)
plt.plot(y)
plt.title("bez sterownika")
#plt.show()

F = sterownik(1, 1, Aa, Ba)
Anowe = Aa - Ba*F
ynowe = odpSkok(0, 1, Anowe, Ba, Ca, Da)
plt.subplot(223)
plt.plot(ynowe)
plt.title("c1 = 1 c2 = 1")


def uk(x0, u, A, B, C, D, F):
    x = []
    y = []
    uk = []
    print F
    for i in range(20):
        xk = A*x0 + B*u
        #yk = C*xk + D*u
        x0 = xk
        uk_buf = -np.dot(F,xk)


        #x.append(xk)
        #uk.append(uk_buf)
        uk = np.append(uk, uk_buf)
    print xk
    return uk

uk = uk(x0, 1, Anowe, Ba, Ca, Da, F)
plt.subplot(224)
plt.plot(uk)
plt.title("uk")
plt.show()