#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt


def passeBas(x, Q=1):
    """
    Filtre passe-bas en pulsation réduite *x* = omega/omega0, facteur de
    qualité *Q*.
    """

    return 1 / (1 - x ** 2 + x / Q * 1j)


def passeHaut(x, Q=1):

    return -x ** 2 / (1 - x ** 2 + x / Q * 1j)


def passeBande(x, Q=1):

    return 1 / (1 + Q * (x - 1 / x) * 1j)


def coupeBande(x, Q=1):

    return (1 - x ** 2) / (1 - x ** 2 + x / Q * 1j)


def gainNphase(f, dB=True):
    """
    Retourne le gain (éventuellement en dB) et la phase [rad] d'un
    filtre de fonction de transfert complexe *f*.
    """

    g = np.abs(f)                        # Gain
    if dB:                              # [dB]
        g = 20 * np.log10(g)
    p = np.angle(f)                      # [rad]

    return g, p


def asympGain(x, pentes=(0, -40)):

    lx = np.log10(x)
    return np.where(lx < 0, pentes[0] * lx, pentes[1] * lx)


def asympPhase(x, phases=(0, -np.pi)):

    return np.where(x < 1, phases[0], phases[1])


def diagBode(x, filtres, labels,
             title='', plim=None, gAsymp=None, pAsymp=None):
    """
    Trace le diagramme de Bode -- gain [dB] et phase [rad] -- des filtres
    de fonction de transfert complexe *filtres* en fonction de la pulsation
    réduite *x*.
    """

    fig, (axg, axp) = plt.subplots(2, 1, sharex=True)        # 2×1 axes
    axg.set(xscale='log',                                    # Axe des gains
            ylabel='Gain [dB]')
    axp.set(xlabel=r'x = $\omega$/$\omega_0$', xscale='log', # Axe des phases
            ylabel='Phase [rad]')

    lstyles = ['--', '-', '-.', ':']
    for f, label, ls in zip(filtres, labels, lstyles):   # Tracé des courbes
        g, p = gainNphase(f, dB=True)       # Calcul du gain et de la phase
        axg.plot(x, g, lw=2, ls=ls, label=f"Q={label}")  # Gain
        axp.plot(x, p, lw=2, ls=ls)                      # Phase

    # Asymptotes
    if gAsymp is not None:              # Gain
        axg.plot(x, asympGain(x, gAsymp), 'k:', lw=2)
    if pAsymp is not None:              # Phase
        # axp.plot(x, asympPhase(x,pAsymp), 'k:')
        pass

    axg.legend(fontsize='small')

    # Labels des phases
    axp.set_yticks(np.arange(-2, 2.1) * np.pi / 2)
    axp.set_yticks(np.arange(-4, 4.1) * np.pi / 4, minor=True)
    axp.set_yticklabels([r'$-\pi$', r'$-\pi/2$', r'$0$', r'$\pi/2$', r'$\pi$'])
    # Domaine des phases
    if plim is not None:
        axp.set_ylim(plim)

    # Ajouter les grilles
    for ax in (axg, axp):
        ax.grid()                           # x et y, majors
        ax.grid(which='minor')              # x et y, minors

    # Ajustements fins
    gmin, gmax = axg.get_ylim()
    axg.set_ylim(gmin, max(gmax, 3))

    fig.subplots_adjust(hspace=0.1)
    axg.xaxis.set_major_formatter(plt.matplotlib.ticker.ScalarFormatter())
    plt.setp(axg.get_xticklabels(), visible=False)

    if title:
        axg.set_title(title)

    return fig


if __name__ == '__main__':

    x = np.logspace(-1, 1, 1000)              # de 0.1 à 10 en 1000 pas

    # Facteurs de qualité
    qs = [0.25, 1 / np.sqrt(2), 5]            # Valeurs numériques
    labels = [0.25, r'$1/\sqrt{2}$', 5]      # Labels

    # Calcul des fonctions de transfert complexes
    pbs = [ passeBas(x, Q=q) for q in qs ]
    phs = [ passeHaut(x, Q=q) for q in qs ]
    pcs = [ passeBande(x, Q=q) for q in qs ]
    cbs = [ coupeBande(x, Q=q) for q in qs ]

    # Création des 4 diagrammes de Bode
    figPB = diagBode(x, pbs, labels, title='Filtre passe-bas',
                     plim=(-np.pi, 0),
                     gAsymp=(0, -40), pAsymp=(0, -np.pi))
    figPH = diagBode(x, phs, labels, title='Filtre passe-haut',
                     plim=(0, np.pi),
                     gAsymp=(40, 0), pAsymp=(np.pi, 0))
    figPC = diagBode(x, pcs, labels, title='Filtre passe-bande',
                     plim=(-np.pi / 2, np.pi / 2),
                     gAsymp=(20, -20), pAsymp=(np.pi / 2, -np.pi / 2))
    figCB = diagBode(x, cbs, labels, title='Filtre coupe-bande',
                     plim=(-np.pi / 2, np.pi / 2),
                     gAsymp=(0, 0), pAsymp=(0, 0))

    plt.show()