#!/usr/bin/env python3
"""
This Test Driver computes the total potential energy and forces acting on a
triclinic periodic 3D box of atoms, supplied by the user in the form of an
extended xyz file which lies in the Test's directory.  The extended xyz file
must be of the following form:
Natoms
Lattice=" Ax Ay Az Bx By Bz Cx Cy Cz "
Element positionx positiony positionz
Element positionx positiony positionz
.
.
.
where A, B, and C are the simulation box edge vectors (periodic boundary
conditions).  Element is an integer which species the atomic species code used
by LAMMPS; the order these species codes should start from 1 and corresponds to
the chemical species according to the order of elements specified in
pipeline.stdin.tpl:  the first species given in pipeline.stdin.tpl should
correspond to species 1, the second to species 2, etc.  Note that all entries
in the xyz file are separated by single spaces.

Author: Daniel S. Karls, University of Minnesota (karl0100 |AT| umn DOT edu)
"""
import os
import itertools
import subprocess
import shutil
import re
import json
from collections import namedtuple, OrderedDict

import numpy

import xyz_utils


def stripquotes(matchobj):
    return matchobj.group(1)


def run_lammps(infile, outfile):
    """Run LAMMPS with given input file and write the output to outfile"""
    with open(outfile, "w") as outfl:
        try:
            lammps_process = subprocess.check_call(
                ["lammps", "-in", infile], shell=False, stdout=outfl
            )
        except subprocess.CalledProcessError:
            extrainfo = ""
            try:
                with open("log.lammps") as f:
                    extrainfo = f.read()
            except IOError:
                extrainfo = "no log file"
            raise Exception("LAMMPS did not exit properly:\n" + extrainfo)


def get_isolated_atom_energy(model, element, mass):
    ISOLATED_ATOM_LAMMPS_INPUT_TEMPLATE = open(
        os.path.join(THIS_DIR, "isolated_atom.lammps.tpl")
    ).read()
    templated_input = "isolated_atom.lammps." + element + ".in"
    lammps_output = "isolated_atom.lammps." + element + ".out"
    with open(templated_input, "w") as in_file:
        in_file.write(
            ISOLATED_ATOM_LAMMPS_INPUT_TEMPLATE.format(
                modelname=model, symbol=element, mass=mass
            )
        )

    run_lammps(templated_input, lammps_output)

    # Get potential energy
    with open(lammps_output) as outfl:
        output = outfl.read()
    try:
        energy = re.search("Isolated atom energy: ([0-9e.\-]+) eV", output).group(1)
        energy = float(energy)
    except AttributeError:
        raise Exception("Failed to find the potential energy")

    return energy


# Read elements, masses, and model from stdin
elementinput = input("Elements: ")
print(elementinput)
elements = elementinput.split(" ")
numelements = len(elements)
if numelements < 1:
    raise RuntimeError("ERROR: No species were listed! Exiting...")
massesinput = input("Masses in g/mol (must match order of elements above): ")
masses = massesinput.split(" ")
print(masses)
model = input("Model: ")
print(model)
xyz_src = input("XYZ file: ")
print(xyz_src)
print("")

# Some directories we need
THIS_DIR = os.path.dirname(__file__)
MAIN_LAMMPS_INPUT_TEMPLATE = open(os.path.join(THIS_DIR, "main.lammps.tpl")).read()
INDIR = os.path.join("output", "lammps_inputs")
OUTDIR = os.path.join("output", "lammps_output_log")
DUMPDIR = os.path.join("output", "lammps_dump")

# Ensure the directories we need are created
try:
    os.makedirs(INDIR)
except OSError:
    pass
try:
    os.makedirs(OUTDIR)
except OSError:
    pass
try:
    os.makedirs(DUMPDIR)
except OSError:
    pass

# Files LAMMPS uses or produces
infile = os.path.join(INDIR, "in.lammps")
outfile = os.path.join(OUTDIR, "lammps.log")
dumpfile = os.path.join(DUMPDIR, "lammps.dump")
lammps_xyzfile = os.path.join(INDIR, "triclinicpbc_lammps.xyz")

# Create results dictionary
results = OrderedDict()
results[
    "property-id"
] = "tag:staff@noreply.openkim.org,2014-04-15:property/configuration-nonorthogonal-periodic-3d-cell-fixed-particles-fixed"
results["instance-id"] = 1

### REGULAR EXPRESSIONS FOR MATCHING LAMMPS OUTPUT

# Finds potential energy
POTENG_MATCH = re.compile(
    r"""
        v_pe_metal      # MAGIC Word
        .*$\n           # until end of line
        \s*             # possible leading whitespace
        ([0-9e.\-]+)    # out float, grouped
        """,
    flags=re.VERBOSE | re.MULTILINE,
)

