from abc import abstractmethod, ABC
from typing import TypeVar, Generic, Optional, Callable

import numpy as np
from matplotlib import cm
from matplotlib.animation import FuncAnimation
from matplotlib.axes import Axes
from matplotlib.figure import Figure

from diffeq_solver_tk.differential_equation_metadata import OrdinaryDifferentialEquationMetadata, BoundedEquationMetadata
from diffeq_solver_tk.util import export_bounded_equation_solution, export_ode_solution

T = TypeVar('T')


class DifferentialEquationService(Generic[T], ABC):
    solve: Callable[[T], np.ndarray]
    main_figure: Optional[Figure]
    metadata: Optional[T] = None
    solution: Optional[np.ndarray] = None

    def __init__(self, solve: Callable[[T], np.ndarray], main_figure: Optional[Figure] = None):
        self.solve = solve
        self.main_figure = main_figure

    def compute_and_update_solution(self, metadata: T):
        self.solution = self.solve(metadata)
        self.metadata = metadata
        self.render_current_solution()

    def clear_solution(self):
        self.metadata = None
        self.solution = None
        self.clear_figure()

    def clear_figure(self):
        self.main_figure.clf()

    @abstractmethod
    def export_solution(self, table_path: str):
        raise NotImplementedError

    @abstractmethod
    def render_current_solution(self):
        raise NotImplementedError


class OrdinaryDifferentialEquationService(DifferentialEquationService[OrdinaryDifferentialEquationMetadata], ABC):
    def render_current_solution(self):
        if self.main_figure is None:
            return
        domain = np.linspace(0, self.metadata.time, self.metadata.samples)
        ax = self.main_figure.add_subplot()
        ax.plot(domain, self.solution)
        ax.set_xlabel("t [time]")
        ax.set_ylabel("x(t)")
        ax.plot()
        ax.set_title("Solution: x(t)")

    def export_solution(self, table_path: str):
        export_ode_solution(self.solution, self.metadata, table_path)


class BoundedEquationService(DifferentialEquationService[BoundedEquationMetadata], ABC):
    __animation: Optional[FuncAnimation] = None
    __is_animation_playing: bool = False

    def render_current_solution(self):
        if self.main_figure is None:
            return
        self.clear_figure()
        self.clear_animation()
        K, N = self.solution.shape
        time_domain = np.linspace(0, self.metadata.time, K)
        space_domain = np.linspace(0, self.metadata.length, N)
        x, t = np.meshgrid(space_domain, time_domain)
        ax = self.main_figure.add_subplot(projection='3d')
        ax.set_xlabel('x [length]')
        ax.set_ylabel('t [time]')
        ax.set_zlabel('u(x, t)')
        surf = ax.plot_surface(x, t, self.solution, cmap=cm.coolwarm, linewidth=0, antialiased=False)
        self.main_figure.colorbar(surf, shrink=0.5, aspect=5)

    def generate_animation(self):
        if self.main_figure is None:
            return
        self.clear_figure()
        K, N = self.solution.shape
        ax = self.main_figure.add_subplot()
        Axes.set_xlim(ax, left=0, right=self.metadata.length)
        Axes.set_ylim(ax, bottom=np.min(self.solution), top=np.max(self.solution))
        ax.set_xlabel("x [length]")
        ax.set_ylabel("u(x, t)")
        ax.set_title("Solution: u(x, t)")
        space_domain = np.linspace(0, self.metadata.length, self.metadata.samples)
        line, = ax.plot(space_domain, self.solution[0], "-o", markersize=4)
        time_text = ax.text(0.82, 0.92, '', transform=ax.transAxes)
        dt = self.metadata.time / K

        def update_plot(k: int):
            line.set_ydata(self.solution[k])
            time_text.set_text("t = %.3f" % (k * dt))
            return line, time_text

        self.__animation = FuncAnimation(self.main_figure, update_plot, frames=K, blit=True, interval=20)

    def toggle_animation(self):
        if not self.__animation:
            self.generate_animation()
            self.__is_animation_playing = True
        elif self.is_animation_playing():
            self.__animation.pause()
            self.__is_animation_playing = False
        else:
            self.__animation.resume()
            self.__is_animation_playing = True

    def is_animation_playing(self):
        return self.__is_animation_playing

    def export_solution(self, table_path: str):
        export_bounded_equation_solution(self.solution, self.metadata, table_path)

    def clear_animation(self):
        if self.main_figure is None or self.__animation is None:
            return
        self.__animation.pause()
        self.__animation = None
        self.__is_animation_playing = False
        self.main_figure.clf()
