#! /usr/bin/env python
'''
Tool to gather ensemble energies and per-residue decomposed energies from
gromacs md files (edr and trr for decomposition)
The program execute gromacs commands from inside

examples of execution:
1. To gather energies from edr:
 python getEnergies.py energy -o energy.dat

2. to obtained per-residue decompose contributes:
 python getEnergies.py decompose -gmxrc /opt/gromacs-4.6.7/bin/GMXRC -o energydec.dat -res "416,417,418,419,420,421,422,423"

For decomposition, configuration as in mdpName='md-prod-out.mdp' is used.
Rerun is performed for the trajectory: ext='trr',pref='*?MD*'
template index file is: ext='ndx',pref='*?sol'
top file is: 'top',pref='*?sol'
gro file for getting aton umber for residue is: ext='gro',pref='*?sol'
'''

import argparse
import collections
import fnmatch
import glob
import logging
import numpy as np
import pandas
import os
import re
import shutil
import subprocess
import sys
from multiprocessing import cpu_count
from panedr import edr_to_df
from subprocess import (PIPE, Popen)

# Container for the files
Files = collections.namedtuple(
    "FILES", ("gro", "ndx", "trr", "top", "mdp", "tpr"))


def main(args):
    if args.mode == 'energy':
        process_energies(args, args.outName)

    # Per-residue energy decomposition
    if args.mode == 'decompose':
        decompose(args)
    logging.info('SUCCESSFUL COMPLETION OF THE PROGRAM')


def process_energies(args, outName):
    """
    Read and format energies using a edr file contains in
    `dataDir` and write it in `outName`.

    :params dataDir: Directory where the job is execute.
    :params outName: Name of the output file.
    """
    path_edr = get_edr_file(args)
    frames = get_energy(path_edr)
    frames.rename(index=str, columns={'Kinetic En.': 'Kinetic_Energy'}, inplace=True)
    labs2print = [
        'Time', 'Potential', 'Kinetic_Energy', 'Temperature', 'ele',
        'vdw', 'Ligand-Ligenv-ele', 'Ligand-Ligenv-vdw']
    writeOut(frames, outName, labs2print)


def get_energy(paths, listRes=['Ligand']):
    """
    Read Energies from .edr files and
    use the `panedr` library to parse them.

    :params paths:  Path to the edr files.
    :returns: Pandas dataframe.
    """
    if not isinstance(paths, list):
        df = edr_to_df(paths)
    else:
        # concatenate the data frames and reduce with the mean function
        rs = [edr_to_df(p) for p in paths]
        df = pandas.concat(rs)
        df.groupby(df.index).mean()

    # Reindex dataframe using sequential integers
    df.reset_index(inplace=True)

    # Electrostatic Energy
    df['ele'] = sum_available_columns(
        df, ['Coulomb-14', 'Coulomb (SR)', 'Coulomb (LR)', 'Coul. recip.'])

    # Van der Waals terms
    df['vdw'] = sum_available_columns(
        df, ['LJ-14', 'LJ (SR)', 'LJ (LR)'])

    return extract_ligand_info(df, listRes)


def decompose(args):
    """
    Make a decomposition of the energy into some if its residue components.
    """
    if args.resList is None:
        msg = 'TERMINATED. List of residues not provided.'
        logging.info(msg)
        return None

    gmx = search_commands(['gmx', 'gmx_mpi'], args.gmxEnv)
    if gmx is None:
        log_and_quit('gmx executable was not found')

    # parse MD mdp
    args_dict = vars(args)
    mdpIn = search_file_in_args(args_dict, ext='mdp', pref='md-prod-out')
    mdp_dict = parseMdp(mdpIn)

    # search for gromacs output files
    gro = search_file_in_args(args_dict, ext='gro', pref='*sol')
    ndx = search_file_in_args(args_dict, ext='ndx', pref='*-sol')
    trr = search_file_in_args(args_dict, ext='trr', pref='*MD.part*')
    top = search_file_in_args(args_dict, ext='top', pref='*-sol')
    files = Files(gro, ndx, trr, top, mdpIn, None)

    # rerun the molecular dynamics
    energy_files = decomp(mdp_dict, args, files, gmx)

    energy_analysis(args, energy_files)


def decomp(
         mdp_dict, args, files, gmx, ligGroup='Ligand', output_prefix='decompose'):
    """
    Decompose the energy into its components for different residues.
    """
    # create a dictionary with the index of the atoms that belong
    # to a given residue
    dict_residues = create_residue_dict(files.gro)

    def compute_decomposition(res, folder):
        workdir = create_workdir(args.dataDir, folder)

        copy_include_files(args.dataDir, workdir)

        # Generate new mdp file including residues
        new_mdp_file = create_new_mdp_file(mdp_dict, res, workdir, ligGroup)

        # create new ndx file
        new_ndx_file = create_new_ndx_file(dict_residues, res, workdir, files.ndx)

        # Generate new tpr file
        files_tpr = Files(files.gro, new_ndx_file, files.trr, files.top, new_mdp_file, None)
        new_tpr_file = create_new_tpr_file(files_tpr, workdir, gmx, args.gmxEnv)

        return rerun_md(new_tpr_file, files.trr, workdir, gmx, args.gmxEnv)

    # it is only possible to compute with gromacs 64 energy groups of a time
    residues = chunksOf(args.resList, 62)

    return [compute_decomposition(res, folder='chunk_{}'.format(i))
            for i, res in enumerate(residues)]


