from ase import Atoms
import numpy as np

# If you do not have a GPU, you could use CPUNEP to
# create the trajectory instead, see calorine documentation
from calorine.calculators import GPUNEP

""" Setup structures """

# Primitive cell (from https://next-gen.materialsproject.org/materials/mp-1487#crystal_structure)
NiAl_primitive = Atoms('NiAl', positions=[(0, 0, 0), (1/2, 1/2, 1/2)])
NiAl_primitive.set_cell([2.86, 2.86, 2.86], scale_atoms=True)
NiAl_primitive.set_pbc((True, True, True))

# Ordered supercell
NiAl_supercell = NiAl_primitive.repeat((24, 24, 24))
NiAl_supercell.info = {'dirname': 'ordered'}

# Disordered supercell
NiAl_disordered_supercell = NiAl_supercell.copy()
NiAl_disordered_supercell.info = {'dirname': 'disordered'}
for i in range(17240):
    random_id_1 = np.random.randint(0, len(NiAl_disordered_supercell))
    random_id_2 = np.random.randint(0, len(NiAl_disordered_supercell))
    NiAl_disordered_supercell.positions[[
        random_id_1, random_id_2]] = NiAl_disordered_supercell.positions[[random_id_2, random_id_1]]

""" Run MD """
# Run e.g. wget https://zenodo.org/records/10081677/files/UNEP-v1-model-without-ZBL.txt to
# retreive the NEP model

model_file = 'UNEP-v1-model-without-ZBL.txt'
temperature = 300
nsteps = 100000

for conf in [NiAl_supercell, NiAl_disordered_supercell]:
    dirname = conf.info['dirname']
    calc = GPUNEP(model_file, command='/path/to/gpumd_executable')
    calc.set_directory(dirname)
    conf.calc = calc

    parameters = [('time_step', 5.0),
                  # First run short NVT equilibration
                  ('velocity', 2 * temperature),
                  ('ensemble', ('nvt_lan', temperature, temperature, 100)),
                  ('dump_thermo', 100),
                  ('dump_restart', nsteps//5),
                  ('run', nsteps),
                  # Then run NVE simulation and dump positions and velocities every 5th frame
                  ('ensemble', 'nve'),
                  ('dump_thermo', 100),
                  ('dump_exyz', (5, 1)),
                  ('run', nsteps)]
    calc.run_custom_md(parameters)
