# %% imports
import logging

import numpy as np
from matplotlib import pyplot as plt
from utils import AnimatedPlot, dir_assets

# %% logging
log = logging.Logger(__name__)
log.setLevel(logging.INFO)


# %%
class TxLinePlot(AnimatedPlot):
    x = np.linspace(-2 * np.pi, 2 * np.pi, 500)

    def update(self, t: float):
        v = 2 * np.pi
        self.ax.clear()

        self.ax.plot(self.x, np.real(np.exp(1j * (t * v - self.x))), label="$V$")

        self.ax.autoscale(enable=True, axis="x", tight=True)
        self.ax.set_ylim(-1.1, 1.1)

        self.ax.set_xlabel("z")
        self.ax.set_ylabel("V")
        self.ax.set_title(f"$t={t:0.2f}$")

        def format_func(value, tick_number):
            # find number of multiples of pi/2
            N = int(np.round(2 * value / np.pi))
            if N == 0:
                return "0"
            elif np.abs(N) == 1:
                return r"${0}\pi/2$".format("-" if N < 0 else "")
            elif np.abs(N) == 2:
                return r"${0}\pi$".format("-" if N < 0 else "")
            elif N % 2 > 0:
                return r"${0}\pi/2$".format(N)
            else:
                return r"${0}\pi$".format(N // 2)

        self.ax.xaxis.set_major_locator(plt.MultipleLocator(np.pi / 2))
        self.ax.xaxis.set_major_formatter(plt.FuncFormatter(format_func))
        self.ax.grid(True)
        self.fig.tight_layout()


class SuperpositionPlot(AnimatedPlot):
    x = np.linspace(-2 * np.pi, 2 * np.pi, 500)

    def update(self, t: float):
        v = 2 * np.pi
        self.ax.clear()

        vf = np.real(np.exp(1j * (t * v - self.x)))
        vr = np.real(np.exp(1j * (t * v + self.x)))

        self.ax.plot(self.x, vf, label="$V_f$")
        self.ax.plot(self.x, vr, label="$V_r$")
        self.ax.plot(self.x, vf + vr, color="black", linestyle="dotted", label="$V$")

        self.ax.autoscale(enable=True, axis="x", tight=True)
        self.ax.set_ylim(-2.2, 2.2)

        self.ax.set_xlabel("z")
        self.ax.set_ylabel("V")
        self.ax.set_title(f"$t={t:0.2f}$")

        def pi_ticks(value, tick_number):
            # find number of multiples of pi/2
            N = int(np.round(2 * value / np.pi))
            if N == 0:
                return "0"
            elif np.abs(N) == 2:
                # +/- 1 * pi
                return "$" + ("-" if N < 0 else "") + r"\pi$"
            elif N % 2 == 0:
                return "$" + f"{N // 2}" + r"\pi$"
            else:
                return "$" + ("-" if N < 0 else "") + r"\frac{" + f"{np.abs(N)}" + r"}{2}\pi$"

        self.ax.xaxis.set_major_locator(plt.MultipleLocator(np.pi / 2))
        self.ax.xaxis.set_major_formatter(plt.FuncFormatter(pi_ticks))
        self.ax.grid(True)
        self.ax.legend(loc="lower right")
        self.fig.tight_layout()


class ReflectionPlot(AnimatedPlot):
    x = np.linspace(-4 * np.pi, 0, 500)

    def __init__(self, zl: complex, **kwargs):
        super().__init__(**kwargs)

        self.zl = zl
        self.z0 = 1

    def update(self, t: float):
        v = 2 * np.pi
        self.ax.clear()

        gamma = 1 if self.zl == np.inf else (self.zl - self.z0) / (self.zl + self.z0)

        vf = np.exp(1j * (t * v - self.x))
        vr = gamma * np.exp(1j * (t * v + self.x))

        self.ax.plot(self.x, np.real(vf), label="$V_f$")
        self.ax.plot(self.x, np.real(vr), label="$V_r$")
        self.ax.plot(self.x, np.real(vf + vr), color="black", linestyle="dotted", label="$V$")
        envelope = np.abs(vf + vr)
        self.ax.fill_between(self.x, envelope, -envelope, color="black", alpha=0.1)

        self.ax.autoscale(enable=True, axis="x", tight=True)
        self.ax.set_ylim(-2.2, 2.2)

        self.ax.set_xlabel("z")
        self.ax.set_ylabel("V")
        self.ax.set_title(f"$Z_L/Z_0={self.zl}$, $t={t:0.2f}$")

        def pi_ticks(value, tick_number):
            # find number of multiples of pi/2
            N = int(np.round(2 * value / np.pi))
            if N == 0:
                return "0"
            elif np.abs(N) == 2:
                # +/- 1 * pi
                return "$" + ("-" if N < 0 else "") + r"\pi$"
            elif N % 2 == 0:
                return "$" + f"{N // 2}" + r"\pi$"
            else:
                return "$" + ("-" if N < 0 else "") + r"\frac{" + f"{np.abs(N)}" + r"}{2}\pi$"

        self.ax.xaxis.set_major_locator(plt.MultipleLocator(np.pi / 2))
        self.ax.xaxis.set_major_formatter(plt.FuncFormatter(pi_ticks))
        self.ax.grid(True)
        self.ax.legend(loc="lower left")
        self.fig.tight_layout()


# class SmithPlot:
#     def __init__(self, )


def generate():
    TxLinePlot().save(dir_assets / "tx_lines" / "tx_line.gif")
    SuperpositionPlot().save(dir_assets / "tx_lines" / "superposition.gif")
    ReflectionPlot(zl=1).save(dir_assets / "tx_lines" / "reflection_matched.gif")
    ReflectionPlot(zl=np.inf).save(dir_assets / "tx_lines" / "reflection_open.gif")
    ReflectionPlot(zl=0).save(dir_assets / "tx_lines" / "reflection_short.gif")
    ReflectionPlot(zl=0.5).save(dir_assets / "tx_lines" / "reflection_real.gif")
    ReflectionPlot(zl=1 + 1j).save(dir_assets / "tx_lines" / "reflection_complex.gif")


# %%
if __name__ == "__main__":
    generate()