import numpy as np
import espressomd
from espressomd import shapes
from espressomd import visualization
from threading import Thread
from Molecule import Molecule

box_l               = 264.0
np_radius           = 5.0
res_radius          = 20.0
struct_file         = '1n5u.struct'
#pot_file            = 'gold_5/1n5u_gold_5_0.pot'
pot_file            = 'test_pot.pot'
use_display         = True

kT                  = 1.0
# Based on the viscosity of water
#gamma               = 6300
gamma               = 1.0
skin_depth          = 6.0

# Each time step is 6.228769e-13 seconds
ts                  = 6.228769e-13

steps               = 10000
warm_steps          = 2000
warm_force_cap      = 2000.0
init_cap            = 10.0
bound_dist          = 4.0

#1fs
#warm_time_step      = 0.7
warm_time_step      = 0.025


mol_repusion        = 2000.0;
mol_repusion_cutoff = 1.0

#datetime_tag = datetime.datetime.now().strftime("%d-%b-%Y_%H:%M:%S")

print("##############################")
print("#                            #")
print("#      Initializing Run      #")
print("#                            #")
print("##############################")

espressomd.assert_features(['EXTERNAL_FORCES', 'MASS', 'EXCLUSIONS', 'TABULATED', 'HAT'])

system = espressomd.system.System (box_l = [ box_l, box_l, box_l ], periodicity = [ True, True, True ])

# This needs to be optimized to speed up the simulation
system.cell_system.skin = skin_depth
system.cell_system.set_domain_decomposition()
system.thermostat.set_langevin (kT = kT, gamma = gamma, seed = 42)

# What does this do?
system.set_random_state_PRNG()

# Insert spherical nanoparticle of type 0
system.constraints.add(
    particle_type = 0,
    penetrable = False,
    shape = shapes.Sphere(
        center = system.box_l / 2.0,
        radius = np_radius,
        direction = 1
    )
)

# Insert spherical reservior of type 1
#system.constraints.add(
#    particle_type = 1,
#    penetrable = True,
#    shape = shapes.Sphere(
#        center = system.box_l / 2.0,
#        radius = res_radius,
#        direction = 1
#    )
#)

# Load molecule
molecule = Molecule(system, struct_file, res_radius)

# Insert molecules into the system
system.part.add(pos = molecule.positions(), type = molecule.types(), id = molecule.ids(), mass = molecule.masses())

# Molecule - Molecule repulsion
for first, second in molecule.type_pairs():
    system.non_bonded_inter[first, second].hat.set_params(F_max = mol_repusion, cutoff = mol_repusion_cutoff)

# Nanoparticle interaction
for first, rmin, rmax, energy, force in molecule.np_bonds(pot_file):
    system.non_bonded_inter[0, first].tabulated.set_params(min = rmin, max = rmax, energy = energy, force = force)

# Exclusions
for first, second in molecule.exclusion_pairs():
    system.part[first].add_exclusion(second)

# Bonds
for first, second, bond in molecule.bond_pairs():
    system.part[first].add_bond((bond, second))

# Angle Bonds
for first, second, third, angle in molecule.angle_pairs():
    system.part[second].add_bond((angle, first, third))

# Print system information
print("Starting: max_cut_bonded {}".format(system.max_cut_bonded))
print("Starting: max_cut_nonbonded {}".format(system.max_cut_nonbonded))

for key, value in system.cell_system.get_state().items():
    print("Starting: {} {}".format(key, value))

# Integration loop
print("############################")
print("#                          #")
print("#      Warming System      #")
print("#                          #")
print("############################")

#handle = open("results_{}.out".format(datetime_tag), 'w')
#step   = 0
#ts     = 6.33e-13

#h5 = h5md.H5md(filename = "trajectory_{}.h5".format(datetime_tag), write_pos = True, write_vel = True)

size = molecule.size()
mols = molecule.ntotal()
visualizer = None

x = 0
y = 0
z = 0
minDist = 0
l = system.box_l[0] / 2
isBound = np.zeros((mols), dtype = bool)

def main_thread():
    system.time_step = warm_time_step

    for force in np.linspace(init_cap, warm_force_cap, warm_steps):
        print(f"Starting: force_cap = {force}")
        if visualizer:
            visualizer.update()
        system.force_cap = force
        system.integrator.run (1000)


        minDist = 1000.0
        for i in np.arange(mols):
            for j in np.arange(size):
                x = system.part[i * size + j].pos_folded[0]
                y = system.part[i * size + j].pos_folded[1]
                z = system.part[i * size + j].pos_folded[2]
                ssd = ((x-l)**2 + (y-l)**2 + (z-l)**2)**0.5

                if ssd < minDist:
                    minDist = ssd

                if ssd < np_radius + bound_dist:
                    isBound[i] = True
                    break
            else:
                isBound[i] = False
        print(f"Starting: bound = {np.sum(isBound)}, minDist = {minDist}")

    system.force_cap = 0

    while True:
        if visualizer:
            visualizer.update()

        system.integrator.run (steps)

        minDist = 1000.0

        for i in np.arange(mols):
            for j in np.arange(size):
                x = system.part[i * size + j].pos_folded[0]
                y = system.part[i * size + j].pos_folded[1]
                z = system.part[i * size + j].pos_folded[2]
                ssd = ((x-l)**2 + (y-l)**2 + (z-l)**2)**0.5

                if ssd < minDist:
                    minDist = ssd

                if ssd < np_radius + bound_dist:
                    isBound[i] = True
                    break
            else:
                isBound[i] = False

        print(f"{system.time * ts * warm_time_step:.3e}, {minDist}, {np.sum(isBound)}")


if use_display:
    visualizer = visualization.openGLLive(system)
    thread = Thread (target = main_thread)
    thread.daemon = True
    thread.start()
    visualizer.start()
else:
    main_thread()