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")

    # obstacle - cylinder
    centerA = [21.0, 7.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 + "/cylinder.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 = int(sys.argv[1])
directory = "output/sim"+str(sim_no)
os.makedirs(directory)
vtk_directory = directory+"/vtk"
os.makedirs(vtk_directory)

boxX = 28.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.00004
soft_a = 0.0003
soft_n = 1.2
soft_cutoff = 0.1
mem_a = 0.03
mem_n = 2.5
mem_cutoff = 0.2

# 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("mem_a = " + str(mem_a) + "\n")
output_file.write("mem_n = " + str(mem_n) + "\n")
output_file.write("mem_cutoff = " + str(mem_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.9, kv=0.5, resize=[rbc_radius, rbc_radius, rbc_radius], normal = True)

# creating the RBCs
rbc1 = oif.OifCell(cell_type=type, particle_type=0, origin=[14.0,6.0,7.0], particle_mass=0.5,
                   rotate=[np.pi/2, 0.0, 0.0])
rbc2 = oif.OifCell(cell_type=type, particle_type=1, origin=[6.0,9.0,7.0], particle_mass=0.5,
                   rotate=[np.pi/2, 0.0, 0.0])

# 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)
system.non_bonded_inter[1, 100].soft_sphere.set_params(a=soft_a, n=soft_n, cutoff=soft_cutoff, offset=0.0)

# cell-cell interactions
system.non_bonded_inter[0,1].membrane_collision.set_params(a=mem_a, n=mem_n, cutoff=mem_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
max_cycle = 100
max_vel = -1
integr_steps = 500
for i in range(1, max_cycle+1):
    system.integrator.run(steps=integr_steps)
    rbc1.output_vtk_pos_folded(file_name=vtk_directory + "/cell1_" + str(i) + ".vtk")
    rbc2.output_vtk_pos_folded(file_name=vtk_directory + "/cell2_" + str(i) + ".vtk")
    curr_vel = oif.norm(lbf[20,13,7].velocity)
    if curr_vel>max_vel:
        max_vel = curr_vel
    print "time: ", str(i * system.time_step*integr_steps), " fluid velocity: ", curr_vel
print "Simulation completed."

# output simulation data to file
output_file = open(directory+"/parameters.dat", "a")
output_file.write("max fluid velocity = " + str(max_vel) + "\n")
output_file.write("Tfinal = " + str(max_cycle * system.time_step*integr_steps) + "\n")
output_file.close()