from interacticle import MoleculeCollection, LennardJones, Coulomb, Simulation

import interacticle.visualizer
from interacticle.visual import WireBox, MolecularLinks, VelocityHistogram, TemperaturePlot

from library import Argon

setup = MoleculeCollection()

setup.potential(LennardJones(39.948, 0.3395, 0.971))

for x in range(10):
    for y in range(10):
        for z in range(10):
            setup.add(Argon, (2+0.4*x,2+0.4*y,2+0.4*z))

def target_temperature(step, temperature):
    if step < 4e4: # equilibration for 20 ps
        return None
    elif step >= 4e4 and step <= 2e5: # Heat up to 300 Kelvin in 20 ps and keep there until t=100ps
        return min(max(temperature, (step-4e4) / 4e4 * 300), 300)
    else: # Cool down to 10 Kelvin within 20 ps
        return max(min(temperature, (4e4 - (step-2e5)) / 4e4 * 300), 10)

setup.domain_size = 8
setup.tau = 0.0005
setup.cutoff = 0.3395*2.5
setup.target_temperature = target_temperature

setup.max_lennard_jones = 200
setup.max_coulomb = 200

simulation = Simulation(setup, opengl = True)
simulation.verbose = False

histogram = VelocityHistogram(simulation, [1.1,0.5], [1,0.5])
temperature = TemperaturePlot(simulation, [1.1,0], [1,0.5])

interacticle.visualizer.simulate(setup, simulation, [ WireBox(setup.domain_size), MolecularLinks(simulation), histogram, temperature ], 100)