import sympy as sm
import numpy as np
import matplotlib.pyplot as plt

sm.init_printing()

print("PROBLEM 1:")
print("part a:")
s, z, t = sm.symbols("s z t")

J, B, RHO, RADIUS, L, RESIST, T = sm.symbols(
    "J, B, rho,  r,  L,  R, T", Real=True, Positive=True
)
K_T, K_B = sm.symbols("K_T,  K_b")

# loop gain
Loop_Gain = K_T * 1 / (J * s + B)

# transfer function from motor current command to linear displacement (ignoring feedback)
X_over_I = Loop_Gain * RHO / s * RADIUS
sm.pprint(X_over_I)
X_over_I = X_over_I.expand().simplify()
sm.pprint(X_over_I)

print("part b:")
# recall ZOH_eq
# G(z) = (1-z**-1) Z{L**-1{G(s)/s}}
print("Finding inverse laplace of G/s")
G_s = X_over_I
sm.pprint(G_s / s)
g_t = sm.inverse_laplace_transform(G_s / s, s, t)
sm.pprint(g_t.expand())

# make z substitution subs

G_z = (
    K_T * T * RADIUS * RHO / B * (z / (z - 1) ** 2)
    - J * K_T * RADIUS * RHO / B**2
    + J * K_T * RADIUS * RHO / B**2 * (z / (z - sm.exp(-B / J * T)))
)

sm.pprint(G_z.simplify())

# part c was on the board
print("part c")
F = sm.Matrix([[0, RHO], [0, -B / J]])
G = sm.Matrix([0, K_T / J])

# part d
print("part d")
A = sm.exp(F * T)
B = sm.integrate(A, (T, 0, T))
sm.pprint(F)
sm.pprint(G)
sm.pprint(A)
sm.pprint(B)