"""
    zplane(B,A) plots the zeros and poles with the unit circle for reference.
    Each zero is represented with a 'o' and each pole with a 'x' on the plot.
    Multiple zeros and poles are indicated by the multiplicity number shown
    to the upper right of the zero or pole. B and A are row vectors containing 
    transfer function polynomial coefficients of B(z)/A(z).

    Author: Antonin Novak
    Laboratoire d'Acoustique de l'Université du Mans
    (LAUM, UMR CNRS 6613), 72085 Le Mans, France.
    email address: antonin.novak@univ-lemans.fr
    Website: https://ant-novak.com
    January 2020; Last revision: 03-Feb-2021
"""


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


def show_roots(roots, ax, **parameters):
    """ show_roots plots the roots inside a unit circle.
        Multiple roots are grouped together and a number is placed next to it

    Args:
        roots (numpy array): array of roots to plot
        ax (figure axis): axis of the figure to plot the roots
        **parameters:
    """
    roots_unique, roots_counts = np.unique(roots, return_counts=True)
    for root, count in zip(roots_unique, roots_counts):
        ax.scatter(root.real, root.imag, s=50, **parameters)
        if count > 1:
            ax.text(root.real, root.imag, " $^" + str(count) + "$")


def zplane(b, a, axis=None):
    """zplane(B, A, axis) plots on the axis (if ax not provided, creates a new figure).
    B and A are numpy arrays containing transfer function polynomial coefficients of B(z)/A(z).

    Args:
        b (numpy array): transfer function polynomial coefficients of B(z)
        a (numpy array): transfer function polynomial coefficients of A(z)
        axis (figure axis, optional): axis to plot to. Defaults to None.
    """
    if axis == None:
        # get a new figure/plot
        _, ax = plt.subplots()
    else:
        ax = axis

    # set axis properties
    ax.set_aspect(1)
    ax.set(xlim=[-1.5, 1.5], ylim=[-1.5, 1.5])
    ax.set(xlabel='real(z)', ylabel='imag(z)')
    ax.grid(True)

    # draw the unit circle
    ax.add_artist(plt.Circle((0, 0), 1, linestyle='--', color='black', fill=False, lw=2))

    # plot zeros and poles
    z, p, _ = signal.tf2zpk(b, a)
    show_roots(z, ax, marker='o', facecolor='white', edgecolors="blue", lw=2)
    show_roots(p, ax, marker='x', facecolor='red', lw=2)

    # if new figure was created by this function, show the plot
    if axis == None:
        plt.show()


if __name__ == '__main__':
    # butterworth high-pass filter
    b, a = signal.butter(2, 0.2, 'high')
    # plot the zeros and poles
    zplane(b, a)