#!/usr/bin/env python

# pmx  Copyright Notice
# ============================
#
# The pmx source code is copyrighted, but you can freely use and
# copy it as long as you don't change or remove any of the copyright
# notices.
#
# ----------------------------------------------------------------------
# pmx is Copyright (C) 2006-2013 by Daniel Seeliger
#
#                        All Rights Reserved
#
# Permission to use, copy, modify, distribute, and distribute modified
# versions of this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both the copyright notice and
# this permission notice appear in supporting documentation, and that
# the name of Daniel Seeliger not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# DANIEL SEELIGER DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
# SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS.  IN NO EVENT SHALL DANIEL SEELIGER BE LIABLE FOR ANY
# SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
# CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
# ----------------------------------------------------------------------

"""Various utility functions and classes
"""

import os
import sys
import types
import numpy as np
import re
import logging
from glob import glob


# =============
# File IO utils
# =============

def clean_gromacs_backup_files( path ):
    """Clean gromacs backup files.

    Parameters
    ----------
    path: str
        path where to delete files starting with #
    """

    toclean = glob('{0}/*#'.format(path)) 
    for clean in toclean:
        os.remove(clean)     

def copy_files_folders( source, target ):
    """Copy files or folders.

    Parameters
    ----------
    source: str
        path to files/folders
    target: str
        target path
    """

    cmd = 'cp {0} {1} -r'.format(source,target)
    if os.path.isdir(target):
        cmd = 'cp {0} {1}/. -r'.format(source,target)
    os.system(cmd)

def create_folder( path, fname=False ):
    """Create a folder.

    Parameters
    ----------
    path: str
        full path to the folder. Alternatively, if fname is given,
        folder fname will be created in this path
    fname: str, optional
        name of the folder to create
    """ 
    if fname==False:
        if not os.path.exists(path):
            os.makedirs(path)
    elif not os.path.exists(path+'/'+fname):
        os.makedirs(path+'/'+fname)

def show_ff(gmxlib=None):
    """Prints the list of forcefields available in GMXLIB.

    Parameters
    ----------
    gmxlib : str, optional
        Path to force field library. If not set explicitly, it is taken from
        the enviroment variable GMXLIB.
    """

    if gmxlib is None:
        gmxlib = os.environ['GMXLIB']

    print('Available Force Fields in $GMXLIB:\n')
    print('  [i]    {0:40}{1}\n'.format('name', 'description'), end='')
    print('  ---    {0:40}{1}\n'.format('----', '-----------'), end='')

    ffs = [d for d in glob('{}/*.ff'.format(gmxlib)) if os.path.isdir(d)]
    ffs = natural_sort(ffs)

    for i, ff in enumerate(ffs):
        # get ff name
        f = os.path.basename(ff).split('.')[0]
        # print info
        docfile = ff + '/forcefield.doc'
        docstring = [l for l in open(docfile, 'r').readlines()][0]
        print('  [{0}]    {1:40}{2}'.format(i+1, f, docstring), end='')


def ffopen(filename, mode='r', backup=True):

    if mode == 'w':
        if os.path.isfile(filename):
            if backup:
                print('Backing up %s to %s~' % (filename, filename))
                os.rename(filename, filename+'~')
        try:
            fp = open(filename, 'w')
            return fp
        except:
            print('Error: Could not open file %s' % filename)

    elif mode == 'r':
        try:
            fp = open(filename, 'r')
            return fp
        except:
            print('No such file %s' % filename)

    else:
        return open(filename, mode)

# this function echanges netmount with home
def remove_netmount( string ): 
    regexp = re.compile('netmount')
    string = regexp.sub('home',string)
    return(string)

def get_ff_path(ff, verbose=False):
    """Get path of force field.

    Parameters
    ----------
    ff : str
        force field name

    Returns
    -------
    ff_path: str
        absolute path to the force field
    """
    ff_path = None
    # if ff is in current folder, take that path
    if os.path.isdir(ff):
        ff_path = ff
    # otherwise look into GMXLIB and GMXDATA
    else:
        # if GMXLIB is defined, look in there first
        if "GMXLIB" in os.environ:
            gmxlib = os.environ.get('GMXLIB')
            p = os.path.join(gmxlib, ff)
            pff = p+'.ff'
            if os.path.isdir(p):
                ff_path = p
            elif os.path.isdir(pff):
                ff_path = pff
        # then look into GMXDATA/top
        if "GMXDATA" in os.environ:
            gmxdata = "%s/top" % os.environ.get('GMXDATA')
            p = os.path.join(gmxdata, ff)
            pff = p+'.ff'
            if os.path.isdir(p):
                ff_path = p
            elif os.path.isdir(pff):
                ff_path = pff
        # if path still not defined, we have a problem
        if ff_path is None:
            raise ValueError('forcefield path "%s" not found' % ff)

    if verbose is True:
        print('Opening forcefield: %s' % ff_path)

    return ff_path


