import numpy as np
from scipy.signal import convolve


class StaggeredArray2D:
    def __init__(self, x_n, y_n):
        """
        creates a staggered array
        :param x_n: x size of the array
        :param y_n: y size of the array
        """
        self.x_n = x_n
        self.y_n = y_n
        self.p = np.zeros((x_n, y_n), dtype=np.float)

        self.u_x = np.zeros((x_n + 1, y_n), dtype=np.float)
        self.u_y = np.zeros((x_n, y_n + 1), dtype=np.float)

        self.has_fluid = np.zeros((x_n, y_n), dtype=np.bool)
        self.density = np.zeros((x_n, y_n), dtype=np.float)

    def get_velocity(self, x, y):
        """
        get mid point value for the coordinates
        :param x: x coordinate
        :param y: y coordinate
        :return:
        """
        assert 0 <= x < self.x_n, 'x value out of bounds!'
        assert 0 <= y < self.y_n, 'y value out of bounds!'

        lower_x = self.u_x[x, y]
        upper_x = self.u_x[x + 1, y]

        lower_y = self.u_y[x, y]
        upper_y = self.u_y[x, y + 1]

        return (lower_x + upper_x) / 2.0, (lower_y + upper_y) / 2.0

    def get_velocity_arrays(self):
        c_x = np.array([[0.5], [0.5]])
        u_x = convolve(self.u_x, c_x, mode='valid')

        c_y = np.array([[0.5, 0.5]])
        u_y = convolve(self.u_y, c_y, mode='valid')

        return u_x, u_y


class StaggeredArray3D:
    def __init__(self, x_n, y_n, z_n):
        """
        creates a staggered array
        :param x_n: x size of the array
        :param y_n: y size of the array
        :param z_n: z size of the array
        """

        self.x_n = x_n
        self.y_n = y_n
        self.z_n = z_n

        self.p = np.zeros((x_n, y_n, z_n), dtype=np.float)

        self.u_x = np.zeros((x_n + 1, y_n, z_n), dtype=np.float)
        self.u_y = np.zeros((x_n, y_n + 1, z_n), dtype=np.float)
        self.u_z = np.zeros((x_n, y_n, z_n + 1), dtype=np.float)

        self.has_fluid = np.zeros((x_n, y_n, z_n), dtype=np.bool)

    def get_velocity(self, x, y, z):
        """
        get mid point value for the coordinates
        :param x: x coordinate
        :param y: y coordinate
        :param z: z coordinate
        :return:
        """
        assert 0 <= x < self.x_n, 'x value out of bounds!'
        assert 0 <= y < self.y_n, 'y value out of bounds!'
        assert 0 <= z < self.z_n, 'y value out of bounds!'

        lower_x = self.u_x[x, y, z]
        upper_x = self.u_x[x + 1, y, z]

        lower_y = self.u_y[x, y, z]
        upper_y = self.u_y[x, y + 1, z]

        lower_z = self.u_z[x, y, z]
        upper_z = self.u_z[x, y, z + 1]

        return (lower_x + upper_x) / 2.0, (lower_y + upper_y) / 2.0, (lower_z + upper_z) / 2.0


if __name__ == '__main__':
    sa = StaggeredArray2D(10, 10)

    for x in range(11):
        for y in range(10):
            sa.u_x[x, y] = y

    for x in range(10):
        for y in range(11):
            sa.u_y[x, y] = x

    ux, uy = sa.get_velocity_arrays()

    ux2, uy2 = sa.get_velocity(0, 0)