def rerun_md(tpr_file, trr_file, workdir, gmx, gmx_env):
    """
    Rerun the molecular dynamics and create decomposition of the energy.
    """
    mdrun = set_gmx_mpi_run(gmx)
    cmd = ['-s', tpr_file, '-rerun', trr_file, '-deffnm', 'decompose']
    rs, err = call_subprocess(mdrun + cmd, gmx_env, workdir)
    logging.error(err)
    logging.info(rs)

    energy_file = os.path.join(workdir, 'decompose.edr')
    if not os.path.exists(energy_file):
        msg = 'Something went wrong in the rerun decomposition analysis'
        log_and_quit(msg)

    return energy_file


def energy_analysis(args, energy_files):
    """Analysis of energy decomposition files after rerun"""
    df = get_energy(energy_files)

    write_decomposition_ouput(df, args.outName, args.resList)


def create_new_tpr_file(files, workdir, gmx, gmx_env):
    """
    Call gromacs grompp `http://manual.gromacs.org/programs/gmx-grompp.html`.
    """
    outTpr = os.path.join(workdir, 'decompose.tpr')
    cmd = [gmx, 'grompp', '-f', files.mdp, '-c', files.gro, '-p',
           files.top, '-n', files.ndx, '-o', outTpr, '-maxwarn', '2']
    rs, err = call_subprocess(cmd, gmx_env, cwd=workdir)
    logging.error(err)

    if not os.path.exists(outTpr):
        msg = 'Something went wrong in the creation of the tpr file  \
        for decomposition analysis'
        log_and_quit(msg)

    return outTpr


def set_gmx_mpi_run(gmx):
    """
    Try to run gmx mdrun using MPI see:
    `http://manual.gromacs.org/documentation/5.1/user-guide/mdrun-performance.html`
    """
    if gmx == "gmx_mpi":
        nranks = compute_mpi_ranks()
        gmx_bin = ['mpirun'] + nranks + [gmx]
        cmd = gmx_bin + ['mdrun']
    else:
        cmd = [gmx, 'mdrun', '-nice', '0', '-ntomp', '16']

    return cmd


def compute_mpi_ranks(ranks=''):
    """
    guess a reasonable default for the mpi ranks.
    """
    cpus = cpu_count()
    if cpus > 4:
        s = "-np 4 --map-by ppr:2:socket -v --display-map --display-allocation"
        ranks = s.split()

    return ranks


def create_new_ndx_file(residues_dict, residues, workdir, ndx_file):
    """
    Write a new ndx file mapping the residue numbers to their
    correspoding atoms, using a array `residues_dict` that contains
    in each row the lower and upper limit of the range of
    the atoms contained in a given residue.
    """
    new = ''
    for key in residues:
        new += '[ {} ]\n'.format(key)
        data = np.array2string(
            np.arange(*residues_dict[key]),
            formatter={'int': lambda s: '{:>7d}'.format(s)})[1:-1]
        new += ' {}\n'.format(data)

    new_ndx_file = os.path.join(workdir, 'decompose.ndx')
    with open(ndx_file, 'r') as f_inp, open(new_ndx_file, 'w') as f_out:
        old = f_inp.read()
        f_out.write(old + new)

    return new_ndx_file


def create_new_mdp_file(mdp_dict, residues, workdir, ligGroup):
    """
    Create a new input mdp file using the previous `mdp_dict` data.
    """

    listkeys = [
        'include', 'define', 'cutoff-scheme', 'ns-type', 'pbc',
        'periodic-molecules', 'rlist', 'rlistlong', 'nstcalclr',
        'coulombtype', 'coulomb-modifier', 'rcoulomb-switch', 'rcoulomb',
        'epsilon-r', 'epsilon-rf', 'vdw-type', 'vdw-modifier', 'rvdw-switch',
        'rvdw', 'DispCorr', 'table-extension', 'energygrp-table',
        'fourierspacing', 'fourier-nx', 'fourier-ny', 'fourier-nz',
        'pme-order', 'ewald-rtol', 'ewald-geometry', 'epsilon-surface',
        'optimize-fft', 'implicit-solvent', 'QMMM', 'constraints',
        'constraint-algorithm', 'continuation', 'Shake-SOR', 'shake-tol',
        'lincsorder', 'lincs-iter', 'lincs-warnangle', 'morse',
        'nwall', 'wall-type', 'wall-r-linpot', 'wall-atomtype', 'wall-density',
        'wall-ewald-zfac', 'pull', 'rotation', 'disre', 'orire',
        'free-energy', 'simulated-tempering']

    newKeys = {
        'nstxout': '0', 'nstvout': '0', 'nstfout': '0', 'nstlog': '0',
        'nstcalcenergy': '1', 'nstenergy': '1', 'nstxtcout': '0',
        'xtc-precision': '0', 'xtc-grps': '', 'nstlist': '1'}

    # the [1:-1] removes the [ ] from the beginning and the end
    str_residues = np.array_str(residues, max_line_width=1000)[1:-1]
    energygrps = '{} {}'.format(ligGroup, str_residues)

    newKeys['energygrps'] = energygrps

    new_input = ''
    fmt = "{:15s} = {}\n"
    for mdpKey in newKeys:
        new_input += fmt.format(mdpKey, newKeys[mdpKey])

    for mdpKey in (x for x in listkeys if x in mdp_dict):
        new_input += fmt.format(mdpKey, mdp_dict[mdpKey])

    mdp_file = os.path.join(workdir, "decompose.mdp")
    with open(mdp_file, 'w') as f:
        f.write(new_input)

    return mdp_file