def get_pmxdata(fname):
    """Returns the absolute path for the file with 'fname' inside the pmx
    data folder.
    """
    pmxhome = os.path.dirname(os.path.realpath(__file__))
    pmxdata = os.path.join(pmxhome, 'data')
    filepath = '%s/%s' % (pmxdata, fname)
    if os.path.exists(filepath):
        return filepath
    else:
        raise IOError('File %s cannot be found in the following path:\n'
                      '%s' % (fname, pmxdata))



def get_mtp_file(residue, ff):
    """Returns the correct mutres mtp file depending on the nature of the
    residue.

    Parameters
    ----------
    residue : Molecule instance
        the residue that will be mutated
    ff : str
        the force field to use

    Returns
    -------
    mtp_file : str
        path to the mtp file
    """
    # Determine which mtp file to use
    ffpath = get_ff_path(ff=ff)
    # DNA mutation
    if residue.moltype == 'dna':
        mtp_file = os.path.join(ffpath, 'mutres_dna.mtp')
    # RNA mutation
    elif residue.moltype == 'rna':
        mtp_file = os.path.join(ffpath, 'mutres_rna.mtp')
    # Protein mutation
    elif residue.moltype == 'protein':
        mtp_file = os.path.join(ffpath, 'mutres.mtp')
    else:
        raise(ValueError, 'Cannot undertand mutation type needed '
                          'from the input strcture provided')
    return mtp_file


def ff_selection(gmxlib=None):
    print('Choose a force field:\n')
    print('  [i]    {0:40}{1}\n'.format('name', 'description'), end='')
    print('  ---    {0:40}{1}\n'.format('----', '-----------'), end='')

    # if no gmxlib argument passed, then get $GMXLIB
    if gmxlib is None:
        gmxlib = os.environ['GMXLIB']

    ffs = [d for d in glob('{}/*.ff'.format(gmxlib)) if os.path.isdir(d)]
    ffs = natural_sort(ffs)

    ffs_dict = {}
    for i, ff in enumerate(ffs):
        # get ff name
        f = os.path.basename(ff).split('.')[0]
        # store in dict
        ffs_dict[i+1] = f
        # print info
        docfile = ff + '/forcefield.doc'
        docstring = [l for l in open(docfile, 'r').readlines()][0]
        print('  [{0}]    {1:40}{2}'.format(i+1, f, docstring), end='')

    # allow choice of ff
    print('')
    print('Enter the index [i] of the chosen forcefield: ', end='')
    i = int(input())
    ff = ffs_dict[i]

    return ff


def listFiles(dir='./', ext=None, abs=True, backups=False):

    """ returns a list of files in
    directory dir, optionally only
    certain file types"""

    if dir[-1] != os.sep:
        dir += os.sep
    if dir[0] == '~':
        home = os.environ.get('HOME')
        dir = os.path.join(home, dir[2:])
    dir = os.path.abspath(dir)+os.sep

    l = os.listdir(dir)

    fl = []
    if not ext:
        for f in l:
            if os.path.isfile(dir+f):
                if backups:
                    if f[0] == '#' or \
                       f[-1] == '~':
                        fl.append(dir+f)
                else:
                    fl.append(dir+f)

    elif type(ext) in [types.ListType, types.TupleType]:
        for ex in ext:
            if backups:
                ff = glob(dir+'#*'+ex+'*')
                ff += glob(dir+'*.'+ex+'~')
            else:
                ff = glob(dir+'*.'+ex)
            fl.extend(ff)

    elif type(ext) == types.StringType:
        if backups:
            print(dir+'*.'+ext+'~')
            fl = glob(dir+'#*.'+ext+'*')
            fl += glob(dir+'*.'+ext+'~')
        else:
            fl.extend(glob(dir+'*.'+ext))

    if not abs:
        new = []
        for f in fl:
            new.append(f.split('/')[-1])
        return new

    return fl


def listDirs(dir='./'):
    """ returns a list of directrories in
    directory dir"""

    if dir[-1] != os.sep:
        dir += os.sep
    if dir[0] == '~':
        home = os.environ.get('HOME')
        dir = os.path.join(home, dir[2:])
    dir = os.path.abspath(dir)+os.sep

    l = os.listdir(dir)

    dl = []

    for f in l:
        if os.path.isdir(dir+f):
            dl.append(dir+f)

    return dl


def killBackups(arg, dirname, fname):

    l = listFiles(dirname, arg[0], arg[1], arg[2])
    if arg[3]:
        for f in l:
            print('%s' % f)


