import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as signal

def plot_block_diagram_axes(ax, label_num, label_den):
    ax.arrow(0.05, 0.5, 0.13, 0, head_width=0.05, head_length=0.02, fc='k', ec='k')
    rect = plt.Rectangle((0.22, 0.3), 0.3, 0.6, fill=None, edgecolor='black', lw=2)
    ax.add_patch(rect)
    ax.text(0.32, 0.5, f'H(p) = ({label_num})\n      ({label_den})', fontsize=10, ha='center', va='center')
    ax.arrow(0.47, 0.5, 0.22, 0, head_width=0.05, head_length=0.02, fc='k', ec='k')
    ax.text(0.01, 0.5, 'E(p)', ha='right', va='center', fontsize=11)
    ax.text(0.70, 0.5, 'S(p)', ha='left', va='center', fontsize=11)
    ax.set_xlim(0, 0.8)
    ax.set_ylim(0, 1)
    ax.axis('off')

def plot_pole_zero_and_step_response_with_block(tf_num, tf_den):
    system = signal.TransferFunction(tf_num, tf_den)
    poles = system.poles
    zeros = system.zeros
    fig = plt.figure(figsize=(10, 8))
    gs = fig.add_gridspec(3, 2, height_ratios=[1,3,3])
    ax0 = fig.add_subplot(gs[0, :])
    label_num = ' + '.join([f'{c} p^{len(tf_num)-1-i}' if i < len(tf_num)-1 else f'{c}' for i, c in enumerate(tf_num)])
    label_den = ' + '.join([f'{c} p^{len(tf_den)-1-i}' if i < len(tf_den)-1 else f'{c}' for i, c in enumerate(tf_den)])
    plot_block_diagram_axes(ax0, label_num, label_den)
    ax0.set_title('Schéma bloc du système')
    ax1 = fig.add_subplot(gs[1, 0])
    ax1.scatter(np.real(zeros), np.imag(zeros), marker='o', label='Zéros')
    ax1.scatter(np.real(poles), np.imag(poles), marker='x', label='Pôles')
    ax1.axhline(0, color='grey', lw=1)
    ax1.axvline(0, color='grey', lw=1)
    ax1.set_xlabel('Re')
    ax1.set_ylabel('Im')
    ax1.set_title('Plan des pôles et zéros')
    ax1.grid(True)
    ax1.legend()
    t, y = signal.step(system)
    ax2 = fig.add_subplot(gs[1, 1])
    ax2.plot(t, y)
    ax2.set_xlabel('Temps (s)')
    ax2.set_ylabel('Amplitude')
    ax2.set_title('Réponse indicielle')
    ax2.grid(True)
    plt.tight_layout()
    plt.show()



#exercice 1 du cours sur la Stabilité

plot_pole_zero_and_step_response_with_block([1], [1,1,-1,3]) #coeff des puissances décroissantes

plot_pole_zero_and_step_response_with_block([3], [8, 16 ,10,2])