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

"""This module contains the Model class. It can use
GROMACS routines to read and write structure files. Moreover it
allows to modify structure files in various ways. E.g.:

1. Rename atoms, residues, chains
2. Delete or add atoms, residues and chains

By default, all residues are renumbered from 1. This option can be deactivated
by setting renumber_residues to False.

The Model instance contains:

* model.atoms       -> list of atoms
* model.residues    -> list of residues
* model.chains      -> list of chains
* model.chdic       -> chain dictionary (chdic['A'] returns chain A)

Examples
--------
Basic usage:

    >>> model = Model('input.pdb')

Some useful methods:

    >>> # returns all backbone atoms
    >>> model.fetch_atoms(['CA','N','C'])
    >>> # returns all carbon atoms
    >>> model.fetch_atoms('C',how='byelem')
    >>> # return all atoms except hydrogens
    >>> model.fetch_atoms('H',how='byelem',inv=True)

    >>> # return all ALA,TRP and CYS residues
    >>> model.fetch_residues(['ALA','TRP','CYS'])

    >>> # returns the first 10 residues
    >>> rl = model.residues[:10]
    >>> # return the last residue from chain A
    >>> rl = model.chdic['A'].residues[-1]
    >>> # returns a list with the first residues of each chain
    >>> rl = map(lamda m: m.residues[0], model.chains)
    >>> # remove chain A
    >>> del model['A']
    >>> # write new structure file
    >>> model.write('output.pdb')

"""

import sys
import copy
import numpy as np
from . import _pmx as _p
from . import library
from . import chain
from .atomselection import Atomselection
from .molecule import Molecule
from .atom import Atom
from string import digits


__all__ = ['Model']


