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:
            self.plt.remove()

    def plot(self, model, normalize=True, density=False, save_path=''):
        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')

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


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

        u = copy(model.u)
        u.resize(shape)
        v = copy(model.v)
        v.resize(shape)
        p = copy(model.p)
        p.resize(shape)

        print(shape, u.shape, v.shape)

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

        plt.contourf(x, y, p)
        # self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
        # TODO: allow using streamplot
        self.plt = plt.quiver(
            x,
            y,
            u,
            v,
        )
        self.colorbar = plt.colorbar(label='Pressure')

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

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