import numpy as np
import matplotlib.pyplot as plt
from copy import copy
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:
            try:
                self.plt.remove()
            except:
                pass

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

        axes.set_title('Velocity field (normalized)')
        plt.suptitle(f'Avg mass source per grid point = {model.avg_error()}')
        plt.xlabel('X')
        plt.ylabel('Y')

        shape = model.p.shape

        u = np.zeros(shape, dtype=float)
        v = np.zeros(shape, dtype=float)

        for i in range(shape[0]):
            for j in range(shape[1]):
                u[i][j] = 0.5 * (model.u[i][j] + model.u[i][j + 1])
                v[i][j] = 0.5 * (model.v[i][j] + model.v[i + 1][j])

        assert not v[0, :].any()
        assert not v[-1, :].any()

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

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

        density = density or int(min(model.p.shape) / 40)

        plt.contourf(x, y, model.p, cmap='inferno')
        self.colorbar = plt.colorbar(label='Pressure')

        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(list(x * model.step for x in model.bfs_shape)), color='gray'))

        plotter = plt.streamplot if streamplot else plt.quiver
        self.plt = plotter(
            x[::density, ::density],
            y[::density, ::density],
            u[::density, ::density],
            v[::density, ::density],
            color='black'
        )


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

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