import sympy as sm
import numpy as np

sm.init_printing()

s = sm.symbols("s")
z = sm.symbols("z")
T = sm.symbols("T")

#########################################################
print("PROBLEM 2:")
print("Part a:")

ZOH = (1 - sm.exp(-s * T)) / (s * T)

G_1_s = 1 / s * ZOH

bilinear_s = 2 / T * (z - 1) / (z + 1)

G_1_k = G_1_s.subs({s: bilinear_s}).simplify()

print("G_1_k = ")
sm.pprint(G_1_k)

print("Part b:")

theta_s_r_s = ZOH * 1 / s**2
theta_k_r_k = theta_s_r_s.subs({s: bilinear_s}).simplify()
print("theta_k_r_k = ")
sm.pprint(theta_k_r_k)


#########################################################
print("PROBLEM 3:")
A = sm.Matrix([[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]])
B = sm.Matrix([0, 0, 0, 1])
C = sm.Matrix([1, 0, 0, 0]).transpose()
D = sm.Matrix([1])

print("System Matricies A, B, C, D")
sm.pprint(A)
sm.pprint(B)
sm.pprint(C)
sm.pprint(D)

# Recursive Solution
k = sm.symbols("k", integer=True, real=True, positive=True)


def y(k, u, x_0):
    term_1 = C * A**k * x_0

    term_2 = sm.Matrix([0, 0, 0, 0])
    for j in range(np.size(u)):
        term_2 = term_2 + A ** (k - j - 1) * B * u[j]
    term_2 = C * term_2

    term_3 = D * u[-1]

    return term_1 + term_2 + term_3


print("Part c:")
x_0 = sm.Matrix([2, 1, 3, 0])
u = sm.Matrix([0])

output = y(k, u, x_0)
output = output[0].expand().simplify()
print("y(k) = ")
sm.pprint(output)


print("Part d:")
x_0 = sm.Matrix([0, 0, 0, 0])
u = sm.Matrix([2, 1, 3, 0])
output = y(k, u, x_0)
output = output[0].expand().simplify()
print("y(k) = ")
sm.pprint(output)

print("Part e:")
print(
    "These are the same exact response between parts C and D. This makes sense, becasue we defined our states as just being delays in a chain. The result is that the input at timestep k trickles down through each state in k+1, k+2, and k+3. This means that our states save our input in a way, s.t. loading this initial state mathematically produces an identical result as loading those inputs in one time step at a time. \n \n This is reflected by the two algebraeic expressions being the same."
)

#########################################################
print("PROBLEM 4:")
print("Part a:")

"""
x(k+2) - x(k+1) + 0.25 x(k) = u(k+2)

Applying Z transform:
z^2 X - z X + 0.25 X = z^2 U
X (z^2 - z + 0.25) = z^2 U
X/U = z^2 / (z^2 - z + 0.25)
"""
# Use SymPy to do partial frac
z = sm.symbols("z")

X_U = z**2 / (z**2 - z + 0.25)
sm.pprint(X_U.apart())

print("Part b:")

x_z = (z**3 / (z - 1) - z + 2 * z + z**2 - z) / (z**2 - z + 0.25)
sm.pprint(x_z.apart())

#######################################################
print("PROBLEM 5:")
print("part a:")

G_z = (1 - z**-1) / 5 * (1 / (1 - z**-1) - sm.exp(T) * z / (z - sm.exp(-5 * T)))

sm.pprint(G_z)
sm.pprint(G_z.simplify())
print(sm.latex(G_z.simplify()))

print("part b:")
G_z = G_z.simplify()
alp_0, alp_1, beta_1 = sm.symbols("alpha_0, alpha_1, beta_1")
D_z = (alp_0 + alp_1 * z**-1) / (1 + beta_1 * z**-1)

# Calculate Loop Gain
L_z = G_z * D_z
H_z = 1

C_z_over_R_z = L_z / (1 + L_z * H_z)
sm.pprint(C_z_over_R_z.simplify())
print(sm.latex((C_z_over_R_z.expand().simplify())))

print("part c:")
R_z = 1 / (1 - z**-1)

C_z = C_z_over_R_z * R_z

c_inf = (z - 1) * C_z
sm.pprint(c_inf.simplify())
print(sm.latex(c_inf.simplify()))