import espressomd.lb
import espressomd.visualization_opengl
import espressomd.script_interface

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm
import scipy.ndimage.filters

required_features = ["LB_WALBERLA"]
espressomd.assert_features(required_features)

np.random.seed(42)

BOX_L = 4.
N = 20
AGRID = BOX_L / N

# generate a random landscape
X = np.linspace(0., BOX_L, N)
Y = np.linspace(0., BOX_L, N)
X, Y = np.meshgrid(X, Y)
Z = scipy.ndimage.filters.gaussian_filter(np.random.uniform(size=(N, N)), sigma=3., mode='wrap')
Z = Z - np.min(Z)
Z = np.log(1. + Z)
Z /= np.linalg.norm(Z)

# 3D plot
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, Z, rstride=2, cstride=2, antialiased=True, cmap=matplotlib.cm.viridis)
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$y$')
ax.set_zticks([])
ax.view_init(60, -60)
plt.show()

# ESPResSo system with unthermalized LB fluid
system = espressomd.System(box_l=[N, N, 12.])
system.time_step = 0.01
system.cell_system.skin = 0.4
lb_fluid = espressomd.lb.LBFluidWalberla(agrid=1., density=1., viscosity=1., tau=0.01)
system.actors.add(lb_fluid)

# generate a 3D bitmask to create a velocity-bounce-back LB boundary
Z = 1 + 8 * Z / np.max(Z)
ubb = np.zeros(lb_fluid.shape, dtype=float)
mask = np.zeros(lb_fluid.shape, dtype=int)
for i in range(lb_fluid.shape[0]):
    for j in range(lb_fluid.shape[1]):
        k = int(np.round(Z[j, i]))
        mask[i, j, 0:k] = 1
lb_fluid.call_method(
    'add_boundary_from_shape',
    raster=espressomd.script_interface.array_variant(mask.reshape([-1])),
    velocity=espressomd.script_interface.array_variant(ubb.reshape([-1])))

# configure the visualizer
visualization_mode = 'grid'
if visualization_mode == 'grid':
    params = {
      'LB_draw_boundaries': False,
      'LB_draw_nodes': False,
      'LB_draw_node_boundaries': True,
    }
elif visualization_mode == 'viridis_blocks':
    params = {
      'LB_draw_boundaries': True,
      'LB_draw_nodes': False,
      'LB_draw_node_boundaries': False,
      'rasterize_color': lambda x, y, z: matplotlib.cm.viridis(z / np.max(Z)),
    }
visualizer = espressomd.visualization_opengl.openGLLive(
    system,
    background_color=[1, 1, 1],
    camera_position=[0, 0, 50],
    rasterize_pointsize=30.,
    **params)

visualizer.run(0)