def removeBackups(dir, check=True):

    if dir[-1] != os.sep:
        dir += os.sep
    if dir[0] == '~':
        home = os.environ.get('HOME')
        dir = os.path.join(home, dir[2:])
    dir = os.path.abspath(dir)+os.sep

    os.path.walk(dir, killBackups, (False, True, True, check))


# ======
# Others
# ======
def doLog(fp, s, commandName='ligandHybridTop'):
    l = '{0}__log_> {1}'.format(commandName,s)
#   print(sys.stderr,l)
    print(l)
    fp.write(l+'\n')

def data2gauss(data):
    '''Takes a one dimensional array and fits a Gaussian.

    Parameters
    ----------
    data : list
        1D array of values

    Returns
    -------
    float
        mean of the distribution.
    float
        standard deviation of the distribution.
    float
        height of the curve's peak.
    '''
    m = np.average(data)
    dev = np.std(data)
    A = 1./(dev*np.sqrt(2*np.pi))
    return m, dev, A


def gauss_func(A, mean, dev, x):
    '''Given the parameters of a Gaussian and a range of the x-values, returns
    the y-values of the Gaussian function'''
    x = np.array(x)
    y = A*np.exp(-(((x-mean)**2.)/(2.0*(dev**2.))))
    return y


def natural_sort(l):
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
    return sorted(l, key=alphanum_key)


def multiple_replace(text, word_dic, isfile=False):
    """Takes a text and replace words that match a key in a dictionary with
    the associated value, returning the modified text.

    Parameters
    ----------
    text : str
        input text
    word_dic : dict
        dictionary where keys strings are replaced by their value strings
    isfile : bool, optional
        if isfile is True, the input 'text' will be considered a textfile,
        which will be opened, the patterns replaced, and written again.

    Returns
    -------
    newtext : str
        output text with strings replaced

    Examples
    --------
    >>> a = 'hello, John'
    >>> multiple_replace(a, {'hello':'goodbye', 'John':'Mark'})
    'goodbye, Mark'

    """

    rc = re.compile('|'.join(map(re.escape, word_dic)))
    def translate(match):
        return word_dic[match.group(0)]

    if isfile is False:
        return rc.sub(translate, text)
    elif isfile is True:
        f = open(text, 'r').read()
        newf = rc.sub(translate, f)
        with open(text, 'w') as out:
            out.write(newf)


def list2file(lst, fname):
    """Writes a list of strings to a file.
    """
    with open(fname, 'w') as f:
        for line in lst:
            f.write(line)


def initialise_logger(logname, logfile, level='INFO'):
    if os.path.isfile(logfile):
        os.remove(logfile)
    logger = logging.getLogger(logname)
    if level.upper() == 'INFO':
        logger.setLevel(logging.INFO)
    elif level.upper() == 'DEBUG':
        logger.setLevel(logging.DEBUG)
    else:
        raise ValueError('level can only be INFO or DEBUG')
    hdlr = logging.FileHandler(logfile)
    formatter = logging.Formatter('%(asctime)s [ %(name)s ] '
                                  '[ %(levelname)s ] %(message)s')
    hdlr.setFormatter(formatter)
    logger.addHandler(hdlr)
    return logger


def which(program):
    """Equivalent of UNIX which command. Returns the path of the executable
    if it is found, otherwise returns None.

    Parameters
    ----------
    program : str
        name of executable you are looking for

    Returns
    -------
    path : str
        string of path if found, otherwise None.

    """
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

# ========================
# Custom Exception Classes
# ========================
class UnknownResidueError(Exception):
    """Class for unknown residue Exceptions.
    """
    def __init__(self, s):
        self.s = s

    def __str__(self):
        return repr(self.s)


class mtpError(Exception):
    """Exceptions class for ...
    """
    def __init__(self, s):
        self.s = s

    def __str__(self):
        return repr(self.s)


class MissingTopolParamError(Exception):
    """Raised when certain parameters for terms defined in the topology cannot
    ba found.
    """
    def __init__(self, s, atoms):
        self.s = s
        print('err_> {}'.format(s), file=sys.stderr)
        print('err_>  name      resname      atomtype       atomtypeB       '
              'bondtype       bondtypeB', file=sys.stderr)
        for atom in atoms:
            print('{0.name} {0.resname} {0.atomtype} {0.atomtypeB} {0.type} '
                  '{0.typeB}'.format(atom), file=sys.stderr)
        print('err_> Exiting', file=sys.stderr)

    def __str__(self):
        return repr(self.s)

class importError(Exception):
    """Exceptions class for importing a module.
    """
    def __init__(self, s):
        self.s = "Could not import: {0}".format(s)

    def __str__(self):
        return repr(self.s)

