# %%
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Tuple, Union

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from matplotlib.gridspec import GridSpec
from scipy import signal


# %%
def db20(x):
    return 20 * np.log10(x)


# %%
wpass = 20 * 2 * np.pi
wstop = 30 * 2 * np.pi
gpass = 2
gstop = 30

o, wn = signal.cheb1ord(
    wpass,
    wstop,
    2,
    30,
    analog=True,
)
zpk_cheb1 = signal.cheby1(o, gpass, wn, analog=True, output="zpk", btype="lowpass")

o, wn = signal.cheb2ord(
    wpass,
    wstop,
    gpass,
    gstop,
    analog=True,
)
zpk_cheb2 = signal.cheby2(o, gpass, wn, analog=True, output="zpk", btype="lowpass")

o, wn = signal.buttord(wpass, wstop, gpass, gstop, analog=True)
zpk_butter = signal.butter(o, wn, btype="lowpass", analog=True, output="zpk")

o, wn = signal.ellipord(wpass, wstop, gpass, gstop, analog=True)
zpk_ellip = signal.ellip(o, gpass, gstop, wn, btype="lowpass", analog=True, output="zpk")

# %%
fig = plt.figure(layout="constrained")
gs = GridSpec(2, 2, figure=fig)
ax_zp = fig.add_subplot(gs[:, 0])
ax_f = fig.add_subplot(gs[0, 1])
ax_t = fig.add_subplot(gs[1, 1])

for z, p, k in [
    zpk_cheb1,
    # zpk_cheb2,
    zpk_butter,
    zpk_ellip,
]:
    lines_z = ax_zp.plot(np.real(z), np.imag(z), "o", mfc="none")
    lines_p = ax_zp.plot(np.real(p), np.imag(p), "x", color=lines_z[0].get_color())
    # plt.xlim(-150, 150)
    ax_zp.grid(True)
    ax_zp.set_xlabel("Real(s)")
    ax_zp.set_ylabel("Imag(s)")

    f = np.linspace(0.1 * np.max([wpass, wstop]), 10 * np.max([wpass, wstop]), 1001) / (2 * np.pi)
    ax_f.semilogx(f, db20(signal.freqs_zpk(z, p, k, f * 2 * np.pi)[1]))
    ax_f.grid(True)
    ax_f.set_xlabel("Frequnecy [Hz]")
    ax_f.set_ylabel("Gain [dB]")
    ax_f.set_xlim(f[0], f[-1])

    # t = np.arange(101)
    # x = np.zeros(len(t))
    # x[int(len(x) / 2)] = 1
    # ax_t.plot(t, signal.lfilter(*ba, x))
    # ax_t.set_ylim(-10, 100)
    # ax_t.grid(True)

# %%