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

# ## Convert from XYZ/OUTCAR to RKF
# This example starts from an external multi-frame XYZ trajectory, converts it to `ams.rkf`, and enriches it with molecule and bond information by running a replay job.
#

# ### Initial Imports
#

import os
import subprocess
from pathlib import Path

from scm.plams import AMSJob, Settings, delete_job, file_to_traj, traj_to_rkf


# ### Define XYZ trajectory input
#

# We first write a very short example .xyz file:

xyz_trajectory = """6
frame 1
O        2.01504900       2.16272700       2.99552700
H        2.96307600       2.41195100       2.99244800
H        1.42963800       2.95106300       2.98957900
O        0.95320300       1.47810200       1.19941600
H        1.00821200       2.41101200       0.90355900
H        1.76505900       0.97420300       0.97316100
6
frame 2
O        2.01505786       2.16226811       2.99837195
H        2.95676994       2.39860472       2.98680480
H        1.42248605       2.96034695       2.99335021
O        0.95325699       1.47877442       1.19649020
H        1.01375745       2.41610845       0.89913436
H        1.77197403       0.96978002       0.98074075
6
frame 3
O        2.01512004       2.16202221       3.00146369
H        2.94874909       2.38589811       2.98026036
H        1.41907239       2.96804981       2.99718816
O        0.95316789       1.47931163       1.19345976
H        1.02092464       2.41742729       0.89396582
H        1.77666876       0.96884156       0.98764296
6
frame 4
O        2.01516130       2.16199995       3.00479257
H        2.93983749       2.37412231       2.97284825
H        1.41975474       2.97371643       3.00104071
O        0.95294917       1.47971792       1.19032629
H        1.02967864       2.41481091       0.88814023
H        1.77896035       0.97147237       0.99392690"""

xyz_input = Path("water_frames.xyz")
xyz_input.write_text(xyz_trajectory + "\n")
# ### Convert to XYZ to RKF

# We then convert the .xyz file to a .rkf via an intermediate .traj file:

rkf_output = "water_frames.rkf"

traj_temp = Path("water_frames.traj")
file_to_traj(str(xyz_input), str(traj_temp))
traj_to_rkf(str(traj_temp), rkf_output, task="moleculardynamics", timestep=0.5)


# ### Run Replay Job

# Finally, we add molecule and bonding information by running a replay job:

replay_settings = Settings()
replay_settings.input.ams.task = "replay"
replay_settings.input.lennardjones
replay_settings.input.ams.replay.file = os.path.abspath(rkf_output)
replay_settings.input.ams.properties.molecules = "yes"
replay_settings.input.ams.properties.bondorders = "yes"
replay_settings.runscript.nproc = 1

replay_job = AMSJob(settings=replay_settings, name="replay_for_bond_guessing")
replay_results = replay_job.run()
if not replay_results.ok():
    raise RuntimeError("Replay job failed")

cpkf = Path(os.path.expandvars("$AMSBIN/cpkf"))
if not cpkf.exists():
    raise FileNotFoundError(f"Could not find cpkf at {cpkf}")

subprocess.run(
    ["sh", str(cpkf), replay_results.rkfpath(), rkf_output, "History", "Molecules"],
    check=True,
)
delete_job(replay_job)
traj_temp.unlink(missing_ok=True)