def extract_ligand_info(df, listRes):
    """
    Get the Ligand information from a pandas dataframe `df`.
    """
    for res in listRes:
        names = get_residue_from_columns(res, df.columns)
        df = compute_terms_per_residue(df, names)

    return df


def get_residue_from_columns(name, columns):
    """
    Extract the ligand terms from the column names.

    :param name: Name of the residue
    :param columns: name of the columns in the dataframe
    :return: dictionary containing the residue's names as
           keys and electrostatic and vdw names as values.
    """
    names = collections.defaultdict(list)
    for c in filter(lambda x: name in x, columns):
        v, k = c.split(':')
        names[k].append(v)

    # Sort the list
    for xs in names.values():
        xs.sort()

    return names


def compute_terms_per_residue(df, names):
    """
    Split the names in the electrostatic and van der Waals terms
    and compute the electronic and VDW terms for each one of
    the residue terms specified in the `names` dictionary.

    :param df: Pandas dataframe
    :param names: dictionary of the names use in the dataframe.
    :return: updated dataframe
    """
    for key, vals in names.items():
        sufix_elec = vals[:2]
        sufix_vdw = vals[2:]
        names_elec = ['{}:{}'.format(x, key) for x in sufix_elec]
        names_vdw = ['{}:{}'.format(x, key) for x in sufix_vdw]
        res_elec = '{}-ele'.format(key)
        res_vdw = '{}-vdw'.format(key)

        df[res_elec] = df[names_elec].sum(axis=1)
        df[res_vdw] = df[names_vdw].sum(axis=1)

    return df


def sum_available_columns(df, labels):
    """
    Sum columns if they are in the dataframe
    """
    cols = [x for x in labels if x in df.columns]
    return df[cols].sum(axis=1)


def writeOut(frames, output_file, columns):
    """
    write columns as a table
    """
    with open(output_file, 'w') as f:
        s = '# FRAME ' + frames[columns].to_string()
        f.write(s)


def create_residue_dict(gro_file):
    """
    Read residues from the `gro_file` and create
    an array that map the the residues listed in `res_array`
    with their correspoding atoms.
    """
    def match_int(x):
        return re.match(r"([0-9]+)", x).groups()[0]

    arr = np.loadtxt(gro_file, dtype=np.bytes_, skiprows=2, usecols=0)
    arr = np.array(arr, dtype=np.str)
    xs = np.array(
        [match_int(x) for x in arr if 'SOL' not in x and '.' not in x],
        dtype=np.int32)

    # create range of the atoms for each residue
    unique, counts = np.unique(xs, return_counts=True)
    upper_limits = np.cumsum(counts) + 1
    lower_limits = np.insert(upper_limits[:-1], 0, 1)
    ranges = np.stack((lower_limits, upper_limits), axis=1)

    return dict(zip(unique, ranges))


def write_decomposition_ouput(listFrames, outName, resList):
    """ Write results for decomposition """
    hs1 = ['Time', 'Potential', 'ele', 'vdw']
    hs2 = ['Ligand-{}-vdw'.format(x) for x in resList]
    hs3 = ['Ligand-rest-vdw']
    hs4 = ['Ligand-{}-ele'.format(x) for x in resList]
    hs5 = ['Ligand-rest-ele']

    labs2print = hs1 + hs2 + hs3 + hs4 + hs5
    writeOut(listFrames, outName, labs2print)


def parseResidues(xs):
    """ return an array of the residues index """
    residues = np.array(xs.split(','), dtype=np.int32)

    return residues


def parseMdp(mdpIn):
    """ Read the mdp input file"""
    with open(mdpIn, 'r') as inFile:
        xss = inFile.readlines()

    rs = filter(lambda line: not line.startswith(';') and '=' in line, xss)

    return dict([check_lenght(x.split()[::2]) for x in rs])


def check_lenght(xs):
    """ transform list to tuples fixing the lenght """
    if len(xs) == 1:
        return xs[0], ''
    else:
        return tuple(xs)


def log_and_quit(msg):
    """
    Log a message and quit the program
    """
    logging.error(msg)
    sys.exit(-1)


def availProg(prog, myEnv):
    """ Check if a program is available """
    cmds = (os.path.join(path, prog)
            for path in myEnv["PATH"].split(os.pathsep))

    return any(os.path.isfile(cmd) and os.access(cmd, os.X_OK)
               for cmd in cmds)


def findFile(workdir, ext=None, pref=''):
    """
    Check whether a file starting with `pref` and ending with `ext` exists.
    """
    rs = fnmatch.filter(os.listdir(workdir), "{}.{}".format(pref, ext))
    if rs:
        return os.path.join(workdir, rs[0])

    else:
        logging.error(
            """
file not Found with prefix: {} and ext: {}
in dir: {}""".format(pref, ext, workdir))
        return None


def getGMXEnv(gmxrc):
    """ Use Gromacs environment """
    if gmxrc is None:
        return os.environ
    else:
        command = ['bash', '-c', 'source {} && env'.format(gmxrc)]
        rs = call_subprocess(command)[0]
        lines = rs.decode().splitlines()

        return dict(map(process_line, lines))


def get_edr_file(args):
    """
    """
    if args.edr is None:
        return findFile(args.dataDir, ext='edr', pref='*-MD.part*')
    else:
        return args.edr


def chunksOf(xs, n):
    """Yield successive n-sized chunks from xs"""
    for i in range(0, len(xs), n):
        yield xs[i:i + n]


