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

dir_ = Path(__file__).parent.resolve()
dir_root = dir_ / ".."
dir_assets = dir_root / "assets"

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


class AnimatedPlot(ABC):
    fig: Figure
    ax: Union[Axes, Tuple[Axes]]

    def __init__(self, frames: int = 100):
        self.frames = int(frames)
        self.fig, self.ax = plt.subplots(1, 1)

    def save(self, filename: Union[Path, str], framerate: int = 30):
        log.info(f"Generating animation: {self.__class__}...")
        an = FuncAnimation(
            self.fig,
            self.update,
            frames=np.linspace(0, 1, self.frames, endpoint=False),
            init_func=self.init,
            blit=False,
        )
        an.save(str(filename), writer="pillow", fps=framerate)
        log.info(f"Generating animation: {self.__class__}...Done")

    def init(self) -> None:
        pass

    @abstractmethod
    def update(self, t: float) -> None:
        pass


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()