#!/usr/bin/env python

# ----------------------------------------------------------------------
# LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
# http://lammps.sandia.gov, Sandia National Laboratories
# Steve Plimpton, sjplimp@sandia.gov
# ----------------------------------------------------------------------

# Syntax: nwchem_wrap.py file/zmq ao/pw input_template
# file/zmg = messaging mode, must match LAMMPS messaging mode
# ao/pw = basis set mode, selects between atom-centered and plane-wave
#         the input_template file must correspond to the appropriate basis set mode:
#         the "ao" mode supports the scf and dft modules in NWChem,
#         the "pw" mode supports the nwpw module.
# input_template = NWChem input file used as template, must include a
#                  "geometry" block with the atoms in the simulation, dummy
#                  xyz coordinates should be included (but are not used).
#                  Atom ordering must match LAMMPS input.

# wrapper on NWChem
#   receives message with list of coords
#   creates NWChem inputs
#   invokes NWChem to calculate self-consistent energy of that config
#   reads NWChem outputs
#   sends message with energy, forces, pressure to client

from __future__ import print_function
import sys

version = sys.version_info[0]
if version == 3:
  sys.exit("The CSlib python wrapper does not yet support python 3")

import subprocess
import re
import os
import shutil
from cslib import CSlib

# comment out 2nd line once 1st line is correct for your system

nwchemcmd = "mpirun -np 1 /usr/bin/nwchem nwchem_lammps.nw > "
nwchemcmd = "touch tmp"

# enums matching FixClientMD class in LAMMPS

SETUP,STEP = range(1,2+1)
DIM,PERIODICITY,ORIGIN,BOX,NATOMS,NTYPES,TYPES,COORDS,UNITS,CHARGE = range(1,10+1)
FORCES,ENERGY,VIRIAL,ERROR = range(1,4+1)

# -------------------------------------
# functions

# error message and exit

def error(txt):
  print("ERROR:",txt)
  sys.exit(1)

# -------------------------------------
# read initial input file to setup problem
# return natoms

def nwchem_setup_ao(input):

  template = open(input,'r')

  geometry_block = False
  natoms = 0

  while True:
    line = template.readline()
    if not line: break

    if geometry_block and re.search("end",line):
      geometry_block = False
    if geometry_block and not re.match("#",line) :
      natoms += 1
    if re.search("geometry",line):
      geometry_block = True

  return natoms

# -------------------------------------
# write a new input file for NWChem
# assumes the NWChem input geometry is to be specified in angstroms

def nwchem_input_write_ao(input,coords):

  template = open(input,'r')
  new_input = open("nwchem_lammps.nw",'w')

  geometry_block = False
  i = 0

  while True:
    line = template.readline()
    if not line: break

    if geometry_block and not re.match("#",line) and re.search("end",line):
      geometry_block = False
      if os.path.exists("nwchem_lammps.movecs"):
        # The below is hacky, but one of these lines will be ignored
        # by NWChem depending on if the input file is for scf/dft.
        append = "\nscf\n vectors input nwchem_lammps.movecs\nend\n"
        append2 = "\ndft\n vectors input nwchem_lammps.movecs\nend\n"
        line = line + append + append2

    if geometry_block and not re.match("#",line):
      x = coords[3*i+0]
      y = coords[3*i+1]
      z = coords[3*i+2]
      coord_string = "  %g %g %g \n" % (x,y,z)
      atom_string = line.split()[0]
      line = atom_string + coord_string
      i += 1

    if (not re.match("#",line)) and re.search("geometry",line):
      geometry_block = True
      line = "geometry units angstrom noautosym\n"

    print(line,file=new_input,end='')

  new_input.close()

# -------------------------------------
# read a NWChem output nwchem_lammps.out file

def nwchem_read_ao(natoms, log):

  nwchem_output = open(log, 'r')
  energy_pattern = r"Total \w+ energy"
  gradient_pattern = "x          y          z           x          y          z"

  eout = 0.0
  fout = []

  while True:
    line = nwchem_output.readline()
    if not line: break

    # pattern match for energy
    if re.search(energy_pattern,line):
      eout = float(line.split()[4])

    # pattern match for forces
    if re.search(gradient_pattern, line):
      for i in range(natoms):
        line = nwchem_output.readline()
        forces = line.split()
        fout += [float(forces[5]), float(forces[6]), float(forces[7])]

  # convert units
  hartree2eV = 27.21138602
  bohr2angstrom = 0.52917721092
  eout = eout * hartree2eV
  fout = [i * hartree2eV/bohr2angstrom for i in fout]
  print(eout)

  return eout,fout

