import numpy as np
import sympy as sp
from sympy.abc import z
from sympy.polys.partfrac import apart, apart_list
from sympy import *
from scipy import signal
import matplotlib.pyplot as plt


trm = ((z**2 + z + 1) / (z**3 - 2.55*z**2 + 1.175*z - 0.15))
wspG = [1, -2.55, 1.175, -0.15]

a = apart(trm,z).evalf()

prw1 = a.args[0]
prw2 = a.args[1]
prw3 = a.args[2]

wsp1 = a.args[0].args[0]
wsp2 = a.args[1].args[0]
wsp3 = a.args[2].args[0]


print('GRAF UKLADU: ')
print('\n')
print('            /','------------> ', prw1, '------------->', '\\')
print('U(wejscie) ---','------------> ', prw2, '------------->', '----> Y(wyjscie)')
print('           \\','------------> ', prw3, '------------->', '/')


polaczone = together(a)
print('\n')
print('polaczone:',polaczone)


#MODEL STANOWY========================

A = np.array([[wspG[1],wspG[2],wspG[3]],[1,0,0],[0,1,0]])
print('MACIERZ A: ')
print(A[0])
print(A[1])
print(A[2])

B = np.array([[1,0,0]]).T
C = np.array([[1,1,1]])
D = 0
print('MACIERZ B: ')
print(B)
print('MACIERZ C: ')
print(C)
print('MACIERZ D: ')
print(D)

#MODEL STANOWY========================

#RYSOWANIE============================




#RYSOWANIE============================

#STEROWNIK LQR========================

At = np.matrix.transpose(A)
print('transpose A:')
print(At)
Bt = np.matrix.transpose(B)
P = np.array([[0,1,1],[1,2,3],[1,3,2]])

Q = np.array([[-3,0,0],[0,-3,0],[0,0,-3]])
R = -50

#P_Q = At*(P - P*B*(1/(R+Bt*P*B))*Bt*P)*A

def podstawianieP(P):

    P_BtP = np.matmul(Bt,P)
    P_BtPB = np.matmul(P_BtP,B)
    P_PB = np.matmul(P,B)
    P_PBR = np.matmul(P_PB,1/(P_BtPB+R))
    P_PA = np.matmul(P_PBR,Bt)
    P_PA2 = np.matmul(P_PA,P)
    P_PA2 = P - P_PA2
    P_PA3 = np.matmul(P_PA2,A)
    P_PA4 = np.matmul(At,P_PA3)
    P = Q + P_PA4
    return P


for i in range(20):
    P = podstawianieP(P)
    print('macierz P:')
    print(P)

P_BtP = np.matmul(Bt,P)
P_BtPB = np.matmul(P_BtP,B)
P_PBR = np.matmul(1/(P_BtPB+R),Bt)
P_PA = np.matmul(P_PBR,P)
F = np.matmul(P_PA,A)

noweA = A - np.matmul(B,F)

#STEROWNIK LQR========================

#RYSOWANIE2===========================

name = 'odp skokowa'
name2 = 'odp z regulatorem LQR'
rts = np.roots(wspG)

p1 = np.poly1d([1,-rts[0]])
p2 = np.poly1d([1,-rts[1]])
p3 = np.poly1d([1,-rts[2]])

num = np.poly1d([1,1,1])
den = p1*p2*p3

sys = signal.TransferFunction(num,den)

#=====================================

nowywsp1 = -noweA[0][0]
nowywsp2 = -noweA[0][1]
nowywsp3 = -noweA[0][2]

nowywspG = [1,nowywsp1,nowywsp2,nowywsp3]

nowyRoots = np.roots(nowywspG)
print(nowyRoots)

p1n = np.poly1d([1,nowyRoots[0]])
p2n = np.poly1d([1,nowyRoots[1]])
p3n = np.poly1d([1,nowyRoots[2]])

numn = np.poly1d([1,1,1])
denn = p1n*p2n*p3n

sysn = signal.TransferFunction(numn,denn)



tn,yn = signal.step(sysn)
t,y = signal.step(sys)
plt.figure(1)
plt.plot(t,y,'k-',label=name)
plt.plot(tn,yn,'b-',label=name2)
plt.axis([0,50,0,10000])
plt.legend(loc='best')
plt.show()

#RYSOWANIE2===========================