# series RLC, frequency response
# hz 2024

import numpy
import matplotlib.pyplot as plt

from matplotlib.widgets import Button, Slider, RadioButtons


omegas = numpy.logspace(-9, 9, int(1e4))


def H(omegas, which, R, L, C):
    zr = R
    zl = 1j * omegas * L
    zc = 1 / (1j * omegas * C)
    num = {"R": zr, "C": zc, "L": zl}[which]
    return num / (zr + zl + zc)


# low-, mid-, and high-frequency asymptotes for each freq resp
asymptotes = {
    "R": [
        lambda o, r, l, c: o * r * c,
        lambda o, r, l, c: r / (l * o),
        lambda o, r, l, c: 1,
    ],
    "L": [
        lambda o, r, l, c: l * c * o**2,
        lambda o, r, l, c: 1,
        lambda o, r, l, c: o * l / r,
    ],
    "C": [
        lambda o, r, l, c: 1,
        lambda o, r, l, c: 1 / (l * c * o**2),
        lambda o, r, l, c: 1 / (r * c * o),
    ],
}

asymptotes = {k: [numpy.vectorize(i) for i in v] for k, v in asymptotes.items()}

# an attempt at colorblind-friendly colors
# red, blue, and yellow from okabe-ito color palette
colors = "#d55e00", "#56b4e9", "#f0e442"
labels = "low asymptote", "high asymptote", "mid asymptote"

# set up the UI
fig, ax = plt.subplots()

# constrain the actual plot so the sliders, etc can fit
fig.subplots_adjust(bottom=0.25, right=0.9)

# radio buttons for which graph we want to show
radioax = fig.add_axes([0.92, 0.45, 0.07, 0.06])
radio = RadioButtons(radioax, ("R", "L", "C"))

radioax2 = fig.add_axes([0.92, 0.35, 0.07, 0.06])
radio2 = RadioButtons(radioax2, ("mag", "phase"))

# add sliders for R, L, C
axr = fig.add_axes([0.15, 0.1, 0.65, 0.02])
r_slider = Slider(
    ax=axr,
    label=r"R (ohms)",
    valmin=-9,
    valmax=6,
    valinit=0,
)

axl = fig.add_axes([0.15, 0.075, 0.65, 0.02])
l_slider = Slider(
    ax=axl,
    label=r"L (henries)",
    valmin=-6,
    valmax=6,
    valinit=0,
)

axc = fig.add_axes([0.15, 0.05, 0.65, 0.02])
c_slider = Slider(
    ax=axc,
    label=r"C (farads)",
    valmin=-6,
    valmax=6,
    valinit=0,
)

resetax = fig.add_axes([0.8, 0.0, 0.1, 0.04])
reset_button = Button(resetax, "Reset", hovercolor="0.975")


# this function updates the graph based on the current slider values
def update(val):
    for s in [r_slider, l_slider, c_slider]:
        s.valtext.set_text("10^(%.02f)" % s.val)

    ax.clear()

    r, l, c = 10**r_slider.val, 10**l_slider.val, 10**c_slider.val
    natural_frequency = 1 / numpy.sqrt(l * c)

    # always include the natural frequency so that we don't miss it
    # (the graphs for R can be misleading if we don't sample at/near w0)
    o = numpy.array(sorted([*omegas.tolist(), natural_frequency]))
    freq_response = H(o, radio.value_selected, r, l, c)

    type_ = radio2.value_selected

    # show the location of the natural frequency on the horizontal axis
    ax.semilogx(
        [natural_frequency, natural_frequency],
        [-500, 300] if type_ == "mag" else [-3.14159, 3.14159],
        "#cccccc",
        linestyle="dotted",
    )

    if type_ == "mag":
        vals = 20 * numpy.log10(numpy.abs(freq_response))
        # draw asymptotes first
        for a, color, label in zip(asymptotes[radio.value_selected], colors, labels):
            avals = 20 * numpy.log10(numpy.abs(a(o, r, l, c)))
            ax.semilogx(o, avals, color, label=label)
        # keep the axes the same so we're not shifting things around
        ax.axis([10**-9, 10**9, -300, 200])
        ax.set_title(r"Series RLC: $\left|H_%s(\omega)\right|$" % radio.value_selected)
        ax.legend()
    else:
        vals = numpy.angle(freq_response)
        ax.set_title(r"Series RLC: $\angle(H_%s(\omega))$" % radio.value_selected)

    # now draw the actual function of interest
    ax.semilogx(o, vals, "k", linewidth=4)

    # show Q
    base, exp = ("%.2E" % (numpy.sqrt(l / c) / r)).split("E")
    q = r"%s\times10^{%s}" % (base, exp.lstrip("+"))
    ax.set_xlabel("$Q = %s$" % q)

    fig.canvas.draw_idle()


# this function resets the sliders to their initial values
def reset(event):
    r_slider.reset()
    l_slider.reset()
    c_slider.reset()


# set things up so that the graph updates as the sliders change, etc.
r_slider.on_changed(update)
l_slider.on_changed(update)
c_slider.on_changed(update)
radio.on_clicked(update)
radio2.on_clicked(update)
reset_button.on_clicked(reset)

update(None)
plt.show()