import espressomd
import object_in_fluid as oif

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

import numpy as np
import os, sys


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 A
    tmp_shape = shapes.Rhomboid(corner=[15.0, 1.0, 1.0], a=[5.0, 0.0, 0.0], b=[0.0, boxY - 2.0, 0.0], c=[0.0, 0.0, 2.75],
                    direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/obstacleA.vtk")

    # obstacle B
    tmp_shape = shapes.Rhomboid(corner=[15.0, 1.0, boxZ - 2.75], a=[5.0, 0.0, 0.0], b=[0.0, boxY - 2.0, 0.0],
                                      c=[0.0, 0.0, 2.75], direction=1)
    boundaries.append(tmp_shape)
    oif.output_vtk_rhomboid(rhom_shape=tmp_shape, out_file=vtk_directory + "/obstacleB.vtk")

# this is a script to try the nucleus with non-bonded interactions
sim_no = sys.argv[1]
directory = "output/sim"+str(sim_no)
os.makedirs(directory)
vtk_directory = directory+"/vtk"
os.makedirs(vtk_directory)

boxX = 40.0
boxY = 20.0
boxZ = 20.0
time_step = 0.1
system = espressomd.System(box_l=[boxX, boxY, boxZ])
system.cell_system.skin = 0.2
system.time_step = time_step


# elastic coefficients
membrane_ks = 0.2
membrane_kb = 0.16
membrane_kal = 0.2
membrane_kag = 0.9
membrane_kv = 0.5
membrane_kvisc = 0.0
nucleus_ks = 1.4
nucleus_kb = 1.4
nucleus_kal = 1.4
nucleus_kag = 2.5
nucleus_kv = 3.5
nucleus_kvisc = 0.0

membraneNodes = "input/sphere304nodes.dat"
nucleusNodes = "input/sphere169nodes.dat"
membraneTriangles = "input/sphere304triangles.dat"
nucleusTriangles = "input/sphere169triangles.dat"

# creating templates
membrane_type = oif.OifCellType(nodes_file=membraneNodes, triangles_file=membraneTriangles, check_orientation=False,
                           system=system, ks=membrane_ks, kb=membrane_kb, kal=membrane_kal, kag=membrane_kag,
                           kv=membrane_kv, kvisc=membrane_kvisc, resize=[1.0, 1.0, 1.0], normal=True)
nucleus_type = oif.OifCellType(nodes_file=nucleusNodes, triangles_file=nucleusTriangles, check_orientation=False,
                          system=system, ks=nucleus_ks, kb=nucleus_kb, kal=nucleus_kal, kag=nucleus_kag,
                          kv=nucleus_kv, kvisc=nucleus_kvisc, resize=[0.3, 0.3, 0.3], normal=True)

# creating the cell
membrane = oif.OifCell(cell_type = membrane_type, particle_type=0, origin=[7.0,7.0,7.0])
nucleus = oif.OifCell(cell_type = nucleus_type, particle_type=1, origin=[7.0,7.0,7.0])
system.non_bonded_inter[0,1].soft_sphere.set_params(a = 0.003, n = 3, cutoff = 3, offset = 0.0)

# fluid
lbf = espressomd.lb.LBFluid(agrid = 1, dens = 1.0, visc = 1.5, tau = time_step, fric = 1.5 , ext_force_density = [0.001, 0.0, 0.0])
system.actors.add(lbf)

boundaries = []
create_obstacles()
for boundary in boundaries:
    system.lbboundaries.add(lbboundaries.LBBoundary(shape=boundary))
    system.constraints.add(shape=boundary, particle_type=10, penetrable=False)

# cell-wall interactions
system.non_bonded_inter[0,10].soft_sphere.set_params(a = 0.0001, n = 1.2, cutoff = 0.1, offset = 0.0)
system.non_bonded_inter[1,10].soft_sphere.set_params(a = 0.0001, n = 1.2, cutoff = 0.1, offset = 0.0)

maxCycle = 100
# main integration loop
membrane.output_vtk_pos_folded(file_name="output/sim" + str(sim_no) + "/membrane_0.vtk")
nucleus.output_vtk_pos_folded(file_name="output/sim" + str(sim_no) + "/nucleus_0.vtk")

integr_steps = 500
for i in range(1,maxCycle):
    system.integrator.run(steps=integr_steps)
    membrane.output_vtk_pos_folded(file_name="output/sim" + str(sim_no) + "/membrane_" + str(i) + ".vtk")
    nucleus.output_vtk_pos_folded(file_name="output/sim" + str(sim_no) + "/nucleus_" + str(i) + ".vtk")
    print "time: ", str(i*time_step*integr_steps)