from espressomd import thermostat
from espressomd import interactions
from espressomd import polymer
from espressomd.shapes import SimplePore
from espressomd import visualization
from threading import Thread
from espressomd.io.writer import vtf  # pylint: disable=import-error
import numpy as np

----------
LJ_EPS = 1.0
LJ_SIG = 1.0
LJ_CUT = 2.5 * LJ_SIG

system.non_bonded_inter[0, 0].lennard_jones.set_params(
    epsilon=LJ_EPS, sigma=LJ_SIG, cutoff=LJ_CUT, shift='auto')
-----------
pore = SimplePore(axis=[1, 0, 0],
                  length=10,
                  radius=1.6,
                  smoothing_radius=0.5,
                  center=[25, 5, 5])
system.constraints.add(shape=pore, particle_type=0, penetrable=False)
-----------
warm_steps = 10
LJ_CAP = 0.5
system.force_cap = LJ_CAP
#wca_cap = 1
#system.force_cap = wca_cap
i = 0
#act_min_dist = 0.90 it is just to check while loop
act_min_dist = system.analysis.min_dist()
print(act_min_dist)
system.thermostat.set_langevin(kT=0.0, gamma=1.0)


# slowly ramp un up the cap
while (act_min_dist < 0.82):
    vtf.writevcf(system, outfile)
    print("min_dist: {} \t force cap: {}".format(act_min_dist, LJ_CAP))
    system.integrator.run(warm_steps)
    system.part[:].v = [0, 0, 0] ### with this, you set all the particle's velocities to zero
    # Warmup criterion
    act_min_dist = system.analysis.min_dist()
    LJ_CAP = LJ_CAP * 1.01
    system.force_cap = LJ_CAP


LJ_CAP = 0
system.force_cap = LJ_CAP
system.integrator.run(warm_steps * 10)


system.thermostat.set_langevin(kT=1.0, gamma=1.0)
system.integrator.run(warm_steps * 10)
print("Finished warmup")

system.thermostat.turn_off()
lbf = espressomd.lb.LBFluid(kT=1, seed=123, agrid=1, dens=1, visc=7, tau=0.01, ext_force_density=[0.07,0,0])
system.actors.add(lbf)
system.thermostat.set_lb(LB_fluid=lbf, seed=142, gamma=5)

logging.info("Warming up the system with LB fluid.")
system.integrator.run(100)
logging.info("LB fluid warming finished.")
-------------
print("simulating...")
t_steps = 1000
for t in range(t_steps):
    print("step {} of {}".format(t, t_steps))
    system.integrator.run(warm_steps)
    vtf.writevcf(system, outfile)

outfile.close()

visualizer = visualization.openGLLive(system)

def main_thread():
    while True:
        system.integrator.run(1)
        visualizer.update()

t = Thread(target=main_thread)
t.daemon = True
t.start()
visualizer.start()