#!/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.
# ----------------------------------------------------------------------

"""Module containing fuctions dealing with Gromacs index files.
"""

import re
import sys,os,random
from .parser import readSection
from .model import Model


class IndexGroup:
    """Class for storing information about a Gromacs index group.

    Parameters
    ----------
    name : str
        name of the group
    ids : list, optional
        list of atom indices
    atoms : list, optional
        list of ``Atom`` instances. If this is provided, the ids will be taken from
        the Atom instances.

    Attributes
    ----------
    name : str
        name of the group
    ids : list
        list of atom indices
    """
    def __init__(self, name='', ids=[], atoms=[]):
        self.ids = []
        if atoms:
            self.ids = list(map(lambda a: a.id, atoms))
        else:
            self.ids = ids
        self.name = name
        self.ids.sort()

    def __str__(self):
        s = ''
        s += '[ %s ]\n' % self.name
        count = 0
        for i in range(len(self.ids)):
            if len(str(self.ids[i])) > 6:
                fs = "%d "
            else:
                fs = "%6d "
            s += fs % self.ids[i]
            count += 1
            if count % 15 == 0:
                s += '\n'
        return s

    def read_index_group(self, name, lines):
        if '[' in name:
            name = name[1:-1].strip()
        self.name = name
        key = '[ '+name+' ]'
        ll = readSection(lines, key, '[')
        self.ids = []
        for line in ll:
            self.ids.extend([int(x) for x in line.split()])
        return self

    def select_atoms(self, atom_sel):
        atoms = []
        for idx in self.ids:
            for atom in atom_sel.atoms:
                if atom.id == idx:
                    atoms.append(atom)
        return atoms


class IndexFile:
    """Class for storing ``ndx`` file information. It can be initialsed as an empty
    index file, or by providing an ``ndx`` file as input, or by providing a list
    of names and ``IndexGroup`` instances.

    Parameters
    ----------
    fname : str, optional
        input index file. Default is None.
    names : list, optional
        list of the group names.
    groups : list, optional
        list of IndexGroup instances.

    Attributes
    ----------
    groups : list
        list of IndexGroup instances
    names : list
        list if index group names
    dic : dict
        dictionary containg index group names as keys and ``IndexGroup``
        instances as values
    """

    def __init__(self, fname=None, names=[], groups=[]):
        self.groups = []
        self.names = []
        self.dic = {}
        if fname is not None:
            self.parse(fname)
        if groups:
            for g in groups:
                self.add_group(g)

    def __str__(self):
        s = 'Gromacs index file (%d index groups):\n' % len(self.groups)
        for i, group in enumerate(self.groups):
            s += ' %d  %-20s   :    %d atoms\n' % (i, group.name, len(group.ids))
        return s

    def __getitem__(self, item):
        return self.dic[item]

    def __delitem__(self, item):
        self.delete_group(item)

    def parse(self, fp):
        """Reads an index file.

        Parameters
        ----------
        fp : str|object
            filename of the index file, or the file object of the ndx file
            already opened.
        """
        if hasattr(fp, "read"):
            f = fp.read()
        else:
            f = open(fp).read()
        names = self.__get_index_groups(f)
        lines = f.split('\n')
        for name in names:
            idx = IndexGroup().read_index_group(name, lines)
            self.add_group(idx)
        return self

    def __get_index_groups(self, f):
        x = re.compile('\[ .* \]\n')
        names = x.findall(f)
        r = []
        for name in names:
            r.append(name.strip())
        return r

    def write(self, fn=None, fp=None):
        """Writes the index file.

        Parameters
        ----------
        fn : str, optional
            filename. Default is ``None``. If both ``fn`` and ``fp`` are
            ``None``, the index is printed to screen.
        fp : object, optional
            file object to write the index to. Default is ``None``.
        """
        if not fn and not fp:
            fp = sys.stdout
        if fn:
            fp = open(fn, 'w')
        for gr in self.groups:
            print('{0}\n'.format(str(gr)), file=fp)
        if fn:
            fp.close()

    def add_group(self, group):
        """Adds a group to the IndexFile.

        Parameters
        ----------
        group : IndexGroup
            instance of IndexGroup to add
        """
        if group.name in self.names:
            print("IndexFile has group %s !! " % group.name, file=sys.stderr)
            print("Group %s will be replaced !!" % group.name, file=sys.stderr)
            self.delete_group(group.name)
        self.names.append(group.name)
        self.groups.append(group)
        self.dic[group.name] = group

    def delete_group(self, name):
        """Removes a group from the IndexFile.

        Parameters
        ----------
        name : str
            name of the group to remove
        """
        idx = -1
        for i, group in enumerate(self.groups):
            if group.name == name:
                idx = i
                break
        if idx == -1:
            return
        del self.groups[idx]
        del self.dic[name]
        self.names.remove(name)


def get_index(atom_list=None, residue_list=None, chain_list=None):
    """ return atom indices from a list of atoms/residues/chains"""
    if not atom_list and not residue_list and not chain_list:
        print('Error: Need list~')
        sys.exit(1)
    if atom_list:
        lst = list(map(lambda a: a.id, atom_list))
        return lst
    if residue_list:
        al = []
        list(map(lambda r: al.extend(r.atoms), residue_list))
        return get_index(atom_list=al)
    if chain_list:
        al = []
        list(map(lambda c: al.extend(c.atoms), chain_list))
        return get_index(atom_list=al)


def make_index_group(atomlist, name):
    lst = get_index(atomlist)
    g = IndexGroup(ids=lst, name=name)
    return g


def create_ndx_freeze( pdbFile, ndxFile='index_FREEZE.ndx', groupname='FREEZE' ):
    """Create a group for freezing with all atoms except dummies.

    Parameters
    ----------
    pdbfile : str
            filename of pdb. The group will be constructed based on the atom names.
    ndxFile : str, optional
            filename of ndx. An existing or not existing index file can be provided.
    groupname : str, optional
            name for the new group
    """

    # deal with the index file
    if os.path.isfile(ndxFile): # file exists
        indFile = IndexFile(ndxFile)
        randnum = random.randint(1000,9999)
        if groupname in indFile.names:
            groupname = '{0}{1}'.format(groupname,randnum)
    else: # new file
        indFile = IndexFile()
    
    # create index group
    m = Model(pdbFile)
    i = 0
    ind = []
    for a in m.atoms:
        i+=1
        if a.name.startswith('D') or a.name.startswith('HV'):
            continue
        ind.append(i)
    indGroup = IndexGroup(name=groupname,ids=ind)
    
    # add index group to the file
    indFile.add_group(indGroup)
    
    # write
    indFile.write(fn=ndxFile)
    