def create_workdir(path, folder):
    """create a workdir if it does not exist """
    workdir = os.path.join(path, folder)
    if not os.path.isdir(workdir):
        os.mkdir(workdir)

    return workdir


def copy_include_files(path, workdir):
    """ Copy all the include topology files in the workdir"""
    for p in glob.glob("{}/*itp".format(path)):
        shutil.copy(p, workdir)


def process_line(line):
    """ Split a Line in key, value pairs """
    key, value = line.partition("=")[::2]
    return key, value.rstrip()


def search_file_in_args(args_dict, ext=None, pref=None):
    """
    Search if a file was passed as an argument otherwise look for it
    on the workdir.

    :params args: command line arguments
    :returns: path to file
    """
    file_path = args_dict.get(ext)
    if file_path is not None:
        return file_path
    else:
        return findFile(args_dict['dataDir'], ext=ext, pref=pref)


def call_subprocess(cmd, env=None, cwd=None):
    """
    Execute shell command `cmd`, using the environment `env`
    and wait for the results.
    """
    try:
        p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, env=env, cwd=cwd)
        rs = p.communicate()
        return rs

    except Exception as e:
        msg1 = "Subprocess fails with error: {}".format(e)
        msg2 = "Command: {}\n".format(cmd)
        log_and_quit(msg1 + msg2)


def check_command(cmd, env):
    try:
        with open(os.devnull, 'w') as f:
            r = subprocess.check_output(cmd, stderr=f, env=env)
            return r.split()[0]
    except subprocess.CalledProcessError:
        return None


def search_commands(cmds, env=None):
    xs = [check_command(["which", c], env) for c in cmds]
    rs = [x for x in xs if x is not None]
    if rs:
        return rs[0]
    else:
        return None


if __name__ == "__main__":
    logging.basicConfig(level='INFO')
    parser = argparse.ArgumentParser(
        description='Decompose the energy into its different components per residue')

    # Create separate parsers for the energy and decomposition
    subparsers = parser.add_subparsers(
        help='Parser for both energy and its decomposition',
        dest='mode')
    parser_energy = subparsers.add_parser(
        'energy', help='gather total energy')
    parser_dec = subparsers.add_parser(
        'decompose', help='perform residue decomposition analysis')

    # Arguments for total energy
    parser_energy.add_argument(
        '-edr', required=False,
        help='Gromacs energy output in edr format')

    # Arguments for energy decomposition
    parser_dec.add_argument(
        '-gmxrc', required=False, dest='gmxEnv', type=getGMXEnv,
        help='GMXRC file for environment loading')

    parser_dec.add_argument(
        '-res', '--residues', required=True, dest='resList',
        help='list of residue for which to decompose interaction energies (e.g.1 "1,2,3")',
        type=parseResidues)

    # Gromacs output files
    parser_dec.add_argument(
        '-gro', required=False, help='path to*.gro file')
    parser_dec.add_argument(
        '-ndx', required=False, help='path to*.ndx file')
    parser_dec.add_argument(
        '-trr', required=False, help='path to*.trr file')
    parser_dec.add_argument(
        '-top', required=False, help='path to*.top file')
    parser_dec.add_argument(
        '-mdp', required=False, help='path to*.mdp file')

    # Arguments for both parsers
    for p in [parser_energy, parser_dec]:
        p.add_argument(
            '-d', '--dir', required=False, dest='dataDir', type=os.path.abspath,
            help='directory with MD files to process', default=os.getcwd())
        p.add_argument(
            '-o', '--output', dest='outName', default='energy.out')

    args = parser.parse_args()
    main(args)

#! /usr/bin/env python
'''
Tool to gather ensemble energies and per-residue decomposed energies from
gromacs md files (edr and trr for decomposition)
The program execute gromacs commands from inside

examples of execution:
1. To gather energies from edr:
 python getEnergies.py energy -o energy.dat

2. to obtained per-residue decompose contributes:
 python getEnergies.py decompose -gmxrc /opt/gromacs-4.6.7/bin/GMXRC -o energydec.dat -res "416,417,418,419,420,421,422,423"

For decomposition, configuration as in mdpName='md-prod-out.mdp' is used.
Rerun is performed for the trajectory: ext='trr',pref='*?MD*'
template index file is: ext='ndx',pref='*?sol'
top file is: 'top',pref='*?sol'
gro file for getting aton umber for residue is: ext='gro',pref='*?sol'
'''

import argparse
import collections
import fnmatch
import glob
import logging
import numpy as np
import pandas
import os
import re
import shutil
import subprocess
import sys
from multiprocessing import cpu_count
from panedr import edr_to_df
from subprocess import (PIPE, Popen)

# Container for the files
Files = collections.namedtuple(
    "FILES", ("gro", "ndx", "trr", "top", "mdp", "tpr"))


def main(args):
    if args.mode == 'energy':
        process_energies(args, args.outName)

    # Per-residue energy decomposition
    if args.mode == 'decompose':
        decompose(args)
    logging.info('SUCCESSFUL COMPLETION OF THE PROGRAM')


def process_energies(args, outName):
    """
    Read and format energies using a edr file contains in
    `dataDir` and write it in `outName`.

    :params dataDir: Directory where the job is execute.
    :params outName: Name of the output file.
    """
    path_edr = get_edr_file(args)
    frames = get_energy(path_edr)
    frames.rename(index=str, columns={'Kinetic En.': 'Kinetic_Energy'}, inplace=True)
    labs2print = [
        'Time', 'Potential', 'Kinetic_Energy', 'Temperature', 'ele',
        'vdw', 'Ligand-Ligenv-ele', 'Ligand-Ligenv-vdw']
    writeOut(frames, outName, labs2print)


