#!/usr/bin/env python2.7
 
import argparse, numpy, subprocess, re, sys
import os
 
# Parse command-line arguments
 
parser = argparse.ArgumentParser(description = """
Create a set of NWChem input files for the Ar2 dimer, run NWChem on them, and
parse the output for interaction energy.
""")
parser.add_argument("-m", "--min", dest="minR", type=float, default=2.5,
                    help="minimum Ar separation in A")
parser.add_argument("-M", "--max", dest="maxR", type=float, default=5.5,
                    help="maximum Ar separation in A")
parser.add_argument("-d", "--delta", dest="deltaR", type=float, default=0.5,
                    help="step size for Ar separation in A")
parser.add_argument("-b", "--basis", dest="basis", type=str, default="aug-cc-pvtz",
                    help="basis set used in calculation")
parser.add_argument("-o", "--output", dest="output", type=str, default="output.txt",
                    help="output text file")
 
args = parser.parse_args()
 
nwchem_dir = os.environ["NWCHEM_INSTALL"]
username = os.environ["USER"]
scratch = os.environ["SCRATCH"]
nproc = "1"
 
 
# This command uses the minR, maxR and deltaR arguments to create an array.
positions = numpy.arange(args.minR, args.maxR + 0.00001, args.deltaR)
# If you need a non-uniform set of points, you can define them 
# using the numpy.array command as follows. If you wish to use this approach,
# un-comment the following line and comment out the one above.
#positions = numpy.array([ 3.0, 3.5, 3.75, 4.0, 4.25, 4.5, 4.75, 5.0, 5.5, 6.0 ])
print "Basis set ",args.basis
print "positions = ", positions
print "Output to : ",args.output
output = open(args.output,'w')
 
 
# Conversion factor:
au2cm = 219474.631371
 
# Construct input files
 
input_file_template = """
title "Ar dimer BSSE corrected DFT interaction energy"
 
scratch_dir {scratch}
 
geometry units angstrom "Ar+Ar" noautoz
 Ar1 0 0 0
 Ar2 0 0 {R}
end
 
geometry units angstrom "Ar+ghost" noautoz
 Ar1  0 0 0
 BqAr 0 0 {R}
end
 
basis
 Ar1   library    {basis}
 Ar2   library    {basis}
 BqAr  library Ar {basis}
end
 
dft
  xc pbe0
  direct
  disp
  iterations 60
end
 
set geometry "Ar+Ar"
task dft
unset geometry "Ar+Ar"
 
dft; vectors hcore; end
 
set geometry "Ar+ghost"
task dft
unset geometry "Ar+ghost"
"""
 
# Regular expressions to find the right lines in the file
 
dft_re = re.compile("Total DFT energy =\s+(-?[0-9]+.[0-9]+)")
 
 
# Dictionary to store the energies we find
 
energies = {}
 
# Main loop through different position values
 
for ar_position in positions:
    print "Current position = ",ar_position
    input_file_name = "in-R{0:5.3f}.nw".format(ar_position)
    output_file_name = "out-R{0:5.3f}.out".format(ar_position)
 
    # Write input file
    with open(input_file_name,"w") as INP:
        INP.write(input_file_template.format(R=ar_position, basis=args.basis, 
        user=username, scratch=scratch))
 
    # Run NWChem
    output_file = open(output_file_name, "w")
    subprocess.call(["mpirun","-np",nproc,"nwchem", input_file_name],
                                                stdout=output_file, stderr=subprocess.STDOUT)
    output_file.close()
 
    # Read output file back in
 
    output_file = open(output_file_name, "r")
    dft_energies = []
    for line in output_file:
        dft_match = dft_re.search(line)
        if dft_match:
            dft_energies.append(float(dft_match.group(1)))
    output_file.close()
 
    if not (len(dft_energies) == 2):
        print "Apparent error in output file! {0:d} DFT energies found.".format(len(dft_energies))
        sys.exit(1)
 
    print "dft_energies ", dft_energies     
    energies[ar_position] = dft_energies
    dft_dimer   = energies[ar_position][0]
    dft_monomer = energies[ar_position][1]
    dft_interaction = dft_dimer - 2*dft_monomer
    print "Eint[DFT] = ",dft_interaction*au2cm
 
# prepare plots
 
dft_dimer = numpy.array([energies[p][0] for p in positions])
dft_monomer = numpy.array([energies[p][1] for p in positions])
 
dft_interaction = dft_dimer - 2*dft_monomer
 
# convert to cm-1
dft_interaction = au2cm*dft_interaction
 
output.write(('%10s '*4 + '\n') % 
             ("   R    ","DFT mono","DFT dimr","DFT int "))
 
for i in range(len(positions)):
    output.write(('%10.6f '*4 + '\n') %
    (positions[i],
    dft_monomer[i], dft_dimer[i], dft_interaction[i]))
 
output.close()