import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

figure, axes = plt.subplots()


class Plotter:
    def __init__(self):
        self.plt = None
        self.colorbar = None
        self.patch = None

    def clear(self):
        if self.patch:
            self.patch.remove()
        if self.colorbar:
            self.colorbar.remove()
        if self.plt:
            self.plt.remove()

    def plot(self, model, normalize=True, density=False, save_path=''):
        self.clear()

        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
        axes.set_title('Velocity field (normalized)')
        plt.suptitle(f'Avg mass source per grid point = {model.avg_error()}')
        plt.xlabel('X')
        plt.ylabel('Y')

        u, v = model.u, model.v
        if normalize:
            factor = np.sqrt(u ** 2 + v ** 2)
            u = u / factor
            v = v / factor

        x, y = np.meshgrid(
            np.linspace(0, model.domain_size[1], model.shape[1]),
            np.linspace(0, model.domain_size[0], model.shape[0]),
        )

        density = density or int((max(model.domain_size) / model.step) / 40)

        self.plt = plt.quiver(
            x[::density, ::density],
            y[::density, ::density],
            u[::density, ::density],
            v[::density, ::density],
            model.p[::density, ::density],
            scale=30,
            cmap='inferno'
        )
        self.colorbar = plt.colorbar(label='Pressure')

    def save(self, path):
        return plt.savefig(path, dpi=300)

    def show(self):
        return plt.pause(0.0001)