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 = [16.0, 8.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")

    # obstacle - cylinder C
    centerC = [11.0, 12.0, boxZ/2]
    cyl_centers.append(centerC)
    tmp_shape = shapes.Cylinder(center=centerC, 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 + "/cylinderC.vtk")


# find distance to the closest obstacle
def dist_to_closest_obstacle (x,y):
    min_dist = 100000000.0
    for center in cyl_centers:
        dist = oif.vec_distance([center[0], center[1], 0], [x, y, 0])
        if dist < min_dist:
            min_dist = dist
    if min_dist < cyl_radius + rbc_radius:
        min_dist = 0
    else:
        min_dist = min_dist - cyl_radius - rbc_radius
    return min_dist

if len(sys.argv)!= 3:
    print "2 arguments are expected:"
    print "number of cells: n_cells"
    print "id of the simulation: sim_no"
    print " "
else:
    n_cells = int(sys.argv[1])
    sim_no = sys.argv[2]

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.1

cyl_radius = 2.0
cyl_centers = list()
boundaries = []
create_obstacles()
rbc_radius = 2.0

# random seeding of cells
cell_positions = list()
for k in range(0,n_cells):
    origin_ok = 0
    while origin_ok !=1:
         # generate random position in channel;
         ox = random.random() * (boxX - 2*rbc_radius) + rbc_radius
         oy = random.random() * (boxY - 2*rbc_radius - 2) + 1 + rbc_radius
         oz = random.random() * (boxZ - 2*rbc_radius - 2) + 1 + rbc_radius
         origin_ok = 1

         # check whether cells are outside of obstacles
         if dist_to_closest_obstacle(ox, oy) < 1:
             origin_ok = 0

         # check that it does not collide with other rbc
         if origin_ok == 1:
             for i in range(0,k):
                 dist = oif.vec_distance([ox, oy, oz], cell_positions[i])
                 if dist < 2*rbc_radius:
                     origin_ok = 0
                     break

         # if everything was ok, remember origin
         if origin_ok == 1:
             print ("seeding cell: "+str(k)+" with origin: "+str(ox)+", "+str(oy)+", "+str(oz))
             cell_positions.append([ox,oy,oz])

# creating the template for RBCs
type = oif.OifCellType(nodes_file="input/rbc374nodes.dat", triangles_file="input/rbc374triangles.dat",
                        check_orientation=False, system=system, ks=0.02, kb=0.016, kal=0.02,
                        kag=0.9, kv=0.5, normal=True, resize=[rbc_radius, rbc_radius, rbc_radius])

# creating the RBCs
rbcs = list()
for id, pos in enumerate(cell_positions):
    cell_rbc = oif.OifCell(cell_type=type, particle_type=id, origin=pos,
                           rotate=[random.random()*2*np.pi, random.random()*2*np.pi, random.random()*2*np.pi],
                           particle_mass=0.5)
    rbcs.append(cell_rbc)
for id, rbc in enumerate(rbcs):
    rbc.output_vtk_pos(vtk_directory+"/cell"+str(id)+"_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
for i in range(n_cells):
    system.non_bonded_inter[i, 100].soft_sphere.set_params(a=0.0001, n=1.2, cutoff=0.1, offset=0.0)

# cell-cell interactions
for i in range(n_cells):
     for j in range (n_cells):
         if (i<j):
             system.non_bonded_inter[i, j].membrane_collision.set_params(a=0.0001, n=1.2, cutoff=0.1, offset=0.0)

# fluid
lb_params = {'agrid': 1, 'dens': 1, 'visc': 1.5, 'tau': system.time_step, 'ext_force_density': [0.002, 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 = 100
integr_steps = 500
for i in range(1, maxCycle):
    system.integrator.run(steps=integr_steps)
    for id, rbc in enumerate(rbcs):
        rbc.output_vtk_pos_folded(file_name=vtk_directory + "/cell" + str(id) + "_" + str(i) + ".vtk")
    print "time: ", str(i*system.time_step*integr_steps)
print "Simulation completed."