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

# ## i-PI path integral MD with AMS
#
# This example prepares the files needed to couple [i-PI](https://ipi-code.org/) to AMS through the ASE socket interface.
#
# i-PI can be used to run, for example, path integral molecular dynamics. The example shows how to run thermostatted ring polymer molecular dynamics for a water molecule, together with the UFF force field.
#
# For more information about i-PI, refer to the i-PI documentation and examples.
#
# Note that the provided setup uses a Unix socket and therefore runs on Linux.

# ### Create Input File for i-PI

# First, write the input for i-PI to an xml file.

from pathlib import Path

Path("input.xml").write_text(
    """\
<simulation mode="md" verbosity="medium">
   <output prefix='simulation'>
      <properties stride='10' filename='out'> [ step, time{picosecond}, conserved{electronvolt}, temperature{kelvin}, pressure_cv{megapascal} ] </properties>
      <trajectory filename='pos' stride='10'> positions{angstrom} </trajectory>
      <trajectory filename='xc' stride='10'> x_centroid{angstrom} </trajectory>
      <trajectory filename='vc' stride='10'> v_centroid </trajectory>
      <checkpoint stride='4000'/>
   </output>
   <total_steps> 40 </total_steps>
   <!-- <total_time> 840000000 </total_time> -->
   <prng><seed> 31415 </seed></prng>
   <ffsocket name="ase" mode="unix">
      <address> driver-irpmd-16 </address>
   </ffsocket>
   <system>
      <initialize nbeads='8'>
         <file mode='xyz' units='angstrom'> firstframe.xyz </file>
         <velocities mode="thermal" units='kelvin'> 300 </velocities>
         <cell units='angstrom' mode='abc'> [ 12.412, 12.412, 12.412 ] </cell>
      </initialize>
      <forces><force forcefield="ase"> </force></forces>
      <motion mode='dynamics'>
         <dynamics mode='nvt'>
           <timestep units='femtosecond'> 0.25 </timestep>
           <thermostat mode='pile_g'>
             <tau units='femtosecond'> 100 </tau>
             <pile_lambda> 0.5 </pile_lambda>
           </thermostat> 
         </dynamics>
      </motion>
      <ensemble>
         <temperature units='kelvin'> 300 </temperature>
      </ensemble>
   </system>
</simulation>
    """
)
# Create a simple xyz file with a water molecule:

Path("firstframe.xyz").write_text(
    """\
3

O -0.691870 1.324715 0.212227
H -1.523797 1.656370 -0.135379
H -0.797941 0.389809 0.486593
VEC1 12.412 0. 0.
VEC2 0. 12.412 0.
VEC3 0. 0. 12.412
"""
)
# Create script to launch i-PI server (the variable `ipi_exec` can be adjusted to your installation).

ipi_exec = "i-pi"
Path("run-server.sh").write_text(
    f"""\
#!/usr/bin/env bash
# Example launcher for i-PI server 
{ipi_exec} input.xml
"""
)
# ### Launch the i-PI Server

# Now in a separate window, run `sh run-server.sh`. Once the server is up and running we can proceed with the AMS calculation.
#
# i-PI is external to AMS and must be installed separately. The provided setup uses a Unix socket and therefore runs on Linux and Mac.

print("Start i-PI server by running in a new terminal: 'sh run-server.sh'")


# ### Run AMS Calculator

from ase.calculators.socketio import SocketClient
from ase.io import read
from scm.plams import Settings, from_smiles
from scm.plams.interfaces.adfsuite.ase_calculator import AMSCalculator


atoms = read("firstframe.xyz")

sett = Settings()
sett.input.ams.Task = "SinglePoint"
sett.input.ams.Properties.Gradients = "True"
sett.input.ams.Properties.StressTensor = "False"
sett.input.forcefield.type = "UFF"
sett.runscript.nproc = 1

with AMSCalculator(settings=sett, amsworker=True) as calc:
    atoms.calc = calc
    client = SocketClient(unixsocket="driver-irpmd-16")
    client.run(atoms, use_stress=False)