#!/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=1.,
                    help="minimum Ar separation in A")
parser.add_argument("-M", "--max", dest="maxR", type=float, default=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 = "2"
 
positions = numpy.arange(args.minR, args.maxR + 0.00001, args.deltaR)
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 MP2 interaction energy"
 
scratch_dir {scratch}
 
geometry units angstrom "Ar+Ar" noautoz noautosym
 Ar1 0 0 0
 Ar2 0 0 {R}
end
 
geometry units angstrom "Ar+ghost" noautoz noautosym
 Ar1 0 0 0
 Bq2 0 0 {R}
end
 
basis
 Ar1 library    {basis}
 Ar2 library    {basis}
 Bq2 library Ar {basis}
end
 
set geometry "Ar+Ar"
task mp2
unset geometry "Ar+Ar"
 
scf; vectors atomic; end
 
set geometry "Ar+ghost"
task mp2
unset geometry "Ar+ghost"
"""
 
# Regular expressions to find the right lines in the file
 
scf_re = re.compile("SCF energy\s+(-?[0-9]+.[0-9]+)")
mp2_re = re.compile("Total MP2 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")
    scf_energies = []
    mp2_energies = []
    for line in output_file:
        scf_match = scf_re.search(line)
        if scf_match:
            scf_energies.append(float(scf_match.group(1)))
        else:
            mp2_match = mp2_re.search(line)
            if mp2_match:
                mp2_energies.append(float(mp2_match.group(1)))
    output_file.close()
 
    if not (len(scf_energies) == 2 and len(mp2_energies) == 2):
        print "Apparent error in output file! {:d} SCF energies and {:d} MP2 energies found.".format(
            len(scf_energies), len(mp2_energies))
        sys.exit(1)
 
    energies[ar_position] = scf_energies + mp2_energies
    scf_dimer   = energies[ar_position][0]
    scf_monomer = energies[ar_position][1]
    mp2_dimer   = energies[ar_position][2]
    mp2_monomer = energies[ar_position][3]
    scf_interaction = scf_dimer - 2*scf_monomer
    mp2_interaction = mp2_dimer - 2*mp2_monomer
    print "Eint[SCF] = ",scf_interaction*au2cm
    print "Eint[MP2] = ",mp2_interaction*au2cm
 
# prepare plots
 
scf_dimer = numpy.array([energies[p][0] for p in positions])
scf_monomer = numpy.array([energies[p][1] for p in positions])
mp2_dimer = numpy.array([energies[p][2] for p in positions])
mp2_monomer = numpy.array([energies[p][3] for p in positions])
 
scf_interaction = scf_dimer - 2*scf_monomer
mp2_interaction = mp2_dimer - 2*mp2_monomer
 
# convert to cm-1
scf_interaction = au2cm*scf_interaction
mp2_interaction = au2cm*mp2_interaction
 
output.write(('%10s '*7 + '\n') % 
             ("   R    ","SCF mono","SCF dimr",
              "SCF int ","MP2 mono","MP2 dimr","MP2 int "))
 
for i in range(len(positions)):
    output.write(('%10.6f '*7 + '\n') %
    (positions[i],
    scf_monomer[i], scf_dimer[i], scf_interaction[i],
    mp2_monomer[i], mp2_dimer[i], mp2_interaction[i]))
 
output.close()