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

# ## Extract Frames from ams.rkf Trajectory
#
# This example first runs a short toy molecular dynamics simulation of a small water box. It then loads the generated `ams.rkf` file and extracts PLAMS `Molecule` objects from the frames.

# Import the PLAMS tools used for setup, MD, trajectory reading, and structure visualization.
#

from typing import List
from scm.plams import AMSJob, Molecule, Settings, Trajectory, from_smiles, packmol, view


# ### Run Short Water-Box MD Simulation
#
# Build a small water box with Packmol and run a very short ReaxFF MD simulation.
#

water_box = packmol(from_smiles("O"), n_molecules=8, density=1.0)

md_settings = Settings()
md_settings.input.ams.Task = "MolecularDynamics"
md_settings.input.ReaxFF.ForceField = "Water2017.ff"
md_settings.input.ams.MolecularDynamics.NSteps = 20
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="toy_water_box_md", molecule=water_box, settings=md_settings)
md_results = md_job.run()
if not md_results.ok():
    raise RuntimeError("Toy MD run did not finish successfully.")

rkf_path = md_results.rkfpath()
print(f"Trajectory file: {rkf_path}")

view(water_box, direction="tilt_z", width=300, height=300, picture_path="picture1.png")


# ### Extract First Five Frames
#
# Use `Trajectory` to iterate over frames in `ams.rkf`. Each frame is returned as a PLAMS `Molecule`.
#

trajectory = Trajectory(rkf_path)

for frame_number, frame_molecule in enumerate(trajectory, start=1):
    print(f"Frame {frame_number}")
    print(frame_molecule)
    if frame_number == 5:
        break

view(frame_molecule, direction="tilt_z", width=300, height=300, picture_path="picture2.png")


# ### Extract Specific Frame
#
# Use `get_history_molecule` to load one selected frame directly.
#

n_entries = int(md_results.readrkf("History", "nEntries"))
selected_frame = min(10, n_entries)
selected_molecule = md_results.get_history_molecule(selected_frame)

print(f"Extracting frame {selected_frame} of {n_entries}")
print(selected_molecule)
view(selected_molecule, direction="tilt_z", width=300, height=300, picture_path="picture3.png")