##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Kyle A. Beauchamp
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#############################################################################


##############################################################################
# Imports
##############################################################################

import os
from distutils.spawn import find_executable as _find_executable
import numpy as np

from mdtraj.utils import import_
from mdtraj.utils import enter_temp_directory

##############################################################################
# Globals
##############################################################################

# Possible names for the external commands -- these are expected
# to be found in the PATH.
SHIFTX2 = ['shiftx2.py']
SPARTA_PLUS = ['sparta+', 'SPARTA+', 'SPARTA+.linux']
PPM = ['ppm_linux_64.exe']

__all__ = ['chemical_shifts_shiftx2', 'chemical_shifts_ppm', 'chemical_shifts_spartaplus', "reindex_dataframe_by_atoms"]


def find_executable(names):
    for possible in names:
        result = _find_executable(possible)
        if result is not None:
            return result
    return None


##############################################################################
# Code
##############################################################################

def compute_chemical_shifts(trj, model="shiftx2", **kwargs):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    model : str, optional, default="shiftx2"
        The program to use for calculating chemical shifts.  Must be one
        of shiftx2, ppm, or sparta+

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have the appropriate chemical soft programs installed
    and in your executable path.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference, see docstrings for chemical_shifts_*
    for the various possible models.
    """
    if model == "shiftx2":
        return chemical_shifts_shiftx2(trj, **kwargs)
    elif model == "ppm":
        return chemical_shifts_ppm(trj, **kwargs)
    elif model == "sparta+":
        return chemical_shifts_spartaplus(trj, **kwargs)
    else:
        raise(ValueError("model must be one of shiftx2, ppm, or sparta+"))


def chemical_shifts_shiftx2(trj, pH=5.0, temperature=298.00):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    pH : float, optional, default=5.0
        pH value which gets passed to the ShiftX2 predictor.
    temperature : float, optional, default=298.00
        Temperature which gets passed to the ShiftX2 predictor.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb' -p %.1f -t %.2f" % (binary, pH, temperature)

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results


def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_table("./bb_details.dat", index_col=False, header=None, sep="\s*").drop([3], axis=1)

        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d


def _get_lines_to_skip(filename):
    """Determine the number of comment lines in a SPARTA+ output file."""
    format_string = """FORMAT %4d %4s %4s %9.3f %9.3f %9.3f %9.3f %9.3f %9.3f"""
    handle = open(filename)
    for i, line in enumerate(handle):
        if line.find(format_string) != -1:
            return i + 2

    raise(Exception("No format string found in SPARTA+ file!"))


def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".  
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    
    if rename_HN:
        results.name[results.name == "HN"] = "H"
    
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results


def reindex_dataframe_by_atoms(trj, frame):
    """Reindex chemical shift output to use atom number (serial) indexing.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Returns
    -------
    new_frame : pandas.DataFrame
        Dataframe containing results, with index consisting of atom
        indices (AKA the 'serial' entry in a PDB).  Columns correspond to
        each frame in trj.

    Notes
    -----
    Be aware that this function may DROP predictions if the atom naming
    is different between the input trajectory and the output of various
    chemical shift prediction tools.
    """

    top, bonds = trj.top.to_dataframe()
    top["serial"] = top.index
    top = top.set_index(["resSeq", "name"])

    new_frame = frame.copy()

    new_frame["serial"] = top.ix[new_frame.index].serial
    new_frame = new_frame.dropna().reset_index().set_index("serial").drop(["resSeq", "name"], axis=1)

    return new_frame