def get_energy(paths, listRes=['Ligand']):
    """
    Read Energies from .edr files and
    use the `panedr` library to parse them.

    :params paths:  Path to the edr files.
    :returns: Pandas dataframe.
    """
    if not isinstance(paths, list):
        df = edr_to_df(paths)
    else:
        # concatenate the data frames and reduce with the mean function
        rs = [edr_to_df(p) for p in paths]
        df = pandas.concat(rs)
        df.groupby(df.index).mean()

    # Reindex dataframe using sequential integers
    df.reset_index(inplace=True)

    # Electrostatic Energy
    df['ele'] = sum_available_columns(
        df, ['Coulomb-14', 'Coulomb (SR)', 'Coulomb (LR)', 'Coul. recip.'])

    # Van der Waals terms
    df['vdw'] = sum_available_columns(
        df, ['LJ-14', 'LJ (SR)', 'LJ (LR)'])

    return extract_ligand_info(df, listRes)


def decompose(args):
    """
    Make a decomposition of the energy into some if its residue components.
    """
    if args.resList is None:
        msg = 'TERMINATED. List of residues not provided.'
        logging.info(msg)
        return None

    gmx = search_commands(['gmx', 'gmx_mpi'], args.gmxEnv)
    if gmx is None:
        log_and_quit('gmx executable was not found')

    # parse MD mdp
    args_dict = vars(args)
    mdpIn = search_file_in_args(args_dict, ext='mdp', pref='md-prod-out')
    mdp_dict = parseMdp(mdpIn)

    # search for gromacs output files
    gro = search_file_in_args(args_dict, ext='gro', pref='*sol')
    ndx = search_file_in_args(args_dict, ext='ndx', pref='*-sol')
    trr = search_file_in_args(args_dict, ext='trr', pref='*MD.part*')
    top = search_file_in_args(args_dict, ext='top', pref='*-sol')
    files = Files(gro, ndx, trr, top, mdpIn, None)

    # rerun the molecular dynamics
    energy_files = decomp(mdp_dict, args, files, gmx)

    energy_analysis(args, energy_files)


def decomp(
         mdp_dict, args, files, gmx, ligGroup='Ligand', output_prefix='decompose'):
    """
    Decompose the energy into its components for different residues.
    """
    # create a dictionary with the index of the atoms that belong
    # to a given residue
    dict_residues = create_residue_dict(files.gro)

    def compute_decomposition(res, folder):
        workdir = create_workdir(args.dataDir, folder)

        copy_include_files(args.dataDir, workdir)

        # Generate new mdp file including residues
        new_mdp_file = create_new_mdp_file(mdp_dict, res, workdir, ligGroup)

        # create new ndx file
        new_ndx_file = create_new_ndx_file(dict_residues, res, workdir, files.ndx)

        # Generate new tpr file
        files_tpr = Files(files.gro, new_ndx_file, files.trr, files.top, new_mdp_file, None)
        new_tpr_file = create_new_tpr_file(files_tpr, workdir, gmx, args.gmxEnv)

        return rerun_md(new_tpr_file, files.trr, workdir, gmx, args.gmxEnv)

    # it is only possible to compute with gromacs 64 energy groups of a time
    residues = chunksOf(args.resList, 62)

    return [compute_decomposition(res, folder='chunk_{}'.format(i))
            for i, res in enumerate(residues)]


def rerun_md(tpr_file, trr_file, workdir, gmx, gmx_env):
    """
    Rerun the molecular dynamics and create decomposition of the energy.
    """
    mdrun = set_gmx_mpi_run(gmx)
    cmd = ['-s', tpr_file, '-rerun', trr_file, '-deffnm', 'decompose']
    rs, err = call_subprocess(mdrun + cmd, gmx_env, workdir)
    logging.error(err)
    logging.info(rs)

    energy_file = os.path.join(workdir, 'decompose.edr')
    if not os.path.exists(energy_file):
        msg = 'Something went wrong in the rerun decomposition analysis'
        log_and_quit(msg)

    return energy_file


def energy_analysis(args, energy_files):
    """Analysis of energy decomposition files after rerun"""
    df = get_energy(energy_files)

    write_decomposition_ouput(df, args.outName, args.resList)


def create_new_tpr_file(files, workdir, gmx, gmx_env):
    """
    Call gromacs grompp `http://manual.gromacs.org/programs/gmx-grompp.html`.
    """
    outTpr = os.path.join(workdir, 'decompose.tpr')
    cmd = [gmx, 'grompp', '-f', files.mdp, '-c', files.gro, '-p',
           files.top, '-n', files.ndx, '-o', outTpr, '-maxwarn', '2']
    rs, err = call_subprocess(cmd, gmx_env, cwd=workdir)
    logging.error(err)

    if not os.path.exists(outTpr):
        msg = 'Something went wrong in the creation of the tpr file  \
        for decomposition analysis'
        log_and_quit(msg)

    return outTpr


def set_gmx_mpi_run(gmx):
    """
    Try to run gmx mdrun using MPI see:
    `http://manual.gromacs.org/documentation/5.1/user-guide/mdrun-performance.html`
    """
    if gmx == "gmx_mpi":
        nranks = compute_mpi_ranks()
        gmx_bin = ['mpirun'] + nranks + [gmx]
        cmd = gmx_bin + ['mdrun']
    else:
        cmd = [gmx, 'mdrun', '-nice', '0', '-ntomp', '16']

    return cmd


