import espressomd
import espressomd.lb

from espressomd import lbboundaries
from espressomd import shapes
from espressomd import interactions

import object_in_fluid as oif
import numpy as np
import os, sys
import random


def create_obstacles():
    # bottom of the channel
    tmp_shape = shapes.Rhomboid(corner=[0.0, 0.0, 0.0], a=[boxX, 0.0, 0.0], b=[0.0, boxY, 0.0], c=[0.0, 0.0, 1.0],
                                direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/wallBottom.vtk")

    # top of the channel
    tmp_shape = shapes.Rhomboid(corner=[0.0, 0.0, boxZ - 1], a=[boxX, 0.0, 0.0], b=[0.0, boxY, 0.0], c=[0.0, 0.0, 1.0],
                                direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/wallTop.vtk")

    # front wall of the channel
    tmp_shape = shapes.Rhomboid(corner=[0.0, 0.0, 0.0], a=[boxX, 0.0, 0.0], b=[0.0, 1.0, 0.0], c=[0.0, 0.0, boxZ],
                                direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/wallFront.vtk")

    # back wall of the channel
    tmp_shape = shapes.Rhomboid(corner=[0.0, boxY - 1.0, 0.0], a=[boxX, 0.0, 0.0], b=[0.0, 1.0, 0.0],
                                c=[0.0, 0.0, boxZ], direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/wallBack.vtk")

    # obstacle - cylinder A
    centerA = [11.0, 2.0, boxZ/2]
    cyl_centers.append(centerA)
    tmp_shape = shapes.Cylinder(center=centerA, axis=[0.0, 0.0, 1.0], length=boxZ, radius=cyl_radius, direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_cylinder(cyl_shape=tmp_shape, n=20, out_file=vtk_directory + "/cylinderA.vtk")

    # obstacle - cylinder B
    centerB = [11.0, 12.0, boxZ/2]
    cyl_centers.append(centerB)
    tmp_shape = shapes.Cylinder(center=centerB, axis=[0.0, 0.0, 1.0], length=boxZ, radius=cyl_radius, direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_cylinder(cyl_shape=tmp_shape, n=20, out_file=vtk_directory + "/cylinderB.vtk")

def max_velocity():
    max_vel = -1
    for i in range(0, int(boxX)):
        for j in range(0, int(boxY)):
            for k in range(0, int(boxZ)):
                vel = oif.norm(lbf[i,j,k].velocity)
                if vel > max_vel:
                    max_vel = vel
                    iout = i
                    jout = j
                    kout = k
    return [max_vel, iout, jout, kout]


sim_no = sys.argv[1]
directory = "output/sim"+str(sim_no)
os.makedirs(directory)
vtk_directory = directory+"/vtk"
os.makedirs(vtk_directory)

boxX = 22.0
boxY = 14.0
boxZ = 14.0
system = espressomd.System(box_l=[boxX, boxY, boxZ])
system.cell_system.skin = 0.2
system.time_step = 0.2

cyl_radius = 3.0
cyl_centers = list()
boundaries = []
create_obstacles()
rbc_radius = 3.91
rbc_friction = 1.82

fluid_force = 0.0002
soft_a = 0.0003
soft_n = 1.2
soft_cutoff = 0.1

# output simulation parameters to a file
output_file = open(directory+"/parameters.dat", "w")
output_file.write("sim_no = " + str(sim_no) + "\n")
output_file.write("ext_force = " + str(fluid_force) + "\n")
output_file.write("soft_a = " + str(soft_a) + "\n")
output_file.write("soft_n = " + str(soft_n) + "\n")
output_file.write("soft_cutoff = " + str(soft_cutoff) + "\n")
output_file.write("----------------------------------------\n")
output_file.close()

# creating the template for RBC
type = oif.OifCellType(nodes_file="input/rbc374nodes.dat", triangles_file="input/rbc374triangles.dat",
                        check_orientation=False, system=system, ks=0.0127, kb=0.0002, kal=0.0254,
                        kag=0.5, kv=0.9, resize=[rbc_radius, rbc_radius, rbc_radius])

# creating the RBC
cell_rbc = oif.OifCell(cell_type=type, particle_type=0, origin=[4.0,7.0,7.0], particle_mass=0.5, rotate = [np.pi/2, 0.0, 0.0])
cell_rbc.output_vtk_pos(vtk_directory+"/cell_0.vtk")

# creating boundaries
for boundary in boundaries:
    system.lbboundaries.add(lbboundaries.LBBoundary(shape=boundary))
    system.constraints.add(shape=boundary, particle_type=100, penetrable=False)

# cell-wall interactions
system.non_bonded_inter[0, 100].soft_sphere.set_params(a=soft_a, n=soft_n, cutoff=soft_cutoff, offset=0.0)

# fluid
lb_params = {'agrid': 1, 'dens': 1, 'visc': 1.5, 'tau': system.time_step, 'ext_force_density': [fluid_force, 0.0, 0.0]}
lbf = espressomd.lb.LBFluid(**lb_params)
system.actors.add(lbf)
system.thermostat.set_lb(LB_fluid=lbf, gamma=1.5)

# main integration loop
maxCycle = 500
integr_steps = 500
for i in range(1, maxCycle):
    system.integrator.run(steps=integr_steps)
    cell_rbc.output_vtk_pos_folded(file_name=vtk_directory + "/cell_" + str(i) + ".vtk")
    print "time: ", str(i * system.time_step * integr_steps), " max fluid velocity: ", max_velocity()
print "Simulation completed."