class Model(Atomselection):
    """Model Class.

    Parameters
    ----------
    filename : str
        filename of input structure
    pdbline : bool(?)
        what is pdbline?
    renumber_atoms : bool, optional
        renumber all atoms from 1. Default is True.
    renumber_residues : bool, optional
        renumber all residues from 1. In this way, each residue will have a
        unique ID, also across chains. Default is True.
    rename_atoms : bool, optional
        rename atoms so to conform to Gromacs format. Default is False.
    scale_coords : A|nm, optional
        whether to enforce the units of the coordinates to be in A or in nm.
        By default, PDB coordinates are assumed to be in Angstrom (A) and
        GRO coordinates in nanometers (nm). If you read a PDB file but would
        like operate on nm coordinates, then select "nm". Viceversa, if you
        read a GRO file but would like to work on A coordinates, select "A".
        Note that if you read a PDB file in A coordinates and select "A",
        nothing happens (same for GRO and "nm" selection).
    bPDBTER : bool
        whether to recognize TER lines and other chain breaks, e.g.
        discontinuous residue indices(?). Default is True.
    bNoNewID : bool
        whether to assign new chain IDs? If True, new chain IDs starting
        with 'pmx' will be assigned(?). Only relevant if bPDBTER is True.
        Default is True.
    bPDBGAP : bool
        whether search for gaps in the chain to assign new chain IDs.
    bPDBMASS : bool
        whether to guess masses from the atom library (will fail for
        complex atom naming, e.g. ND will not be interpreted as nitrogen)

    Attributes
    ----------
    title : str
        title of model. Default is 'PMX MODEL'.
    filename : str
        filename from which the Model was imported, otherwise None.
    chains : list
        list of Chain instances
    chdic : dict
        dict with chain IDs as keys and Chain instances as values
    residues : list
        list of molecules/residues
    unity : str
        coordinates unit, either 'A' or 'nm'.
    box : 2d array
        3x3 array containing the box vectors. See Gromacs manual, Table 3.1
    moltype : str
        Type of system: protein, dna, rna, or unknown if organic molecule or
        a mix of molecules are in the system.
    """
    def __init__(self, filename=None, pdbline=None, renumber_atoms=True,
                 renumber_residues=True, rename_atoms=False, scale_coords=None,
                 bPDBTER=True, bNoNewID=True, bPDBGAP=False, bPDBMASS=False,
                 **kwargs):

        Atomselection.__init__(self)
        self.title = 'PMX MODEL'
        self.filename = filename
        self.chains = []
        self.chdic = {}
        self.residues = []
        self.name = None  # FIXME/QUESTION: self.name not used anywhere? remove?
        self.id = 0
        self.have_bonds = 0  # FIXME/QUESTION: same as above: never used -> remove?
        self.box = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
        self.unity = 'A'
        for key, val in kwargs.items():
            setattr(self, key, val)

        if filename is not None:
            self.read(filename=filename, bPDBTER=bPDBTER, bNoNewID=bNoNewID, bPDBGAP=bPDBGAP, bPDBMASS=bPDBMASS)
        if pdbline is not None:
            self.__readPDB(pdbline=pdbline)
        if self.atoms:
            self.unity = self.atoms[0].unity
            self.make_chains()
            self.make_residues()
        if self.residues and not self.atoms:
            self.al_from_resl()
            self.make_chains()
            self.make_residues()
        if self.chains and not self.residues:
            self.resl_from_chains()
            self.al_from_resl()
            self.make_chains()
            self.make_residues()
        if self.chdic and not self.chains:
            for key, val in self.chdic.items():
                self.chains.append(val)
            if not self.atoms and not self.residues:
                self.resl_from_chains()
                self.al_from_resl()
        if renumber_atoms is True:
            self.renumber_atoms()
        if renumber_residues is True:
            self.renumber_residues()
        if rename_atoms is True:
            self.rename_atoms_to_gmx()
        if scale_coords is not None:
            if scale_coords == 'A':
                self.nm2a()
            elif scale_coords == 'nm':
                self.a2nm()
            else:
                raise ValueError('unknown unit %s for coordinates' % scale_coords)

        self.assign_moltype()



    def __str__(self):
        s = '< Model: moltype=%s, nchain=%d, nres=%d, natom=%d >' %\
            (self.moltype, len(self.chains), len(self.residues),
             len(self.atoms))
        return s

    def writePIR(self, filename, title=""):
        """Prints sequence to screen in PIR format"""
        fp = open(filename, "w")
        if not title:
            title = '_'.join(self.title.split())
        print('>P1;%s' % title, file=fp)
        print('sequence:::::::::', file=fp)
        for i in range(len(self.chains) - 1):
            print('%s/' % self.chains[i].get_sequence(), file=fp)
        print('%s*' % self.chains[-1].get_sequence(), file=fp)
        fp.close()

    def writeFASTA(self, filename, title=""):
        """Prints sequence to screen in FASTA format"""
        fp = open(filename, "w")
        if not title:
            title = '_'.join(self.title.split())
        if len(self.chains) == 1:
            print('> %s' % title, file=fp)
            print(self.chains[0].get_sequence(), file=fp)
        else:
            for chain in self.chains:
                print('> %s_chain_%s' % (title, chain.id), file=fp)
                print(chain.get_sequence(), file=fp)

    # FIXME/TODO: this function is overwriting the write function inherited from
    # Atomselection. Should we keep only one of the 2?
    # writePIR and writeFASTA could be moved to Atomselection so that all
    # write functions are in one place
    def write(self, fn, title='', nr=1, bPDBTER=False, bAssignChainIDs=False):
        """Writes Model to file. The format is deduced from the filename
        extension given, and available options are '.gro', '.pdb', '.fasta',
        '.pir'.

        Parameters
        ----------
        fn : str
            filename
        title : str, optional
            title of the Model to write in fn
        nr : int, optional
            what is nr?
        bPDBTER : bool, optional
            whether to separate chains with TER entries. Only relevant when
            writing PDB files. Default is False.
        bAssignChainIDs : bool, optional
            whether to write the chains IDs to file. Only relevant when
            writing PDB files. Default is False.
        """
        ext = fn.split('.')[-1]
        if ext == 'pdb':
            self.writePDB(fn, title, nr, bPDBTER, bAssignChainIDs)
        elif ext == 'gro':
            self.writeGRO(fn, title)
        elif ext == 'pir':
            self.writePIR(fn, title)
        elif ext == 'fasta':
            self.writeFASTA(fn, title)
        else:
            print('pmx_Error_> Can only write pdb or gro!', file=sys.stderr)
            sys.exit(1)

    def make_chains(self):
        """Initialises the Chain instances from the input"""
        self.chains = []
        self.chdic = {}
        cur_chain = None
        ch = None
        for atom in self.atoms:
            if atom.chain_id == cur_chain:
                if ch:
                    atom.chain = ch
                    ch.atoms.append(atom)
                else:
                    ch = chain.Chain()
                    cur_chain = atom.chain_id
                    ch.id = atom.chain_id
                    atom.chain = ch
                    ch.atoms.append(atom)
            else:
                if ch:
                    self.chains.append(ch)
                    ch = chain.Chain()
                    cur_chain = atom.chain_id
                    ch.id = cur_chain
                    atom.chain = ch
                    ch.atoms.append(atom)
                else:
                    ch = chain.Chain()
                    cur_chain = atom.chain_id
                    ch.id = atom.chain_id
                    atom.chain = ch
                    ch.atoms.append(atom)

        self.chains.append(ch)
        for ch in self.chains:
            ch.model = self
            idx = ch.id
            self.chdic[idx] = ch

    def make_residues(self):
        """Initialises the Molecule instances from the input"""
        self.residues = []
        for ch in self.chains:
            cur_mol = None
            mol = None
            for atom in ch.atoms:
                if atom.resnr == cur_mol:
                    if mol:
                        mol.atoms.append(atom)
                    else:
                        mol = Molecule()
                        cur_mol = atom.resnr
                        mol.resname = atom.resname
                        mol.id = cur_mol
                        mol.atoms.append(atom)
                else:
                    if mol:
                        mol.model = self
                        mol.chain = ch
                        ch.residues.append(mol)
                        self.residues.append(mol)
                        mol = Molecule()
                        cur_mol = atom.resnr
                        mol.resname = atom.resname
                        mol.id = cur_mol
                        mol.atoms.append(atom)
                    else:
                        mol = Molecule()
                        cur_mol = atom.resnr
                        mol.resname = atom.resname
                        mol.id = cur_mol
                        mol.atoms.append(atom)
            self.residues.append(mol)
            ch.residues.append(mol)
        for r in self.residues:
            r.assign_moltype()
            for atom in r.atoms:
                atom.molecule = r
                atom.model = self
                r.model = self
        for ch in self.chains:
            for r in ch.residues:
                r.chain = ch
                r.chain_id = ch.id

    def __readPDB(self, fname=None, pdbline=None):
        """Reads a PDB file"""
        if pdbline:
            lines = pdbline.split('\n')
        else:
            lines = open(fname, 'r').readlines()
        for line in lines:
            if line[:4] == 'ATOM' or line[:6] == 'HETATM':
                a = Atom().readPDBString(line)
                self.atoms.append(a)
            if line[:6] == 'CRYST1':
                self.box = _p.box_from_cryst1(line)
        self.make_chains()
        self.make_residues()
        self.unity = 'A'
        return self

    def __check_if_gap( self, atC, atN ):
        if atC==None:
            return(False)
        if atN.name != 'N':
            return(False)
        d = atC - atN
        if d > 1.7: # bond 
            return(True)
        return(False)

    def __compareWithoutLastChar(self, str1, str2): 
        if isinstance(str1,int) and isinstance(str2,int): # e.g. 52, 53
            return(False)
 
        if isinstance(str1,int): # e.g. 52, 52A
            if str1==int(str2[0:-1]): # 52, 52A
                return(True)
            if str1==int(str2[0:-1])-1: # 52, 53A
                return(True)
        elif isinstance(str2,int): # e.g. 52A, 52
            if int(str1[0:-1])==str2: # 52A, 52
                return(True)
            if int(str1[0:-1])==str2-1: # 52A, 53
                return(True)
        else: # e.g. 52A, 52B
            if int(str1[0:-1])==int(str2[0:-1]): # 52A, 52B
                return(True)
            if int(str1[0:-1])==int(str2[0:-1])-1: # 52A, 53B
                return(True)
        return(False)

    def _getChainIDs( self, lines ):
        """from the pdb lines extract chain IDs """
        usedChainIDs = []
        for line in lines:
            if (line[:4] == 'ATOM') or (line[:6] == 'HETATM'):
                a = Atom().readPDBString(line)
                if( a.chain_id not in usedChainIDs ):
                    usedChainIDs.append( a.chain_id )
        return( usedChainIDs )
        
    # TODO: make readPDB and readPDBTER a single function. It seems like
    # readPDBTER is more general PDB reader?
    def __readPDBTER(self, fname=None, pdbline=None, bNoNewID=True, bPDBGAP=False, bPDBMASS=False):
        """Reads a PDB file with more options than __readPDB ?"""
        if pdbline:
            lines = pdbline.split('\n')
        else:
            lines = open(fname, 'r').readlines()

        chainIDstringInit = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnoprstuvwxyz123456789'
        chainIDstring = copy.deepcopy(chainIDstringInit)
        bNewChain = True
        chainID = ' '
        prevID = ' '
        prevAtomName = ' '
        prevResID = 0
        prevResName = ' '
        usedChainIDs = []
        atomcount = 1
        prevCatom = None

        # get already used chain IDs
        usedChainIDs = self._getChainIDs( lines )

        for line in lines:
            if 'TER' in line:
                bNewChain = True
            if (line[:4] == 'ATOM') or (line[:6] == 'HETATM'):
                a = Atom().readPDBString(line, origID=atomcount)
                atomcount += 1
                # identify chain change by ID (when no TER is there)
                if (a.chain_id != prevID):
                    bNewChain = True
                if (self.__check_if_gap( prevCatom,a )==True and bPDBGAP==True):
                    bNewChain = True
                if a.resname=='ACE' and prevResName!='ACE':
                    bNewChain = True
                if (a.resnr != prevResID):
                    try:
                        if self.__compareWithoutLastChar(prevResID,a.resnr)==True: # there are some special cases where residues are named, e.g. 52, 52A, 52B, ...
                            bNewChain = False
                        elif a.resnr != prevResID+1:
                            bNewChain = True
                        if (prevAtomName == 'OC2') or (prevAtomName == 'OXT') or (prevAtomName == 'OT2'):
                            bNewChain = True
                        if (prevAtomName == 'HH33') and ((prevResName=='NME') or (prevResName=='NAC') or (prevResName=='CT3')): # NME cap
                            bNewChain = True
                        # do not assign new chain IDs for waters, ions 
                        if (a.resname=='WAT') or (a.resname=='SOL') or (a.resname=='TIP3') or (a.resname=='HOH') \
                           or (a.resname=='NA') or (a.resname=='CL') \
                           or (a.resname=='NaJ') or (a.resname=='Na') or (a.resname=='Cl') or (a.resname=='K') or (a.resname=='KJ') \
                           or (a.resname=='MG') or (a.resname=='Mg') or (a.resname=='CA') or (a.resname=='Ca') or (a.resname=='CaJ') \
                           or (a.resname=='ZN') or (a.resname=='Zn'): # add other ions when needed
                            bNewChain = False
                            chainID = ''
                    except TypeError:
                        bNewChain = False
                        chainID = ''
                # the option below overwrites all the others:
                # when a non-empty ID of the next chain is the same as the previous, then do no change
                if (a.chain_id!='') and (a.chain_id!=' ') and (a.chain_id==prevID):
                    bNewChain = False
                prevID = a.chain_id
                prevResID = a.resnr
                prevAtomName = a.name
                prevResName = a.resname
                if a.name == 'C':
                    prevCatom = a
                if bNewChain==True:
                    if ((a.chain_id==' ') or (a.chain_id==chainID) or (a.chain_id in usedChainIDs) and bNoNewID==False):
                        # find a new chain id
                        bFound = False
                        while bFound==False:
                            if len(chainIDstring)==0: # used up all the IDs
                                chainIDstring = copy.deepcopy(chainIDstringInit)
                                chainID = "pmxX"
                                break
                            foo = chainIDstring[0]
                            chainIDstring = chainIDstring.lstrip(chainIDstring[0])
                            if foo not in usedChainIDs:
                                bFound=True
                                chainID = foo
                                if bNoNewID==True:
                                    chainID = "pmx"+foo
                                usedChainIDs.append(chainID)
                    else:
                        chainID = a.chain_id
                        usedChainIDs.append(chainID)
                a.chain_id = chainID
                self.atoms.append(a)
                bNewChain = False
            if line[:6] == 'CRYST1':
                self.box = _p.box_from_cryst1(line)

        ##### now fix chain IDs that have been newly created #####
        newChainDict = {}
        chainIDstring = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnoprstuvwxyz123456789'
        for a in self.atoms:
            # chain with a new ID
            if 'pmx' in a.chain_id:
                # this ID has already been encountered
                if a.chain_id in newChainDict.keys():
                    a.chain_id = newChainDict[a.chain_id]
                # ID not yet encountered
                else:
                    # find a suitable ID
                    bFound = False
                    while bFound==False:
                        if len(chainIDstring)==0: # used up all the IDs
                            chainIDstring = copy.deepcopy(chainIDstringInit)
                            newChainDict[a.chain_id] = "X"
                            a.chain_id = "X"
                            break
                        foo = chainIDstring[0]
                        chainIDstring = chainIDstring.lstrip(chainIDstring[0])
                        # found
                        if foo not in usedChainIDs:
                            bFound=True
                            usedChainIDs.append(foo)
                            newChainDict[a.chain_id] = foo
                            a.chain_id = foo

        self.make_chains()
        self.make_residues()
        self.unity = 'A'

        if bPDBMASS==True:
            assign_masses_to_model( self )

        return self

    def __readGRO(self, filename):
        """Reads a GRO file"""
        l = open(filename).readlines()
        # first line is name/comment
        name = l[0].rstrip()
        self.title = name
        # next line is number of atoms
        natoms = int(l[1])
        atoms_parsed = 0
        while atoms_parsed != natoms:
            line = l[atoms_parsed+2]
            resid = int(line[:5])
            resname = line[5:9].strip()
            name = line[10:15].strip()
            idx = int(line[15:20])
            rest = line[20:].split()
            assert len(rest) in [3, 6]
            x = float(rest[0])
            y = float(rest[1])
            z = float(rest[2])
            coords = [x, y, z]
            if len(rest) == 6:
                vx = float(rest[3])
                vy = float(rest[4])
                vz = float(rest[5])
                vel = [vx, vy, vz]
            else:
                vel = [0, 0, 0]
            a = Atom(id=idx, name=name, resname=resname,
                     resnr=resid, x=coords, v=vel, unity='nm')
            a.get_symbol()
            self.atoms.append(a)
            atoms_parsed += 1
        box_line = [float(i) for i in l[-1].split()]
        assert len(box_line) in [3, 9]
        box = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
        box[0][0] = box_line[0]
        box[1][1] = box_line[1]
        box[2][2] = box_line[2]
        if len(box_line) == 3:
            box[0][1] = 0
            box[0][2] = 0
            box[1][0] = 0
            box[1][2] = 0
            box[2][0] = 0
            box[2][1] = 0
        else:
            box[0][1] = box_line[3]
            box[0][2] = box_line[4]
            box[1][0] = box_line[5]
            box[1][2] = box_line[6]
            box[2][0] = box_line[7]
            box[2][1] = box_line[8]
        self.box = box
        self.make_chains()
        self.make_residues()
        self.unity = 'nm'
        return self

    def assign_moltype(self):
        """Identifies what type of molecule the Model is: protein, dna, or rna.
        If it is a mix, or if it is an organic molecule, "unknown" is
        assigned to self.moltype.
        """
        residues = set([r.resname for r in self.residues])

        # do not consider water and ions
        residues -= library._water
        residues -= library._ions

        # determine type
        # if 'residues' is not empty AND it's a subset of 'library._protein_residues'
        if bool(residues) and residues <= library._protein_residues:
            self.moltype = 'protein'
        elif bool(residues) and residues <= library._dna_residues:
            self.moltype = 'dna'
        elif bool(residues) and residues <= library._rna_residues:
            self.moltype = 'rna'
        else:
            self.moltype = 'unknown'

    def read(self, filename, bPDBTER=False, bNoNewID=True, bPDBGAP=False, bPDBMASS=False):
        """PDB/GRO file reader.

        Parameters
        ----------
        filename : str
            name of input file
        bPDBTER : bool, optional
            whether the file contains TER records?. Default is False.
        bNoNewID : bool, optional
            whether to assign new chain IDs. If True, new chain IDs starting
            with 'pmx' will be assigned(?). Only relevant if bPDBTER is also
            True. Default is True.
        bPDBGAP : bool
            whether search for gaps in the chain to assign new chain IDs.
        """
        ext = filename.split('.')[-1]
        if ext == 'pdb':
            if bPDBTER is True:
                return self.__readPDBTER(fname=filename,
                                         pdbline=None,
                                         bNoNewID=bNoNewID, bPDBGAP=bPDBGAP, bPDBMASS=bPDBMASS)
            else:
                return self.__readPDB(fname=filename)
        elif ext == 'gro':
            return self.__readGRO(filename)
        else:
            raise IOError('ERROR: Can only read pdb or gro!')

    def renumber_residues(self):
        """Renumbers all residues from 1."""
        for i, res in enumerate(self.residues):
            res.set_orig_resid(res.id)
            res.set_resid(i+1)

    # TODO/FIXME: should add/remove/append atoms all be only once in
    # Atomselection? At the moment they are repeated in Molecule, Chain,
    # and Model
    def remove_atom(self, atom):
        """Removes an Atom instance.

        Parameters
        ----------
        atom : Atom
            Atom instance to remove
        """
        m = atom.molecule
        m.remove_atom(atom)

    def remove_residue(self, residue, renumber_atoms=True, renumber_residues=True):
        """Removes a Molecule/residue instance.

        Parameters
        ----------
        residue : Molecule
            Molecule instance to remove
        renumber_atoms : bool, optional
            whether to renumber the atoms of the Model after removing the
            residue. Default is True.
        renumber_residues : bool, optional
            whether to renumber the residues of the Model after removing the
            residue. Default is True.
        """
        ch = residue.chain
        ch.remove_residue(residue, renumber_atoms=renumber_atoms,
                          renumber_residues=renumber_residues)

    def remove_chain(self, key):
        """Removes a Chain instance given the chain ID.

        Parameters
        ----------
        key : str
            ID of the chain to remove
        """
        if key not in self.chdic:
            print('No chain %s to remove....' % key)
            print('No changes applied.')
            return
        for ch in self.chains:
            if ch.id == key:
                idx = self.chains.index(ch)
                del self.chains[idx]
                del self.chdic[key]
        self.resl_from_chains()
        self.al_from_resl()
        self.renumber_residues()
        self.renumber_atoms()

    def __delitem__(self, key):
        self.remove_chain(key)

    def insert_residue(self, pos, res, chain_id):
        """Inserts a residue in Model.

        Parameters
        ----------
        pos : int
            position/index where to insert the residue
        res : Molecule
            Molecule instance to insert containig the residue of interest
        chain_id : str
            ID of the chain where to insert the residue
        """
        ch = self.chdic[chain_id]
        ch.insert_residue(pos, res)

    def replace_residue(self, residue, new, bKeepResNum=False):
        """Replaces a residue.

        Parameters
        ----------
        residue : Molecule
            residue to replace
        new : Molecule
            residue to insert
        bKeepResNum : bool, optional
            whether to keep residue ID of the residue that is inserted.
            Default is False
        """
        ch = residue.chain
        ch.replace_residue(residue, new, bKeepResNum)

    def insert_chain(self, pos, new_chain):
        """Inserts a Chain in Model.

        Parameters
        ----------
        pos : int
            index where to insert the chain within the list of chains
        new_chain : Chain
            instance of the Chain to insert in Model
        """
        if new_chain.id in self.chdic:
            print('Chain identifier %s already in use!' % new_chain.id)
            print('Changing chain identifier to 0')
            new_chain.set_chain_id('0')
        self.chains.insert(pos, new_chain)
        self.resl_from_chains()
        self.al_from_resl()
        self.make_chains()
        self.make_residues()
        self.renumber_atoms()
        self.renumber_residues()

    def append(self, new_chain):
        """ we assume chain is a Chain"""
        idx = len(self.chains)
        self.insert_chain(idx, new_chain)

    def fetch_residue(self, idx, chain=None):
        """Get a residue based on its index, or index and chain.

        Parameters
        ----------
        idx : int
            ID of the residue to fetch.
        chain : str, optional
            chain ID of the residue. This is needed when you have multiple
            chains and you have not renumbered the residues.

        Returns
        -------
        residue : Molecule
            Molecule instance of the residue found.
        """

        #########################
        # generate some residue id lists
        valid_resids = []
        for r in self.residues:
            if isinstance(r.id,str):
                valid_resids.append(r.id.replace(" ",""))
            else:
                valid_resids.append(r.id)

        if chain is not None:        
            valid_chresids = []
            for r in self.chdic[chain].residues:
                if isinstance(r.id,str):
                    valid_chresids.append(r.id.replace(" ",""))
                else:
                    valid_chresids.append(r.id)
        ##########################


        # check idx is a valid selection
        if idx not in valid_resids:#[r.id for r in self.residues]:
            raise ValueError('resid %s not found in Model residues' % idx)

        if chain is None:
            # check selection is unique
            if [r.id for r in self.residues].count(idx) > 1:
                raise ValueError('idx choice %s results in non-unique selection' % idx)
        else:
            # check chain is a valid selection
            if chain not in [c.id for c in self.chains]:
                raise ValueError('chain ID "%s" not found in Model chains' % chain)
            # check idx+chain is a valid selection
            if idx not in valid_chresids:#[r.id for r in self.chdic[chain].residues]:
                raise ValueError('resid %s not found in chain "%s"' % (idx, chain))
            # check selection is unique
            if [r.id for r in self.chdic[chain].residues].count(idx) > 1:
                raise ValueError('idx choice %s for chain "%s" results in non-unique selection' % (idx, chain))

        # then find and return the residue
        if chain is None:
            for r in self.residues:
                if isinstance(r.id,int):
                    if r.id == idx:
                        return r
                else:
                    if r.id.replace(" ","")==idx:
                        return r
        elif chain is not None:
            for r in self.chdic[chain].residues:
                if isinstance(r.id,int):
                    if r.id == idx:
                        return r
                else:
                    if r.id.replace(" ","")==idx:
                        return r
            

    def fetch_residues(self, key, inv=False):
        """Gets residues using a list of residue names.

        Parameters
        ----------
        key : str or list
            resname
        inv : bool
            invert selection. Default is False. If True, it finds all residues
            with resnames different from those in key.

        Returns
        -------
        residues : list
            list of Molecule instances with the residues
        """
        if not hasattr(key, "append"):
            key = [key]
        result = []
        if not inv:
            for r in self.residues:
                if r.resname in key:
                    result.append(r)
        else:
            for r in self.residues:
                if r.resname not in key:
                    result.append(r)
        return result

    def al_from_resl(self):
        self.atoms = []
        for r in self.residues:
            for atom in r.atoms:
                self.atoms.append(atom)

    def resl_from_chains(self):
        self.residues = []
        for ch in self.chains:
            for r in ch.residues:
                self.residues.append(r)

    def copy(self):
        return copy.deepcopy(self)

    def get_mol2_types(self):
        if self.atoms[0].symbol == '':
            self.get_symbol()
        for ch in self.chains:
            ch.get_mol2_types()

    def get_mol2_resname(self):
        for ch in self.chains:
            ch.get_mol2_resname()

    def get_nterms(self):
        nter = []
        for ch in self.chains:
            first = ch.residues[0]      # first residue
            if first.resname in library._one_letter.keys():
                nter.append(first)
        return nter

    def get_cterms(self):
        cter = []
        for ch in self.chains:
            last = ch.residues[-1]      # last residue
            if last.resname in library._one_letter.keys():
                cter.append(last)
        return last

    def rename_atoms(self):
        for c in self.chains:
            c.rename_atoms()

    def residue(self, idx):
        return self.residues[idx-1]

    def chain(self, iden):
        return self.chdic[iden]


