#!/usr/bin/env python3
"""
Author:      Daniel S. Karls
Date:        June 2019
Institution: Department of Aerospace Engineering and Mechanics
             University of Minnesota Twin Cities
Email:       karl0100 |AT| umn DOT edu

Description: This file contains a single function that computes
             the energy of an isolated atom of a given species
             for a specific model in eV.

Copyright:   2019 Regents of the University of Minnesota
"""

import os
import subprocess
import re
import textwrap

LAMMPS_INPUT_TEMPLATE = textwrap.dedent(
    r"""kim init {model_name} metal unit_conversion_mode
    boundary f f f

    variable boxextent equal 25.0
    region box block -${{boxextent}} ${{boxextent}} &
                     -${{boxextent}} ${{boxextent}} &
                     -${{boxextent}} ${{boxextent}}

    create_box 1 box

    create_atoms 1  single 0.0 0.0 0.0

    change_box all x scale ${{_u_distance}} &
                   y scale ${{_u_distance}} &
                   z scale ${{_u_distance}} &
                   remap

    mass * 1.0 # Mass unimportant since we're not doing time integration

    atom_modify sort 0 0

    kim interactions {species}

    # Use nsq neighlist method instead of binning since this is a small system
    variable neigh_skin equal 2.0*${{_u_distance}}
    neighbor ${{neigh_skin}} nsq

    # Variables used to rescale the positions and forces so that the quantities in the
    # dumpfile are in the original metal units (angstrom and eV/angstrom) even if we're
    # running with a Simulator Model that uses different units
    variable  pe_metal  equal  "c_thermo_pe/v__u_energy"
    thermo_style custom v_pe_metal

    log {logfile}

    run 0

    print "Isolated atom energy: ${{pe_metal}} eV"
    """)


def compute_isolated_atom_energy(output_dir, model_name, species):

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    lammps_input_file = os.path.join(output_dir, "isolated_atom_energy.lammps")
    lammps_log_file = os.path.join(output_dir, "isolated_atom_energy.lammps.log")

    with open(lammps_input_file, "w") as f:
        f.write(
            LAMMPS_INPUT_TEMPLATE.format(
                model_name=model_name, species=species, logfile=lammps_log_file
            )
        )

    # Run lammps
    try:
        with open(os.devnull, "w") as fnull:
            subprocess.check_call(
                ["lammps", "-i", lammps_input_file], shell=False, stdout=fnull
            )
    except subprocess.CalledProcessError:
        extrainfo = ""
        try:
            with open(lammps_log_file) as f:
                extrainfo = f.read()
        except IOError:
            extrainfo = "LAMMPS log file {} not found".format(lammps_log_file)
        raise RuntimeError("LAMMPS did not exit properly:\n" + extrainfo)

    # Parse LAMMPS output
    with open(lammps_log_file) as f:
        log = f.read()

    isolated_atom_energy = re.search(
        "Isolated atom energy: ([0-9.e-]+) eV", log
    ).groups(1)[0]
    isolated_atom_energy = float(isolated_atom_energy)

    return isolated_atom_energy