from asap3 import *
from asap3.md.verlet import VelocityVerlet
from ase.lattice.cubic import FaceCenteredCubic
from asap3.io.trajectory import *
from ase.parallel import world
from asap3.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
from asap3.constraints import FixAtoms
from asap3.testtools import ReportTest
from ase.io import read
import os

print_version(1)
#set_verbose(1)

ismaster = world.rank == 0
isparallel = world.size != 1
if world.size == 1:
    cpulayout = None
elif world.size == 2:
    cpulayout = [2,1,1]
elif world.size == 3:
    cpulayout = [1,3,1]
elif world.size == 4:
    cpulayout = [2,1,2]
else:
    cpulayout = 'auto'

delete = True
temp = 3000  # K

if ismaster:
    initial = FaceCenteredCubic(size=(10,10,10), symbol="Cu", pbc=(1,0,0))
else:
    initial = None
if isparallel:
    atoms = MakeParallelAtoms(initial, cpulayout)
    atoms.set_constraint(FixAtoms(mask=(atoms.get_ids() < atoms.get_global_number_of_atoms()/2)))
else:
    atoms = initial.copy()
    atoms.set_constraint(FixAtoms(indices=range(len(atoms)//2)))

atoms.set_calculator(EMT())
MaxwellBoltzmannDistribution(atoms, temperature_K=temp)
Stationary(atoms)

dyn = Langevin(atoms, 3*units.fs, temperature_K=temp, friction=0.02)
dyn.run(1000)

temps = [atoms.get_temperature()]
for i in range(10):
    dyn.run(10)
    temps.append(atoms.get_temperature())
    
avgtemp = np.mean(temps)
print("Average temperature T={:.1f}K  (task {})".format(avgtemp, world.rank))
c = atoms.constraints[0]
print("Removed DOF: {}   (task {})".format(c.get_removed_dof(atoms), world.rank))
ReportTest("Mean temperature:", avgtemp, temp, 200)

ReportTest.Summary()
