#!/usr/bin/env python
# coding: utf-8

# ## Convert RKF Trajectory to DCD Format
#
# This example converts `ams.rkf` to DCD and PSF, mapping atoms back into molecules in each frame before writing.
# A short MD run is included first to produce the RKF input file.
#

# ### Initial Imports

from pathlib import Path

from scm.flexmd import PSFTopology, pdb_from_plamsmol
from scm.plams import (
    DCDTrajectoryFile,
    RKFTrajectoryFile,
    AMSJob,
    Settings,
    from_smiles,
    packmol,
    view,
)


# ### Run MD Simulation

# Run a short toy MD simulation to generate an `ams.rkf` trajectory.
#

water_box = packmol(from_smiles("O"), n_molecules=8, density=1.0)
view(water_box, direction="tilt_z", width=300, height=300, picture_path="picture1.png")


md_settings = Settings()
md_settings.input.ams.Task = "MolecularDynamics"
md_settings.input.ReaxFF.ForceField = "Water2017.ff"
md_settings.input.ams.MolecularDynamics.NSteps = 30
md_settings.input.ams.MolecularDynamics.TimeStep = 0.5
md_settings.input.ams.MolecularDynamics.Trajectory.SamplingFreq = 1
md_settings.input.ams.MolecularDynamics.InitialVelocities.Temperature = 300
md_settings.runscript.nproc = 1
md_settings.runscript.preamble_lines = ["export OMP_NUM_THREADS=1"]

md_job = AMSJob(name="md_for_dcd", molecule=water_box, settings=md_settings)
md_results = md_job.run()

rkf_path = Path(md_results.rkfpath())


# ### Convert Frames

# Convert RKF frames to DCD while remapping atoms into molecular entities.
#

rkf = RKFTrajectoryFile(str(rkf_path))
mol = rkf.get_plamsmol()
print(f"NSteps: {len(rkf)}")

pdb = pdb_from_plamsmol(mol)
psf = PSFTopology(pdb=pdb)
psf.write_psf("ams.psf")

dcd = DCDTrajectoryFile("ams.dcd", mode="wb")
for i in range(len(rkf)):
    coords, cell = rkf.read_frame(i)
    mol.from_array(coords)
    mol.map_atoms_to_bonds()
    dcd.write_next(coords=mol.as_array(), cell=cell)

print("Created files 'ams.dcd' and 'ams.psf'")