def compute_mpi_ranks(ranks=''):
    """
    guess a reasonable default for the mpi ranks.
    """
    cpus = cpu_count()
    if cpus > 4:
        s = "-np 4 --map-by ppr:2:socket -v --display-map --display-allocation"
        ranks = s.split()

    return ranks


def create_new_ndx_file(residues_dict, residues, workdir, ndx_file):
    """
    Write a new ndx file mapping the residue numbers to their
    correspoding atoms, using a array `residues_dict` that contains
    in each row the lower and upper limit of the range of
    the atoms contained in a given residue.
    """
    new = ''
    for key in residues:
        new += '[ {} ]\n'.format(key)
        data = np.array2string(
            np.arange(*residues_dict[key]),
            formatter={'int': lambda s: '{:>7d}'.format(s)})[1:-1]
        new += ' {}\n'.format(data)

    new_ndx_file = os.path.join(workdir, 'decompose.ndx')
    with open(ndx_file, 'r') as f_inp, open(new_ndx_file, 'w') as f_out:
        old = f_inp.read()
        f_out.write(old + new)

    return new_ndx_file


def create_new_mdp_file(mdp_dict, residues, workdir, ligGroup):
    """
    Create a new input mdp file using the previous `mdp_dict` data.
    """

    listkeys = [
        'include', 'define', 'cutoff-scheme', 'ns-type', 'pbc',
        'periodic-molecules', 'rlist', 'rlistlong', 'nstcalclr',
        'coulombtype', 'coulomb-modifier', 'rcoulomb-switch', 'rcoulomb',
        'epsilon-r', 'epsilon-rf', 'vdw-type', 'vdw-modifier', 'rvdw-switch',
        'rvdw', 'DispCorr', 'table-extension', 'energygrp-table',
        'fourierspacing', 'fourier-nx', 'fourier-ny', 'fourier-nz',
        'pme-order', 'ewald-rtol', 'ewald-geometry', 'epsilon-surface',
        'optimize-fft', 'implicit-solvent', 'QMMM', 'constraints',
        'constraint-algorithm', 'continuation', 'Shake-SOR', 'shake-tol',
        'lincsorder', 'lincs-iter', 'lincs-warnangle', 'morse',
        'nwall', 'wall-type', 'wall-r-linpot', 'wall-atomtype', 'wall-density',
        'wall-ewald-zfac', 'pull', 'rotation', 'disre', 'orire',
        'free-energy', 'simulated-tempering']

    newKeys = {
        'nstxout': '0', 'nstvout': '0', 'nstfout': '0', 'nstlog': '0',
        'nstcalcenergy': '1', 'nstenergy': '1', 'nstxtcout': '0',
        'xtc-precision': '0', 'xtc-grps': '', 'nstlist': '1'}

    # the [1:-1] removes the [ ] from the beginning and the end
    str_residues = np.array_str(residues, max_line_width=1000)[1:-1]
    energygrps = '{} {}'.format(ligGroup, str_residues)

    newKeys['energygrps'] = energygrps

    new_input = ''
    fmt = "{:15s} = {}\n"
    for mdpKey in newKeys:
        new_input += fmt.format(mdpKey, newKeys[mdpKey])

    for mdpKey in (x for x in listkeys if x in mdp_dict):
        new_input += fmt.format(mdpKey, mdp_dict[mdpKey])

    mdp_file = os.path.join(workdir, "decompose.mdp")
    with open(mdp_file, 'w') as f:
        f.write(new_input)

    return mdp_file


def extract_ligand_info(df, listRes):
    """
    Get the Ligand information from a pandas dataframe `df`.
    """
    for res in listRes:
        names = get_residue_from_columns(res, df.columns)
        df = compute_terms_per_residue(df, names)

    return df


def get_residue_from_columns(name, columns):
    """
    Extract the ligand terms from the column names.

    :param name: Name of the residue
    :param columns: name of the columns in the dataframe
    :return: dictionary containing the residue's names as
           keys and electrostatic and vdw names as values.
    """
    names = collections.defaultdict(list)
    for c in filter(lambda x: name in x, columns):
        v, k = c.split(':')
        names[k].append(v)

    # Sort the list
    for xs in names.values():
        xs.sort()

    return names


def compute_terms_per_residue(df, names):
    """
    Split the names in the electrostatic and van der Waals terms
    and compute the electronic and VDW terms for each one of
    the residue terms specified in the `names` dictionary.

    :param df: Pandas dataframe
    :param names: dictionary of the names use in the dataframe.
    :return: updated dataframe
    """
    for key, vals in names.items():
        sufix_elec = vals[:2]
        sufix_vdw = vals[2:]
        names_elec = ['{}:{}'.format(x, key) for x in sufix_elec]
        names_vdw = ['{}:{}'.format(x, key) for x in sufix_vdw]
        res_elec = '{}-ele'.format(key)
        res_vdw = '{}-vdw'.format(key)

        df[res_elec] = df[names_elec].sum(axis=1)
        df[res_vdw] = df[names_vdw].sum(axis=1)

    return df


def sum_available_columns(df, labels):
    """
    Sum columns if they are in the dataframe
    """
    cols = [x for x in labels if x in df.columns]
    return df[cols].sum(axis=1)


