from __future__ import print_function

import espressomd
from espressomd import shapes
from espressomd import lb
from espressomd import lbboundaries

import numpy as np


print( " " )
print( "===================================================" )
print( "=            Lattice-Boltzmann Fluid              =" )
print( "===================================================" )
print( " " )

print( "Program Information: \n" )
print( espressomd.code_info.features() )
box_l = 32
system = espressomd.System(box_l = [box_l, box_l, box_l])
system.time_step = 0.01
system.cell_system.skin = 0.2

# the force density on the fluid nodes
force_density=0.01

# choose between these two: GPU or CPU (depending on compilation features)
#lbf = lb.LBFluid(agrid=1, dens=1, visc=1, fric=20, tau=0.01, ext_force_density=[force_density, 0, 0])
lbf = lb.LBFluidGPU(agrid=1, dens=1, visc=1, fric=20, tau=0.01, ext_force_density=[force_density, 0, 0])

system.actors.add(lbf)
system.thermostat.set_lb(kT=0.0)

# create the boundary "shape"
upper_wall=shapes.Wall(normal=[0,1,0], dist=1.5)
lower_wall=shapes.Wall(normal=[0,-1,0], dist=-(box_l-1.5))

# from these shapes, define the LB boundary
upper_bound=lbboundaries.LBBoundary(shape=upper_wall)
lower_bound=lbboundaries.LBBoundary(shape=lower_wall)

system.lbboundaries.add(upper_bound)
system.lbboundaries.add(lower_bound)

#system.part.add(pos=0.5*system.box_l, type=0)

max_time=50
for t in range(max_time):
    system.integrator.run(100)

    # the velocity of the LB nodes are stored in lbf[i,j,k].velocity
    vel = lbf[box_l/2, box_l/2, box_l/2].velocity

    print("time: {} velocity:{}".format(system.time, vel))

print("**Simulation finished")