# Finds number of atoms
NATOM_MATCH = re.compile(
    r"""
        ITEM:\ NUMBER\ OF\ ATOMS  # MAGIC WORDS
        .*$\n                     # until end of line
        \s*                       # possible leading whitespace
        (\d+)                     # our count, grouped
        """,
    flags=re.VERBOSE | re.MULTILINE,
)

# Finds the BOX BOUNDS lines in the LAMMPS dumpfile
# BOXBOUNDS_LINES = re.compile(r"""
#        ITEM:\ BOX\ BOUNDS\        #magic words
#        .*?\n               #till end of line, nongreedy dot
#        (.*)                #rest of file
#        """,flags=re.VERBOSE|re.DOTALL)

# Finds the ATOMS lines in the LAMMPS dumpfile
ATOMS_LINES = re.compile(
    r"""
        ITEM:\ ATOMS        #magic words
        .*?\n               #till end of line, nongreedy dot
        (.*)                #rest of file
        """,
    flags=re.VERBOSE | re.DOTALL,
)

# Read xyz file, convert the cell to LAMMPS' convention and create a
# dump file to read in
# NOTE: using a dump file rather than a data file means that it's not
#       sensitive to the specific atom_style in the event that we're
#       running with an SM, e.g. if the atom_style of the SM is
#       'charge', charges of zero will simply be assigned to each atom
#       since they're not specified in the dumpfile
numatoms, orig_cell, species, orig_pos = xyz_utils.read_xyz(xyz_src)
new_cell, new_pos = xyz_utils.convert_xyz_to_lammps_convention(
    numatoms, orig_cell, orig_pos
)

# Curate species from xyz file
species_codes = []
for el in species:
    try:
        species_codes.append(int(el))
    except ValueError:
        # Try to map the element string given to an integer code based on
        # the elements the Test inputted
        if el in elements:
            species_codes.append(elements.index(el) + 1)
        else:
            raise RuntimeError(
                "Found species in xyz file {} that was not "
                "given in the elements inputted".format(xyz_src)
            )

# Calculate isolated atomic energies for each species
isolated_atom_energies = {}
for ind, symbol in enumerate(elements):
    isolated_atom_energies[symbol] = get_isolated_atom_energy(
        model, symbol, masses[ind]
    )
    print(
        "Isolated atomic energy for species {}: {} eV".format(
            symbol, isolated_atom_energies[symbol]
        )
    )
print("")

# Sum together isolated atomic energies of each atom in the xyz file
isolated_atom_energies_sum = sum(
    [isolated_atom_energies[elements[code - 1]] for code in species_codes]
)
print("Sum of isolated atomic energies: {} eV".format(isolated_atom_energies_sum))
print("")

xyz_utils.write_lammps_xyz_dump(
    numatoms, new_cell, species_codes, new_pos, lammps_xyzfile
)

# Extract LAMMPS description for new cell
xlo = 0.0
ylo = 0.0
zlo = 0.0
xhi, _, _, xy, yhi, _, xz, yz, zhi = new_cell

Rinv = xyz_utils.get_inverse_rotation_matrix(orig_cell, new_cell)
Rinv11, Rinv21, Rinv31, Rinv12, Rinv22, Rinv32, Rinv13, Rinv23, Rinv33 = Rinv

# Set atomic masses
mass_string = ""
for type_count in range(numelements):
    mass_string += "variable mass{}_converted equal {}*${{_u_mass}}\n".format(
        type_count + 1, masses[type_count]
    )
    mass_string += "mass {} ${{mass{}_converted}}\n".format(
        type_count + 1, type_count + 1
    )

# Create the LAMMPS input file
with open(infile, "w") as in_file:
    in_file.write(
        MAIN_LAMMPS_INPUT_TEMPLATE.format(
            modelname=model,
            symbol=elementinput,
            numelements=numelements,
            box_xlo=xlo,
            box_xhi=xhi,
            box_ylo=ylo,
            box_yhi=yhi,
            box_zlo=zlo,
            box_zhi=zhi,
            box_xy=xy,
            box_xz=xz,
            box_yz=yz,
            lammps_xyzfile=lammps_xyzfile,
            set_masses=mass_string,
            dumpfile=dumpfile,
        )
    )

# Run LAMMPS
run_lammps(infile, outfile)

### Now process the output and dumpfile for relevant information
# Get potential energy
with open(outfile) as outfl:
    output = outfl.read()
try:
    poteng = {}
    total_energy = POTENG_MATCH.search(output).group(1)
except AttributeError:
    raise Exception("Failed to find the potential energy")

