from __future__ import print_function
import espressomd
import numpy as np
from espressomd import visualization
from threading import Thread

# set up system
box_l = 150.
system = espressomd.System()
system.time_step = 0.01
system.cell_system.skin = 0.5
system.box_l = [box_l, box_l, box_l]

# set up LJ, without offset subtractLJ-bond works, with it does not work!
#system.non_bonded_inter[0, 0].lennard_jones.set_params(epsilon = 1., sigma = 4., cutoff = 4. * np.power(2., 1. / 6.), shift = "auto") 
system.non_bonded_inter[0, 0].lennard_jones.set_params(epsilon = 1., sigma = 4., cutoff = 4. * np.power(2., 1. / 6.), offset = 14., shift = "auto") 

# set up bonds
bond = espressomd.interactions.HarmonicBond(k = 210., r_0 = 3.4)
system.bonded_inter.add(bond)
anglebond = espressomd.interactions.Angle_Harmonic(bend = 750., phi0 = np.pi)
system.bonded_inter.add(anglebond)
subtLJbond = espressomd.interactions.Subt_Lj(k = 1., r = 50.)
system.bonded_inter.add(subtLJbond)

# set up system
for i in range(15):
    system.part.add(id = i, pos = [box_l / 2., box_l / 2., box_l / 2. + (i - 7.) * 3.4], type = 0) 
    if (i > 0):
        system.bonded_inter.add(bond) 
        system.part[i].add_bond((bond, i - 1)) 
        if (i > 1):
            system.part[i - 1].add_bond((anglebond, i - 2, i)) 

# exclude all LJ-interactions in the chain using subt-LJ-bonds
for i in range(15):
    for j in range(i + 1, 15):
        system.part[i].add_bond((subtLJbond, j)) 

# print force on particle should be ~ 0.
print(system.part[7].f)
system.integrator.run(0)
print(system.part[7].f)

# warmup
lj_cap = 5 
system.non_bonded_inter.set_force_cap(lj_cap)
for i in range(1000):
    system.integrator.run(100)
    lj_cap += 1.0
    system.non_bonded_inter.set_force_cap(lj_cap)
system.non_bonded_inter.set_force_cap(0)

# visualize
visualizer = visualization.mayaviLive(system)
def main():
    while True:
        print(system.part[7].f)
        visualizer.update()
        system.integrator.run(10)
        rawInput = raw_input("\t\t+Press Enter to continue...")
t = Thread(target=main)
t.daemon = True
t.start()
visualizer.start()