import matplotlib.pyplot as plt
import numpy as np

"""
Create Your Own Lattice Boltzmann Simulation (With Python)
Philip Mocz (2020) Princeton Univeristy, @PMocz
Simulate flow past cylinder
for an isothermal fluid
"""


def main():
    """ Finite Volume simulation """

    # Simulation parameters
    Nx = 400  # resolution x-dir
    Ny = 100  # resolution y-dir
    rho0 = 100  # average density
    tau = 0.6  # collision timescale
    Nt = 80000  # number of timesteps
    plotRealTime = True  # switch on for plotting as the simulation goes along

    # Lattice speeds / weights
    NL = 9
    idxs = np.arange(NL)
    cxs = np.array([0, 0, 1, 1, 1, 0, -1, -1, -1])
    cys = np.array([0, 1, 1, 0, -1, -1, -1, 0, 1])
    weights = np.array([4 / 9, 1 / 9, 1 / 36, 1 / 9, 1 / 36, 1 / 9, 1 / 36, 1 / 9, 1 / 36])  # sums to 1

    # Initial Conditions
    F = np.ones((Ny, Nx, NL))  # * rho0 / NL
    has_fluid = np.ones((Ny, Nx), dtype=np.bool)
    has_fluid[int(Ny/2):, :] = False
    np.random.seed(42)
    F += 0.01 * np.random.randn(Ny, Nx, NL)
    X, Y = np.meshgrid(range(Nx), range(Ny))
    F[:, :, 3] += 2 * (1 + 0.2 * np.cos(2 * np.pi * X / Nx * 4))
    # F[:, :, 5] += 1
    rho = np.sum(F, 2)
    for i in idxs:
        F[:, :, i] *= rho0 / rho

    # Cylinder boundary
    X, Y = np.meshgrid(range(Nx), range(Ny))
    cylinder = (X - Nx / 4) ** 2 + (Y - Ny / 2) ** 2 < (Ny / 4) ** 2
    inner_cylinder = (X - Nx / 4) ** 2 + (Y - Ny / 2) ** 2 < (Ny / 4 - 2) ** 2
    F[cylinder] = 0
    F[0, :] = 0
    F[Ny - 1, :] = 0
    # F[int(Ny/2):, :] = 0

    has_fluid[cylinder] = False
    has_fluid[0, :] = False
    has_fluid[Ny - 1, :] = False

    # for i in idxs:
    #     F[:, :, i] *= has_fluid

    # Prep figure
    fig = plt.figure(figsize=(4, 2), dpi=80)

    reflection_mapping = [0, 5, 6, 7, 8, 1, 2, 3, 4]
    # Simulation Main Loop
    for it in range(Nt):
        print(it)

        # Drift
        new_has_fluid = np.zeros((Ny, Nx))
        F_sum = np.sum(F, 2)
        for i, cx, cy in zip(idxs, cxs, cys):
            F_part = F[:, :, i] / F_sum
            F_part[F_sum == 0] = 0
            to_move = F_part * (has_fluid * 1.0)
            to_move = (np.roll(to_move, cx, axis=1))
            to_move = (np.roll(to_move, cy, axis=0))

            new_has_fluid += to_move

            F[:, :, i] = np.roll(F[:, :, i], cx, axis=1)
            F[:, :, i] = np.roll(F[:, :, i], cy, axis=0)

        # has_fluid = new_has_fluid > 0.5
        # new_has_fluid[F_sum == 0] += has_fluid[F_sum == 0] * 1.0
        # new_has_fluid[(np.abs(F_sum) < 0.000000001)] = 0

        fluid_sum = np.sum(has_fluid * 1.0)
        has_fluid = (new_has_fluid / np.sum(new_has_fluid * 1.0)) * fluid_sum

        print('fluid_cells: %d' % np.sum(has_fluid * 1))

        # for i in idxs:
        #     F[:, :, i] *= has_fluid

        bndry = np.zeros((Ny, Nx), dtype=np.bool)
        bndry[0, :] = True
        bndry[Ny - 1, :] = True
        # bndry[:, 0] = True
        # bndry[:, Nx - 1] = True
        bndry = np.logical_or(bndry, cylinder)

        # bndry = np.logical_or(bndry, has_fluid < 0.5)

        # Set reflective boundaries
        bndryF = F[bndry, :]
        bndryF = bndryF[:, reflection_mapping]

        sum_f = np.sum(F)
        print('Sum of Forces: %f' % sum_f)

        # sum_f_cyl = np.sum(F[cylinder])
        # print('Sum of Forces in cylinder: %f' % sum_f_cyl)

        # sum_f_inner_cyl = np.sum(F[inner_cylinder])
        # print('Sum of Forces in inner cylinder: %f' % sum_f_inner_cyl)

        # if sum_f > 4000000.000000:
        #     test = 1

        # F[Ny - 1, :, 5] += 0.1
        # F[0, :, 1] -= 0.1
        # F[0, :, 5] += 0.1
        # F[Ny - 1, :, 1] -= 0.1

        # Calculate fluid variables
        rho = np.sum(F, 2)
        ux = np.sum(F * cxs, 2) / rho
        uy = np.sum(F * cys, 2) / rho

        ux[(np.abs(rho) < 0.000000001)] = 0
        uy[(np.abs(rho) < 0.000000001)] = 0

        # print('minimum rho: %f' % np.min(np.abs(rho)))
        # print('Maximum F: %f' % np.max(F))
        # print('Minimum F: %f' % np.min(F))

        # Apply Collision
        Feq = np.zeros(F.shape)
        for i, cx, cy, w in zip(idxs, cxs, cys, weights):
            Feq[:, :, i] = rho * w * (
                        1 + 3 * (cx * ux + cy * uy) + 9 * (cx * ux + cy * uy) ** 2 / 2 - 3 * (ux ** 2 + uy ** 2) / 2)

        F += -(1.0 / tau) * (F - Feq)

        # Apply boundary
        F[bndry, :] = bndryF

        # plot in real time - color 1/2 particles blue, other half red
        if (plotRealTime and (it % 10) == 0) or (it == Nt - 1):
            plt.cla()
            ux[cylinder] = 0
            uy[cylinder] = 0
            vorticity = (np.roll(ux, -1, axis=0) - np.roll(ux, 1, axis=0)) - (
                        np.roll(uy, -1, axis=1) - np.roll(uy, 1, axis=1))
            vorticity[cylinder] = np.nan

            # vorticity *= has_fluid

            cmap = plt.cm.bwr
            cmap.set_bad('black')

            # plt.imshow(vorticity, cmap='bwr')
            plt.imshow(has_fluid * 2.0 - 1.0, cmap='bwr')
            # plt.imshow(bndry * 2.0 - 1.0, cmap='bwr')

            plt.clim(-.1, .1)
            ax = plt.gca()
            ax.invert_yaxis()
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)
            ax.set_aspect('equal')
            plt.pause(0.001)

    # Save figure
    # plt.savefig('latticeboltzmann.png', dpi=240)
    plt.show()

    return 0


if __name__ == "__main__":
    main()