total_energy = float(total_energy)
print("Raw total energy: {} eV".format(total_energy))

total_energy = total_energy - isolated_atom_energies_sum
print(
    "Total energy with isolated atomic energies subtracted: {} eV "
    "".format(total_energy)
)
print("")
print("")

poteng["source-value"] = total_energy
poteng["source-unit"] = "eV"

# Get configuration
with open(dumpfile) as dumpfl:
    dump = dumpfl.read()
# Find the number of atoms
try:
    natoms = int(NATOM_MATCH.search(dump).group(1))
    # results.setdefault(triclinic_file,{})['natoms'] = natoms
except AttributeError:
    raise Exception("Failed to find the number of atoms")

# Process the rest of the dump file, the atom positions, etc
try:
    atomslines = ATOMS_LINES.search(dump).group(1)
except AttributeError:
    raise Exception("Failed to find ATOMS block")


# Set lattice vectors to original orientation
unrelaxedcellvec1 = {}
unrelaxedcellvec2 = {}
unrelaxedcellvec3 = {}
unrelaxedcellvec1["source-unit"] = "angstrom"
unrelaxedcellvec2["source-unit"] = "angstrom"
unrelaxedcellvec3["source-unit"] = "angstrom"
unrelaxedcellvec1["source-value"] = [orig_cell[0], orig_cell[1], orig_cell[2]]
unrelaxedcellvec2["source-value"] = [orig_cell[3], orig_cell[4], orig_cell[5]]
unrelaxedcellvec3["source-value"] = [orig_cell[6], orig_cell[7], orig_cell[8]]

nanflag = 0
# Check poteng for NaN
if numpy.isnan(float(poteng["source-value"])):
    nanflag = 1

atom_count = 0
configspecies = {}
configspecies["source-value"] = []
unrelaxedconfigpos = {}
unrelaxedconfigpos["source-value"] = []
unrelaxedconfigpos["source-unit"] = "angstrom"
unrelaxedconfigforce = {}
unrelaxedconfigforce["source-value"] = []
unrelaxedconfigforce["source-unit"] = "eV/angstrom"
for line in atomslines.split("\n"):
    if line:
        fields = [float(l) for l in line.split()[1:]]
        # Convert LAMMPS 'type' to element string
        fields[0] = elements[int(fields[0]) - 1]
        # Set positions to original values (we do this instead of applying the
        # inverse rotation to the LAMMPS positions because LAMMPS always wraps
        # atoms that are outside of a periodic box.  We want to preserve the
        # original positions in case the user provides an extended xyz which
        # features atoms which are outside of the triclinic box.
        fields[1] = orig_pos[atom_count][0]
        fields[2] = orig_pos[atom_count][1]
        fields[3] = orig_pos[atom_count][2]
        if numpy.isnan(float(fields[1])):
            nanflag = 1
        if numpy.isnan(float(fields[2])):
            nanflag = 1
        if numpy.isnan(float(fields[3])):
            nanflag = 1
        # Rotate forces back to original orientation
        fx = fields[4]
        fy = fields[5]
        fz = fields[6]
        fields[4] = Rinv11 * fx + Rinv12 * fy + Rinv13 * fz
        fields[5] = Rinv21 * fx + Rinv22 * fy + Rinv23 * fz
        fields[6] = Rinv31 * fx + Rinv32 * fy + Rinv33 * fz
        if numpy.isnan(fields[4]):
            nanflag = 1
        if numpy.isnan(fields[5]):
            nanflag = 1
        if numpy.isnan(fields[6]):
            nanflag = 1
        atom_count = atom_count + 1
        configspecies["source-value"].append(fields[0])
        unrelaxedconfigpos["source-value"].append([fields[1], fields[2], fields[3]])
        unrelaxedconfigforce["source-value"].append([fields[4], fields[5], fields[6]])


results["species"] = configspecies
results["unrelaxed-periodic-cell-vector-1"] = unrelaxedcellvec1
results["unrelaxed-periodic-cell-vector-2"] = unrelaxedcellvec2
results["unrelaxed-periodic-cell-vector-3"] = unrelaxedcellvec3
results["unrelaxed-configuration-positions"] = unrelaxedconfigpos
results["unrelaxed-configuration-forces"] = unrelaxedconfigforce
results["unrelaxed-potential-energy"] = poteng

# If none of the reported quantities was NaN, print a results.edn file
if nanflag == 0:
    resultsedn = open("output/results.edn", "w")
    resultsedn.write(
        re.sub(
            '"([0-9.e+-]+)"',
            stripquotes,
            json.dumps(results, separators=(" ", " "), indent=2, sort_keys=False),
        )
    )
    resultsedn.close()