# %% imports
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

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

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


# %%
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