def writeOut(frames, output_file, columns):
    """
    write columns as a table
    """
    with open(output_file, 'w') as f:
        s = '# FRAME ' + frames[columns].to_string()
        f.write(s)


def create_residue_dict(gro_file):
    """
    Read residues from the `gro_file` and create
    an array that map the the residues listed in `res_array`
    with their correspoding atoms.
    """
    def match_int(x):
        return re.match(r"([0-9]+)", x).groups()[0]

    arr = np.loadtxt(gro_file, dtype=np.bytes_, skiprows=2, usecols=0)
    arr = np.array(arr, dtype=np.str)
    xs = np.array(
        [match_int(x) for x in arr if 'SOL' not in x and '.' not in x],
        dtype=np.int32)

    # create range of the atoms for each residue
    unique, counts = np.unique(xs, return_counts=True)
    upper_limits = np.cumsum(counts) + 1
    lower_limits = np.insert(upper_limits[:-1], 0, 1)
    ranges = np.stack((lower_limits, upper_limits), axis=1)

    return dict(zip(unique, ranges))


def write_decomposition_ouput(listFrames, outName, resList):
    """ Write results for decomposition """
    hs1 = ['Time', 'Potential', 'ele', 'vdw']
    hs2 = ['Ligand-{}-vdw'.format(x) for x in resList]
    hs3 = ['Ligand-rest-vdw']
    hs4 = ['Ligand-{}-ele'.format(x) for x in resList]
    hs5 = ['Ligand-rest-ele']

    labs2print = hs1 + hs2 + hs3 + hs4 + hs5
    writeOut(listFrames, outName, labs2print)


def parseResidues(xs):
    """ return an array of the residues index """
    residues = np.array(xs.split(','), dtype=np.int32)

    return residues


def parseMdp(mdpIn):
    """ Read the mdp input file"""
    with open(mdpIn, 'r') as inFile:
        xss = inFile.readlines()

    rs = filter(lambda line: not line.startswith(';') and '=' in line, xss)

    return dict([check_lenght(x.split()[::2]) for x in rs])


def check_lenght(xs):
    """ transform list to tuples fixing the lenght """
    if len(xs) == 1:
        return xs[0], ''
    else:
        return tuple(xs)


def log_and_quit(msg):
    """
    Log a message and quit the program
    """
    logging.error(msg)
    sys.exit(-1)


def availProg(prog, myEnv):
    """ Check if a program is available """
    cmds = (os.path.join(path, prog)
            for path in myEnv["PATH"].split(os.pathsep))

    return any(os.path.isfile(cmd) and os.access(cmd, os.X_OK)
               for cmd in cmds)


def findFile(workdir, ext=None, pref=''):
    """
    Check whether a file starting with `pref` and ending with `ext` exists.
    """
    rs = fnmatch.filter(os.listdir(workdir), "{}.{}".format(pref, ext))
    if rs:
        return os.path.join(workdir, rs[0])

    else:
        logging.error(
            """
file not Found with prefix: {} and ext: {}
in dir: {}""".format(pref, ext, workdir))
        return None


def getGMXEnv(gmxrc):
    """ Use Gromacs environment """
    if gmxrc is None:
        return os.environ
    else:
        command = ['bash', '-c', 'source {} && env'.format(gmxrc)]
        rs = call_subprocess(command)[0]
        lines = rs.decode().splitlines()

        return dict(map(process_line, lines))


def get_edr_file(args):
    """
    """
    if args.edr is None:
        return findFile(args.dataDir, ext='edr', pref='*-MD.part*')
    else:
        return args.edr


def chunksOf(xs, n):
    """Yield successive n-sized chunks from xs"""
    for i in range(0, len(xs), n):
        yield xs[i:i + n]


def create_workdir(path, folder):
    """create a workdir if it does not exist """
    workdir = os.path.join(path, folder)
    if not os.path.isdir(workdir):
        os.mkdir(workdir)

    return workdir


def copy_include_files(path, workdir):
    """ Copy all the include topology files in the workdir"""
    for p in glob.glob("{}/*itp".format(path)):
        shutil.copy(p, workdir)


def process_line(line):
    """ Split a Line in key, value pairs """
    key, value = line.partition("=")[::2]
    return key, value.rstrip()


def search_file_in_args(args_dict, ext=None, pref=None):
    """
    Search if a file was passed as an argument otherwise look for it
    on the workdir.

    :params args: command line arguments
    :returns: path to file
    """
    file_path = args_dict.get(ext)
    if file_path is not None:
        return file_path
    else:
        return findFile(args_dict['dataDir'], ext=ext, pref=pref)


def call_subprocess(cmd, env=None, cwd=None):
    """
    Execute shell command `cmd`, using the environment `env`
    and wait for the results.
    """
    try:
        p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, env=env, cwd=cwd)
        rs = p.communicate()
        return rs

    except Exception as e:
        msg1 = "Subprocess fails with error: {}".format(e)
        msg2 = "Command: {}\n".format(cmd)
        log_and_quit(msg1 + msg2)


def check_command(cmd, env):
    try:
        with open(os.devnull, 'w') as f:
            r = subprocess.check_output(cmd, stderr=f, env=env)
            return r.split()[0]
    except subprocess.CalledProcessError:
        return None


def search_commands(cmds, env=None):
    xs = [check_command(["which", c], env) for c in cmds]
    rs = [x for x in xs if x is not None]
    if rs:
        return rs[0]
    else:
        return None


