import numpy as np


class SIMPLE:
    def __init__(self, domain_size, bfs_size, step, Re, alpha=0.8):
        np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)

        self.Re = Re
        self.nu = 1 / Re
        self.alpha = alpha

        self.domain_size = domain_size
        self.bfs_size = bfs_size

        self.step = step
        self.shape = tuple(int(x / step) for x in domain_size)
        self.bfs_node_size = tuple(int(x / step) for x in bfs_size)

        # Allocations
        self.u = self.allocate_field()
        self.u_star = self.allocate_field()

        self.v = self.allocate_field()
        self.v_star = self.allocate_field()

        self.p = self.allocate_field()
        self.p_star = self.allocate_field(True)
        self.p_prime = self.allocate_field()

        self.d_e = self.allocate_field()
        self.d_n = self.allocate_field()
        self.b = self.allocate_field()

    def allocate_field(self, random=False):
        if random:
            return np.random.rand(*self.shape)
        return np.zeros(shape=self.shape, dtype=float)

    def apply_inflow_boundary(self):
        for i in range(1, self.shape[0] - 1):
            self.u_star[i][0] = 1
            self.v_star[i][0] = 0

    def apply_outflow_boundary(self):
        for i in range(0, self.shape[0]):
            self.u_star[i][-1] = self.u_star[i][-2];

    def assert_positive(self, value):
        '''Assert that the value is nearly positive'''
        assert value > -0.1, f'WARNING: Value must be positive: {value}'
        return value

    def grid(self):
        '''Iterator over all grid points, excluding the obstacle'''
        for i in range(1, self.shape[0] - 1):
            for j in range(1, self.shape[1] - 1):
                # if i > self.bfs_node_size[0] or j > self.bfs_node_size[1]:
                yield (i, j)

    def solve_momentum_equations(self):
        # Momentum along X direction
        for i, j in self.grid():
            u_E = 0.5 * (self.u[i][j] + self.u[i][j + 1])
            u_W = 0.5 * (self.u[i][j] + self.u[i][j - 1])
            v_N = 0.5 * (self.v[i - 1][j] + self.v[i - 1][j + 1])
            v_S = 0.5 * (self.v[i][j] + self.v[i][j + 1])

            a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
            a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
            a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
            a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)

            a_e = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
            A_e = -self.step

            self.d_e[i][j] = A_e / a_e

            self.u_star[i][j] = (
                a_E * self.u[i][j + 1] +
                a_W * self.u[i][j - 1] +
                a_N * self.u[i - 1][j] +
                a_S * self.u[i + 1][j]
            ) / a_e + self.d_e[i][j] * (self.p_star[i][j + 1] - self.p_star[i][j])

        # Momentum along Y direction
        for i, j in self.grid():
            u_E = 0.5 * (self.u[i][j] + self.u[i + 1][j])
            u_W = 0.5 * (self.u[i][j - 1] + self.u[i + 1][j - 1])
            v_N = 0.5 * (self.v[i - 1][j] + self.v[i][j])
            v_S = 0.5 * (self.v[i][j] + self.v[i + 1][j])

            a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
            a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
            a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
            a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)

            a_n = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
            A_n = -self.step

            self.d_n[i][j] = A_n / a_n

            self.v_star[i][j] = (
                a_E * self.v[i][j + 1] +
                a_W * self.v[i][j - 1] +
                a_N * self.v[i - 1][j] +
                a_S * self.v[i + 1][j]
            ) / a_n + self.d_n[i][j] * (self.p_star[i][j] - self.p_star[i + 1][j])

    def apply_sides_boundary(self):
        for j in range(self.shape[1]):
            self.v_star[0][j] = 0
            self.v_star[-2][j] = 0
            # WORKSNICE: self.v_star[-3][j] = 0

    def apply_bfs_boundary(self):
        '''Apply Backwards Facing Step boundary conditions'''
        for i in range(self.bfs_node_size[0]):
            self.u_star[i][self.bfs_node_size[1]] = 0
        for j in range(self.bfs_node_size[1]):
            self.v_star[self.bfs_node_size[0]][j] = 0

    def correct_pressure(self):
        self.p_prime = self.allocate_field()
        for i, j in self.grid():
            a_E = self.assert_positive(-self.d_e[i][j] * self.step)
            a_W = self.assert_positive(-self.d_e[i][j-1] * self.step)
            a_N = self.assert_positive(-self.d_n[i-1][j] * self.step)
            a_S = self.assert_positive(-self.d_n[i][j] * self.step)
            a_P = a_E + a_W + a_N + a_S

            self.b[i][j] = self.step * (
                -(self.u_star[i][j] - self.u_star[i][j-1])
                + (self.v_star[i][j] - self.v_star[i-1][j])
            )

            self.p_prime[i][j] = (
                a_E * self.p_prime[i][j+1] +
                a_W * self.p_prime[i][j-1] +
                a_N * self.p_prime[i-1][j] +
                a_S * self.p_prime[i+1][j] +
                self.b[i][j]
            ) / a_P

        self.p = self.p_star + self.p_prime * self.alpha
        self.p_star = self.p

    def correct_velocities(self):
        for i, j in self.grid():
            self.u[i][j] = self.u_star[i][j] + self.d_e[i][j] * (self.p_prime[i][j + 1] - self.p_prime[i][j])
            self.v[i][j] = self.v_star[i][j] + self.d_n[i][j] * (self.p_prime[i][j] - self.p_prime[i + 1][j])

    def iterate(self):
        self.apply_inflow_boundary()
        self.solve_momentum_equations()

        self.apply_sides_boundary()
        # self.apply_bfs_boundary()

        self.correct_pressure()
        self.correct_velocities()

        # Sides (magick wtf?)
        for j in range(self.shape[1]):
            self.v[1][j] = 0
            self.v[-2][j] = 0

        for i in range(self.shape[0]):
            self.v[i][0] = 0
            self.v[i][1] = 0


    def avg_error(self):
        return np.absolute(self.b).sum()

    def save(self, path):
        print('SAVE', path)
        with open(path, 'wb') as file:
            np.save(file, self.u)
            np.save(file, self.v)
            np.save(file, self.p)
            np.save(file, self.b)

    def load(self, path):
        print('LOAD', path)
        with open(path, 'rb') as file:
            self.u = np.load(file)
            self.v = np.load(file)
            self.p = np.load(file)
            self.b = np.load(file)
            self.p_star = self.p