import numpy as np
from scipy import signal
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd

def zdiagrammes(system):
    """
        Trace les diagrammes de Bode, de Black-Nichols et de Nyquist.

        Cette fonction trace les diagrammes associés à une fonction de transfert rationnelle de numérateur b et de dénomniteur a.

        Parametres
        ----------
        system : scipy.signal TransferFunctionContinuous / StateSpaceContinuous / ZerosPolesGainContinuous
            système à temps continu
        """
    
    num = system.to_tf().num
    den = system.to_tf().den
    w, H = signal.freqz(num,den)

    # Bode diagramm
    gain_db = 20 * np.log10(np.abs(H))
    phase_deg = np.unwrap(np.angle(H)) * 180 / np.pi

    fig = go.Figure()

    # Courbe de gain
    fig.add_trace(
        go.Scatter(
            x=w,
            y=gain_db,
            name="Gain (dB)",
            yaxis="y1"
        )
    )

    # Courbe de phase
    fig.add_trace(
        go.Scatter(
            x=w,
            y=phase_deg,
            name="Phase (deg)",
            yaxis="y2"
        )
    )

    # Set figure properties
    fig.update_layout(
        width = 500,
        height = 500,
        title="Module et phase de la TFTD",
        xaxis=dict(
            title="Pulsation (rad/s)",
        ),
        yaxis=dict(
            title="Gain (dB)",
        ),
        yaxis2=dict(
            title="Phase (°)",
            overlaying="y",
            side="right",
            tickmode="linear",
            dtick=90
        ),
        legend=dict(
            x=0.66,
            y=0.99
        )
    )
    fig.show()