# ==============================================================================
#                                  Functions
# ==============================================================================
def merge_models(*args):
    '''Merges the atoms from all Model objects in the list provided. Atoms will
    be merged based on the order of models in the list.

    Parameters
    ----------
    *args :
        variable length argument containing Model objects.

    Returns
    -------
    m : Model
        new Model object containing all the Models in the list.

    Examples
    --------
    >>> newmodel = merge_models(model1, model2, model3)
    >>> newmodel = merge_models(protein, ligand, cofactor, ions)
    '''
    model = Model()

    if not all(m.unity == args[0].unity for m in args):
        raise ValueError('the Model objects provided do not have the same '
                         'units - convert units to that they are compatible')

    for m in args:
        # operate on a deep copy of the model, otherwise changes in the atoms
        # of the merged model will be reflected in the parent models too
        m_ = copy.deepcopy(m)
        for a in m_.atoms:
            model.atoms.append(a)

    model.unity = model.atoms[0].unity
    model.make_chains()
    model.make_residues()
    model.assign_moltype()
    return model


def assign_masses_to_model(model, topology=None):
    '''Assigns masses to the Model atoms given the ones present in the Topology.

    Parameters
    ----------
    model : Model
        Model object of the molecule.
    topology : Topology
        Topology object of the same molecule. When no topology provided, standard library masses are used.
    '''
    
    if topology!=None:
        for ma, ta in zip(model.atoms, topology.atoms):
            if ma.name != ta.name:
                raise ValueError('mismatch of atom names between Model and '
                                 'Topology objects provided')
            ma.m = ta.m
    else:
        for a in model.atoms:
            aname = a.name
            aname = aname.upper()
            aname = aname.translate(str.maketrans('','',digits))
            if aname not in library._atommass.keys():
                a.m = 0.0
            else:
                a.m = library._atommass[aname]