# -------------------------------------
# read initial planewave input file to setup problem
# return natoms,box
def nwchem_setup_pw(input):

  template = open(input,'r')

  geometry_block = False
  system_block = False
  coord_pattern = r"^\s*\w{1,2}(?:\s+-?(?:\d+.?\d*|\d*.?\d+)){3}"
  natoms = 0
  box = []

  while True:
    line = template.readline()
    if not line: break

    if geometry_block and re.search("system crystal",line):
      system_block = True
      for i in range(3):
        line = template.readline()
        line = re.sub(r'd|D', 'e', line)
        box += [float(line.split()[1])]

    if geometry_block and not system_block and re.match("#",line) and re.search("end",line):
      geometry_block = False

    if system_block and re.search("end",line):
      system_block = False

    if geometry_block and not re.match("#",line) and re.search(coord_pattern,line):
      natoms += 1

    if re.search("geometry",line) and not re.match("#",line):
      geometry_block = True

  return natoms,box

# -------------------------------------
# write a new planewave input file for NWChem
# assumes the NWChem input geometry is to be specified fractional coordinates

def nwchem_input_write_pw(input,coords,box):

  template = open(input,'r')
  new_input = open("nwchem_lammps.nw",'w')

  writing_atoms = False
  geometry_block = False
  system_block = False
  coord_pattern = r"^\s*\w{1,2}(?:\s+-?(?:\d+.?\d*|\d*.?\d+)){3}"
  i = 0

  while True:
    line = template.readline()
    if not line: break

    if geometry_block and re.search("system crystal",line):
      system_block = True

    if geometry_block and not system_block and not re.match("#",line) and re.search("end",line):
      geometry_block = False
      if os.path.exists("nwchem_lammps.movecs"):
        append = "\nnwpw\n vectors input nwchem_lammps.movecs\nend\n"
        line = line + append

    if system_block and re.search("end",line):
      system_block = False

    if geometry_block and not re.match("#",line) and re.search(coord_pattern,line):
      x = coords[3*i+0] / box[0]
      y = coords[3*i+1] / box[1]
      z = coords[3*i+2] / box[2]
      coord_string = "  %g %g %g \n" % (x,y,z)
      atom_string = line.split()[0]
      line = atom_string + coord_string
      i += 1

    if re.search("geometry",line) and not re.match("#",line):
      geometry_block = True

    print(line,file=new_input,end='')

  new_input.close()

# -------------------------------------
# read a NWChem output nwchem_lammps.out file for planewave calculation

def nwchem_read_pw(log):
  nw_output = open(log, 'r')

  eout = 0.0
  sout = []
  fout = []
  reading_forces = False

  while True:
    line = nw_output.readline()
    if not line: break

    # pattern match for energy
    if re.search("PSPW energy",line):
      eout = float(line.split()[4])

    # pattern match for forces
    if re.search("C\.O\.M", line):
      reading_forces = False
    if reading_forces:
      forces = line.split()
      fout += [float(forces[3]), float(forces[4]), float(forces[5])]
    if re.search("Ion Forces",line):
      reading_forces = True

    # pattern match for stress
    if re.search("=== total gradient ===",line):
      stensor = []
      for i in range(3):
        line = nw_output.readline()
        line = line.replace("S =","   ")
        stress = line.split()
        stensor += [float(stress[1]), float(stress[2]), float(stress[3])]
      sxx = stensor[0]
      syy = stensor[4]
      szz = stensor[8]
      sxy = 0.5 * (float(stensor[1]) + float(stensor[3]))
      sxz = 0.5 * (stensor[2] + stensor[6])
      syz = 0.5 * (stensor[5] + stensor[7])
      sout = [sxx,syy,szz,sxy,sxz,syz]

  # convert units
  hartree2eV = 27.21138602
  bohr2angstrom = 0.52917721092
  austress2bar = 294210156.97
  eout = eout * hartree2eV
  fout = [i * hartree2eV/bohr2angstrom for i in fout]
  sout = [i * austress2bar for i in sout]

  return eout,fout,sout

# -------------------------------------
# main program

# command-line args
#
if len(sys.argv) != 4:
  print("Syntax: python nwchem_wrap.py file/zmq ao/pw input_template")
  sys.exit(1)

