'''
Calculating poles and simulating impulse response using Python Control Package
Finn Aakre Haugen (finn@techteach.no)
Updated 2025 06 29
'''

# %% Imports:

import numpy as np
import matplotlib.pyplot as plt
import control

# %% Creating the transfer function:

s = control.tf('s')

H1 = 1/(s + 1)
# H2 = 1/s
# H3 = 1/(s - 1)  # Transfer function
H = H1

# %% Calculation and plotting of poles:

poles = control.poles(H)

plt.close('all')
plt.figure(1, figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.plot(poles.real, poles.imag, 'or', label='Poles')
plt.plot([0, 0], [-1, 1], 'k--')  # Vertical blue line
plt.plot([-3, 3], [0, 0], 'k--')  # Horizontal blue line
plt.legend()
plt.grid()
plt.xlim(-3, 3)
plt.ylim(-1, 1)
plt.xlabel('Real')
plt.ylabel('Imag')
plt.title(f'Poles = {np.around(poles, decimals=2)}')

print (f'Poles of H(s) = {np.around(poles, decimals=2)}')

# %% Simulation and plotting of impulse response:

dt = 0.01
t_start = 0
t_stop = 10
t_array = np.arange(t_start, t_stop+dt, dt)    

(t, y) = control.impulse_response(H, t_array)

plt.subplot(1, 2, 2)
plt.plot(t, y, 'b', label='Impulse response, h')
plt.legend()
plt.grid()
plt.xlabel('t [s]')

# plt.savefig('poles_and_sim_impulse_resp_H1.pdf')
# plt.savefig('poles_and_sim_impulse_resp_H2.pdf')
# plt.savefig('poles_and_sim_impulse_resp_H3.pdf')
plt.show()