def double_box(m1, m2, r=2.5, d=1.5, bLongestAxis=False, verbose=False):
    '''Places two structures (two Model objects) into a single box.
    The box is a rectangular cuboid in which the two structures are placed in
    such a way to minimise the box volume.

    Parameters
    ----------
    m1 : Model
        first structure.
    m2 : Model
        first structure.
    r : float
        distance between the two structures (nm).
    d : float
        distance to the box wall (nm).
    bLongestAxis : bool
        whether to just place structures along the
        longest axis, rather then optimising the placement and minimising
        the volume.
    verbose : bool
        whether to print out information about the new box or not.

    Returns
    -------
    mout : Model
        new Model object where the two input structures are placed within a
        single box.
    '''

    def _translate(m, v, fact=1.0):
        for a in m.atoms:
            a.x[0] = a.x[0] + fact*v[0]
            a.x[1] = a.x[1] + fact*v[1]
            a.x[2] = a.x[2] + fact*v[2]

    def _get_mass(a):
        aname = ''.join(i for i in a.name if not i.isdigit())
        if aname.startswith('Br') or aname.startswith('BR'):
            return(library._atommass['BR'])
        elif aname.startswith('Cl') or aname.startswith('CL'):
            return(library._atommass['CL'])
        elif aname.startswith('D'):
            return(library._atommass[aname[1]])
        else:
            return(library._atommass[aname[0]])
        return(1.0)

    def _principal_axes(m):
        tensor = np.matrix([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
        for a in m.atoms:
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            mass = _get_mass(a)
            tensor[0,0] += mass*(np.power(a.x[1], 2) + np.power(a.x[2], 2))
            tensor[1,1] += mass*(np.power(a.x[0], 2) + np.power(a.x[2], 2))
            tensor[2,2] += mass*(np.power(a.x[0], 2) + np.power(a.x[1], 2))
            tensor[0,1] -= mass*a.x[0]*a.x[1]
            tensor[1,0] = tensor[0, 1]
            tensor[0,2] -= mass*a.x[0]*a.x[2]
            tensor[2,0] = tensor[0, 2]
            tensor[1,2] -= mass*a.x[1]*a.x[2]
            tensor[2,1] = tensor[1, 2]
        evals, evecs = np.linalg.eig(tensor)
        idx = evals.argsort()[::1] # sort descending, because later the rotation matrix will be calculated as transpose
        evals = evals[idx]
        evecs = evecs[:, idx]
        # check if one axis needs to be flipped
        crossprod = np.cross(np.transpose(evecs[:,0]), np.transpose(evecs[:,1]))
        dotprod = np.dot(np.transpose(evecs[:,2]), np.transpose(crossprod))
        if dotprod < 0.0:
            evecs[:,2] *= -1.0
        rotmat = np.transpose(evecs)
        return(rotmat)

    def _rotate(atoms, R):
        for atom in atoms:
            x_old = list(map(lambda x: x, atom.x))
            for i in range(3):
                atom.x[i] = 0.0
                for j in range(3):
                    atom.x[i] += x_old[j]*R[i,j]

    def _get_com_radius(m):
        com = [0.0, 0.0, 0.0]
        radius = 0.0
        mass = 0.0

        # com
        for a in m.atoms:
            a.a2nm()
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            aname = a.name[0]
            if aname.isdigit():
                aname = a.name[1]
            amass = _get_mass(a)

            com[0] += a.x[0]*amass
            com[1] += a.x[1]*amass
            com[2] += a.x[2]*amass
            mass += amass

        com[0] /= mass
        com[1] /= mass
        com[2] /= mass

        # radius
        for a in m.atoms:
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            foo = (a.x[0]-com[0])**2 + (a.x[1]-com[1])**2 + (a.x[2]-com[2])**2
            if foo > radius:
                radius = foo
        radius = np.sqrt(radius)

        return(com, radius)

    def _get_rect_dist(m):
        minx = 9999.999
        maxx = -9999.999
        miny = 9999.999
        maxy = -9999.999
        minz = 9999.999
        maxz = -9999.999
        for a in m.atoms:
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            if a.x[0] < minx:
                minx = a.x[0]
            if a.x[0] > maxx:
                maxx = a.x[0]
            if a.x[1] < miny:
                miny = a.x[1]
            if a.x[1] > maxy:
                maxy = a.x[1]
            if a.x[2] < minz:
                minz = a.x[2]
            if a.x[2] > maxz:
                maxz = a.x[2]
        a = maxx-minx
        b = maxy-miny
        c = maxz-minz
        return (a, b, c)

    def _get_extr(m, i=0):
        minx = 9999.999
        maxx = -9999.999
        for a in m.atoms:
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            if a.x[i] < minx:
                minx = a.x[i]
            if a.x[i] > maxx:
                maxx = a.x[i]
        return(minx, maxx)

    def _translate_sticking_out(m, l, i=0):
        minx = 0.0
        maxx = 0.0
        for a in m.atoms:
            if ('HOH' in a.resname) or ('SOL' in a.resname) or ('WAT' in a.resname):
                continue
            if a.x[i] < minx:
                minx = a.x[i]
            if (a.x[i]-l) > maxx:
                maxx = a.x[i]-l
        return(-1.0*(minx+maxx))

    # ====
    # Main
    # ====

    # com and radii
    com1, rad1 = _get_com_radius(m1)
    com2, rad2 = _get_com_radius(m2)

    # remove COM from the systems
    _translate(m1, com1, fact=-1.0)
    _translate(m2, com2, fact=-1.0)

    # if needed, calculate principal axes
    if bLongestAxis:
        princRot1 = _principal_axes(m1)
        princRot2 = _principal_axes(m2)
        _rotate(m1.atoms, princRot1)
        _rotate(m2.atoms, princRot2)
        a_rect1, b_rect1, c_rect1 = _get_rect_dist(m1)
        a_rect2, b_rect2, c_rect2 = _get_rect_dist(m2)

    # estimate cube's edge for the larger structure
    a_cube = 0.0
    if rad1 > rad2:
        a_cube = 2*rad1+2*d
    else:
        a_cube = 2*rad2+2*d
    if bLongestAxis is True:
        a_rect = a_rect1 + a_rect2 + r + 2.0*d
        b_rect = b_rect1 + b_rect2 + 2.0*d
        c_rect = c_rect1 + c_rect2 + 2.0*d
        if verbose is True:
            print("Cuboid dimensions (nm): ", np.round(a_rect), np.round(b_rect), np.round(c_rect, 2))
    else:
        if verbose is True:
            print("Cube's edge (nm): ", np.round(a_cube, 2))

    # cube's diagonal
    if bLongestAxis is True:
        d_cube = np.sqrt(np.power(a_rect, 2) + np.power(b_rect, 2) + np.power(c_rect, 2))
        if verbose is True:
            print("Cuboid's diagonal (nm): ", np.round(d_cube, 2))
    else:
        d_cube = a_cube*np.sqrt(3.0)
        if verbose is True:
            print("Cube's diagonal (nm): ", np.round(d_cube, 2))

    # check if cube is enough
    if bLongestAxis is True:
        a_box = a_rect
        b_box = b_rect
        c_box = c_rect
    else:
        dist_cube = d_cube - 2.0*rad1 - 2.0*rad2 - r - 2.0*d
        a_box = a_cube
        b_box = a_cube
        c_box = a_cube
        if dist_cube > 0.0:
            if verbose is True:
                print("Having a cube with the diagonal ", np.round(d_cube, 2), "nm is good enough")
        else:
            if verbose is True:
                print("Need to extend the cube")
            d_rect = 2.0*rad1 + 2.0*rad2 + r + 2.0*d
            delta_a_cube = -a_cube + np.sqrt(d_rect**2 - 2.0*a_cube**2)
            a_box = a_cube + delta_a_cube
            if verbose is True:
                print("Rectangle's edge (nm): ", np.round(a_box, 2))

    # translate the larger structure to the middle of the box
    if rad1 > rad2:
        # translate smaller to (0,b/2,c/2) and larger to (maxx2+dist+abs(minx2),b/2,c/2)
        if bLongestAxis is True:
            _translate(m2, [0.0, b_box/2.0, c_box/2.0])
            minx1, maxx1 = _get_extr(m1, 0)
            minx2, maxx2 = _get_extr(m2, 0)
            _translate(m1, [maxx2+r+abs(minx1), b_box/2.0, c_box/2.0])
        else:
            _translate(m1, [a_box/2.0, b_box/2.0, c_box/2.0])
    else:
        if bLongestAxis is True:
            _translate(m1, [0.0, b_box/2.0, c_box/2.0])
            minx1, maxx1 = _get_extr(m1, 0)
            minx2, maxx2 = _get_extr(m2, 0)
            _translate(m2, [maxx1 + r + abs(minx2), b_box/2.0, c_box/2.0])
        else:
            _translate(m2, [a_box/2.0, b_box/2.0, c_box/2.0])

    # create an output
    mout = m1
    chainList = []
    for ch in m1.chains:
        if ch.id not in chainList:
            chainList.append(ch.id)
    chainIDstring = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    chainIDchanges = {}
    for ch in m2.chains:
        if ch.id in chainList:
            # pick a new chain ID
            bFound = False
            while bFound is False:
                foo = chainIDstring[0]
                chainIDstring = chainIDstring.lstrip(chainIDstring[0])
                if foo not in chainList:
                    bFound = True
            chainList.append(foo)
            chainIDchanges[ch.id] = foo
    for a in m2.atoms:
        chID = a.chain_id
        if chID in chainIDchanges:
            a.chain_id = chainIDchanges[chID]
    mout.atoms.extend(m2.atoms)

    # create the box
    mout.box = [[a_box,0.0,0.0], [0.0,b_box,0.0], [0.0,0.0,c_box]]

    # determine translation exactly by checking how much the atoms are sticking out
    xtransl = _translate_sticking_out(mout, a_box, 0)
    ytransl = _translate_sticking_out(mout, b_box, 1)
    ztransl = _translate_sticking_out(mout, c_box, 2)
    # increase translation by a factor 1.01 to move a bit further from walls
    _translate(mout, [1.01*xtransl, 1.01*ytransl, 1.01*ztransl])

    # volume
    V = a_box*b_box*c_box
    if verbose is True:
        print("Volume: ", np.round(V, 2), "nm^3")

    return mout