comm_mode = sys.argv[1]
basis_type = sys.argv[2]
input_template = sys.argv[3]

if comm_mode == "file": cs = CSlib(1,comm_mode,"tmp.couple",None)
elif comm_mode == "zmq": cs = CSlib(1,comm_mode,"*:5555",None)
else:
  print("Syntax: python nwchem_wrap.py file/zmq")
  sys.exit(1)


natoms = 0
box = []
if basis_type == "ao":
  natoms = nwchem_setup_ao(input_template)
elif basis_type == "pw":
  natoms,box = nwchem_setup_pw(input_template)

# initial message for AIMD protocol

msgID,nfield,fieldID,fieldtype,fieldlen = cs.recv()
if msgID != 0: error("Bad initial client/server handshake")
protocol = cs.unpack_string(1)
if protocol != "md": error("Mismatch in client/server protocol")
cs.send(0,0)

# endless server loop

i = 0
if not os.path.exists("nwchem_logs"):
  os.mkdir("nwchem_logs")

while 1:

  # recv message from client
  # msgID = 0 = all-done message

  msgID,nfield,fieldID,fieldtype,fieldlen = cs.recv()
  if msgID < 0: break

  # SETUP receive at beginning of each run
  # required fields: DIM, PERIODICITY, ORIGIN, BOX,
  #                  NATOMS, COORDS
  # optional fields: others in enum above, but NWChem ignores them

  if msgID == SETUP:

    origin = []
    box_lmp = []
    natoms_recv = ntypes_recv = 0
    types = []
    coords = []

    for field in fieldID:
      if field == DIM:
        dim = cs.unpack_int(DIM)
        if dim != 3: error("NWChem only performs 3d simulations")
      elif field == PERIODICITY:
        periodicity = cs.unpack(PERIODICITY,1)
        if basis_type == "ao":
          if periodicity[0] or periodicity[1] or periodicity[2]:
            error("NWChem AO basis wrapper only currently supports fully aperiodic systems")
        elif basis_type == "pw":
          if not periodicity[0] or not periodicity[1] or not periodicity[2]:
            error("NWChem PW basis wrapper only currently supports fully periodic systems")
      elif field == ORIGIN:
        origin = cs.unpack(ORIGIN,1)
      elif field == BOX:
        box_lmp = cs.unpack(BOX,1)
        if (basis_type == "pw"):
          if (box[0] != box_lmp[0] or box[1] != box_lmp[4] or box[2] != box_lmp[8]):
            error("NWChem wrapper mismatch in box dimensions")
      elif field == NATOMS:
        natoms_recv = cs.unpack_int(NATOMS)
        if natoms != natoms_recv:
          error("NWChem wrapper mismatch in number of atoms")
      elif field == COORDS:
        coords = cs.unpack(COORDS,1)

    if not origin or not box_lmp or not natoms or not coords:
      error("Required NWChem wrapper setup field not received");

  # STEP receive at each timestep of run or minimization
  # required fields: COORDS
  # optional fields: ORIGIN, BOX

  elif msgID == STEP:

    coords = []

    for field in fieldID:
      if field == COORDS:
        coords = cs.unpack(COORDS,1)

    if not coords: error("Required NWChem wrapper step field not received");

  else: error("NWChem wrapper received unrecognized message")

  # unpack coords from client
  # create NWChem input

  if basis_type == "ao":
    nwchem_input_write_ao(input_template,coords)
  elif basis_type == "pw":
    nwchem_input_write_pw(input_template,coords,box)

  # invoke NWChem

  i += 1
  log = "nwchem_lammps.out"
  archive = "nwchem_logs/nwchem_lammps" + str(i) + ".out"
  cmd = nwchemcmd + log
  print("\nLaunching NWChem ...")
  print(cmd)
  subprocess.check_output(cmd,stderr=subprocess.STDOUT,shell=True)

  shutil.copyfile(log,archive)

  # process NWChem output

  if basis_type == "ao":
    energy,forces = nwchem_read_ao(natoms,log)
    virial = [0,0,0,0,0,0]
  elif basis_type == "pw":
    energy,forces,virial = nwchem_read_pw(log)

  # return forces, energy to client
  cs.send(msgID,3)
  cs.pack(FORCES,4,3*natoms,forces)
  cs.pack_double(ENERGY,energy)
  cs.pack(VIRIAL,4,6,virial)

# final reply to client

cs.send(0,0)

# clean-up

del cs