if __name__ == "__main__":
    logging.basicConfig(level='INFO')
    parser = argparse.ArgumentParser(
        description='Decompose the energy into its different components per residue')

    # Create separate parsers for the energy and decomposition
    subparsers = parser.add_subparsers(
        help='Parser for both energy and its decomposition',
        dest='mode')
    parser_energy = subparsers.add_parser(
        'energy', help='gather total energy')
    parser_dec = subparsers.add_parser(
        'decompose', help='perform residue decomposition analysis')

    # Arguments for total energy
    parser_energy.add_argument(
        '-edr', required=False,
        help='Gromacs energy output in edr format')

    # Arguments for energy decomposition
    parser_dec.add_argument(
        '-gmxrc', required=False, dest='gmxEnv', type=getGMXEnv,
        help='GMXRC file for environment loading')

    parser_dec.add_argument(
        '-res', '--residues', required=True, dest='resList',
        help='list of residue for which to decompose interaction energies (e.g.1 "1,2,3")',
        type=parseResidues)

    # Gromacs output files
    parser_dec.add_argument(
        '-gro', required=False, help='path to*.gro file')
    parser_dec.add_argument(
        '-ndx', required=False, help='path to*.ndx file')
    parser_dec.add_argument(
        '-trr', required=False, help='path to*.trr file')
    parser_dec.add_argument(
        '-top', required=False, help='path to*.top file')
    parser_dec.add_argument(
        '-mdp', required=False, help='path to*.mdp file')

    # Arguments for both parsers
    for p in [parser_energy, parser_dec]:
        p.add_argument(
            '-d', '--dir', required=False, dest='dataDir', type=os.path.abspath,
            help='directory with MD files to process', default=os.getcwd())
        p.add_argument(
            '-o', '--output', dest='outName', default='energy.out')

    args = parser.parse_args()
    main(args)

#! /usr/bin/env python Tool to gather ensemble energies and per-residue decomposed energies from gromacs md files (edr and trr for decomposition) The program execute gromacs commands from inside  examples of execution: 1. To gather energies from edr:  python getEnergies.py energy -o energy.dat  2. to obtained per-residue decompose contributes:  python getEnergies.py decompose -gmxrc /opt/gromacs-4.6.7/bin/GMXRC -o energydec.dat -res "416,417,418,419,420,421,422,423"  For decomposition, configuration as in mdpName='md-prod-out.mdp' is used. Rerun is performed for the trajectory: ext='trr',pref='*?MD*' template index file is: ext='ndx',pref='*?sol' top file is: 'top',pref='*?sol' gro file for getting aton umber for residue is: ext='gro',pref='*?sol' # Container for the files # Per-residue energy decomposition Read and format energies using a edr file contains in     `dataDir` and write it in `outName`.      :params dataDir: Directory where the job is execute.     :params outName: Name of the output file. Read Energies from .edr files and     use the `panedr` library to parse them.      :params paths:  Path to the edr files.     :returns: Pandas dataframe. # concatenate the data frames and reduce with the mean function # Reindex dataframe using sequential integers # Electrostatic Energy # Van der Waals terms Make a decomposition of the energy into some if its residue components. # parse MD mdp # search for gromacs output files # rerun the molecular dynamics Decompose the energy into its components for different residues. # create a dictionary with the index of the atoms that belong # to a given residue # Generate new mdp file including residues # create new ndx file # Generate new tpr file # it is only possible to compute with gromacs 64 energy groups of a time Rerun the molecular dynamics and create decomposition of the energy. Analysis of energy decomposition files after rerun Call gromacs grompp `http://manual.gromacs.org/programs/gmx-grompp.html`. Try to run gmx mdrun using MPI see:     `http://manual.gromacs.org/documentation/5.1/user-guide/mdrun-performance.html` guess a reasonable default for the mpi ranks. Write a new ndx file mapping the residue numbers to their     correspoding atoms, using a array `residues_dict` that contains     in each row the lower and upper limit of the range of     the atoms contained in a given residue. Create a new input mdp file using the previous `mdp_dict` data. # the [1:-1] removes the [ ] from the beginning and the end Get the Ligand information from a pandas dataframe `df`. Extract the ligand terms from the column names.      :param name: Name of the residue     :param columns: name of the columns in the dataframe     :return: dictionary containing the residue's names as            keys and electrostatic and vdw names as values. # Sort the list Split the names in the electrostatic and van der Waals terms     and compute the electronic and VDW terms for each one of     the residue terms specified in the `names` dictionary.      :param df: Pandas dataframe     :param names: dictionary of the names use in the dataframe.     :return: updated dataframe Sum columns if they are in the dataframe write columns as a table Read residues from the `gro_file` and create     an array that map the the residues listed in `res_array`     with their correspoding atoms. # create range of the atoms for each residue Write results for decomposition return an array of the residues index Read the mdp input file transform list to tuples fixing the lenght Log a message and quit the program Check if a program is available Check whether a file starting with `pref` and ending with `ext` exists. file not Found with prefix: {} and ext: {} in dir: {} Use Gromacs environment Yield successive n-sized chunks from xs create a workdir if it does not exist Copy all the include topology files in the workdir Split a Line in key, value pairs Search if a file was passed as an argument otherwise look for it     on the workdir.      :params args: command line arguments     :returns: path to file Execute shell command `cmd`, using the environment `env`     and wait for the results. # Create separate parsers for the energy and decomposition # Arguments for total energy # Arguments for energy decomposition # Gromacs output files # Arguments for both parsers