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

Examples
--------

Basic usage:

    >>> model = Model().read(args['-f'])             # read structure file
    >>> atom = model.atoms[0]                        # the first atom
    >>> a = Atom(name='DUM',x = [1,2,3], bfac = 20.) # create an atom

Geometric measurements:

    >>> d = atom1 - atom2                       # calculate distance
    >>> a = atom1.angle(atom2,atom3)            # calculate angle
    >>> dih = atom1.dihedral(atom2,atom3,atom4) # calculate dihedral
    >>> print atom                              # print atom in PDB format

"""

import copy
from numpy import pi
from . import library
from . import _pmx as _p
from .library import pdb_format, pdb_format2

__all__ = ['Atom']


class Atom:
    """Atom Class.

    Parameters
    ----------
    line : str, optional
        input line in PDB format. Default is None.
    mol2line : str, optional
        input line in MOL2 format.  Default is None.


    Attributes
    ----------
    id : int
        atom id
    name : str
        atom name
    resname : str
        name of residue atom is part of
    chain_id : str
        ID of the chain atom belongs to
    x : array
        atom coordinates
    occ : float
        occupancy
    bfac : float
        B-factor
    symbol : str
        chemical element of the atom
    unity : str
        unit of the coordinates, either 'A' or 'nm'. Default is 'A'.
    """
    def __init__(self, line=None, mol2line=None, **kwargs):

        self.race = 'ATOM  '
        self.id = 0
        self.orig_id = 0
        self.code = 0
        self.name = ''
        self.altloc = ''
        self.resname = ''
        self.chain_id = ' '
        self.x = [0, 0, 0]
        self.occ = 1.
        self.bfac = 0.
        self.vdw = 0.
        self.vdw14 = 0.
        self.atype = ''
        self.hyb = ''
        self.symbol = ''
        self.bonds = []
        self.b13 = []
        self.b14 = []
        self.connected = []
        self.neighbors = []
        self.order = 0
        self.boxid = 0
        self.molecule = None
        self.chain = None
        self.model = None
        self.cgnr = 0.
        self.v = [0, 0, 0]
        self.f = []
        self.m = 0.
        self.q = 0.
        self.mB = 0.
        self.qB = 0.
        self.type = ''
        self.typeB = ''
        self.resnr = 0
        self.grpnr = ' '
        self.atomtype = ''
        self.atomtypeB = ''
        self.ptype = ''
        self.long_name = ''
        self.unity = 'A'
        for key, val in kwargs.items():
            setattr(self, key, val)
        if line is not None:
            self.readPDBString(line)
        if mol2line is not None:
            self.read_mol2_line(mol2line)

    def __str__(self):
        """Prints the Atom in PDB format"""
        if self.unity == 'nm':
            coords = list(map(lambda x: x*10, self.x))
        else:
            coords = self.x
        if len(self.resname) < 4:
            resname = self.resname+' '  # [0:3]
        else:
            resname = self.resname
        if len(self.name) == 1:
            name = ' '+self.name+'  '
        elif len(self.name) == 2:
            if self.name[0].isdigit():
                name = self.name+'  '
            else:
                name = ' '+self.name+' '
        elif len(self.name) == 3:
            if self.name[0].isdigit():
                name = self.name+' '
            else:
                name = ' '+self.name
        else:
            name = self.name
        idx = self.id % 100000
        try:
            resid = self.resnr % 10000
        except:
            resid = str(self.resnr)
        try:
            s = pdb_format % (self.race, idx, name, self.altloc,
                              resname, self.chain_id, resid,
                              coords[0], coords[1],
                              coords[2], self.occ, self.bfac, self.symbol)
        except:
            s = pdb_format2 % (self.race, idx, name, self.altloc,
                               resname, self.chain_id, resid,
                               coords[0], coords[1], coords[2],
                               self.occ, self.bfac)
        return s

    def __sub__(self, other):
        """ Overloading of the '-' operator for using
        atom1-atom2 instead of atom1.dist(atom2)"""
        return self.dist(other)

    def readPDBString(self, line, origID=0):
        """PDB String to Atom"""

        self.race = line[0:6]
        self.id = int(line[7:11])
        self.orig_id = origID
        self.name = line[12:16].strip()
        self.altloc = line[16]
        self.resname = line[17:21].strip()
        self.chain_id = line[21]
        try:
            self.resnr = int(line[22:27])
        except:
            self.resnr = line[22:27].lstrip().rstrip()  # contains insertion code

        self.x = [float(line[30:38]),
                  float(line[39:46]),
                  float(line[47:54])]
        try:
            self.occ = float(line[55:60])
        except:
            self.occ = 1.
        try:
            self.bfac = float(line[61:66])
        except:
            self.bfac = 0.
        try:
            self.symbol = line[70:73].strip()
        except:
            self.symbol = None
        self.unity = 'A'
        if not self.symbol:
            self.get_symbol()
        return self

    def read_mol2_line(self, line):
        entr = line.split()
        if len(entr) == 9:
            self.id = int(entr[0])
            self.name = entr[1]
            self.x[0] = float(entr[2])
            self.x[1] = float(entr[3])
            self.x[2] = float(entr[4])
            self.atype = entr[5]
            self.resnr = int(entr[6])
            self.resname = entr[7]
            self.q = float(entr[8])
            self.unity = 'A'
            self.symbol = self.atype.split('.')[0]
        elif len(entr) == 10:
            self.id = int(entr[0])
            self.name = entr[1]
            self.x[0] = float(entr[3])
            self.x[1] = float(entr[4])
            self.x[2] = float(entr[5])
            self.atype = entr[6]
            self.resnr = int(entr[7])
            self.resname = entr[8]
            self.q = float(entr[9])
            self.unity = 'A'
            self.symbol = self.atype.split('.')[0]
        else:
            print(line)
            raise ValueError('Error: Cannot convert line to atom')
        return self

    def make_long_name(self):
        """Make extended name to determine element and order.
        """
        # check for aliases first
        ali = library._aliases
        if self.resname in ali and self.name in ali[self.resname]:
            name = ali[self.resname][self.name]
        else:
            name = self.name.strip()
        if len(name) == 4:
            if not name[0].isdigit():
                tmp = name[3]
                name = tmp+name[:-1]
        elif len(name) == 3:
            if name[0].isdigit():
                name += ' '
            else:
                name = ' '+name
        elif len(name) == 2:
            if name[1].isdigit() and name[0] == 'H' and \
                   self.resname in library._aacids:
                name = name[1]+name[0]+'  '
            else:
                name = ' ' + name + ' '
        elif len(name) == 1:
            name = ' '+name+'  '
        self.long_name = name

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

    def get_symbol(self):
        """Set atom element"""
        if self.long_name == '':
            self.make_long_name()
        if self.resname in library._protein_residues or \
           self.resname in library._nucleic_acids:
            self.symbol = self.long_name[1]
        elif self.resname in library._ions:
            self.symbol = self.resname.upper()
            if len(self.symbol) > 2:
                self.symbol = self.symbol[:2]
        elif self.resname in library._water:
            self.symbol = self.name[0]
        else:
            # unknown residue
            c1 = self.long_name[1]
            c2 = self.long_name[2]
            if c1 == 'C':
                self.symbol = 'C'
                if c2.upper() == 'L':
                    self.symbol = 'CL'
            elif c1 == 'O':
                self.symbol = 'O'
            elif c1 == 'N':
                self.symbol = 'N'
            elif c1 == 'H':
                self.symbol = 'H'
            elif c1 == 'S':
                self.symbol = 'S'
            elif c1 == 'P':
                self.symbol = 'P'
            elif c1 == 'B':
                if c2.upper() == 'R':
                    self.symbol = 'BR'
            elif c1 == 'F':
                if c2.upper() == 'E':
                    self.symbol = 'FE'
                else:
                    self.symbol = 'F'
            elif c1 == 'I':
                self.symbol = 'I'
            elif c1 == 'D':
                self.symbol = 'D'
            elif c1 == 'M':
                self.symbol = 'D'
            else:
                self.symbol = 'UN'

    def get_order(self):
        """Get the order (number of bonds to mainchain)"""
        if self.long_name == '':
            self.make_long_name()
        if self.symbol == '':
            self.get_symbol()
        if self.resname not in library._protein_residues:
            print('Sorry, implemented for proteins only')
            return

        el = self.symbol
        x = self.long_name[2]
        if self.name in ['C', 'CA', 'N']:
            self.order = 0
        elif self.name in ['H', 'H1', 'H2', 'H3', '1H', '2H',
                           '3H', 'O', 'O1',
                           'O2', 'OC1', 'OC2', 'OXT',
                           'OT1', 'OT2']:
            self.order = 1
        else:
            if el != 'H':
                if x == 'B':
                    self.order = 1
                elif x == 'G':
                    self.order = 2
                elif x == 'D':
                    self.order = 3
                elif x == 'E':
                    self.order = 4
                elif x == 'Z':
                    self.order = 5
                elif x == 'H':
                    self.order = 6
            elif el == 'H':
                if x == 'A':
                    self.order = 1
                elif x == 'B':
                    self.order = 2
                elif x == 'G':
                    self.order = 3
                elif x == 'D':
                    self.order = 4
                elif x == 'E':
                    self.order = 5
                elif x == 'Z':
                    self.order = 6
                elif x == 'H':
                    self.order = 7

    # ==============
    # Public Methods
    # ==============
    def dist(self, other):
        """Calculates the distance between two atoms. This function is also when
        subtracting two atom instances, as in the example.

        Examples
        --------
        >>> dist = atom1.dist(atom2)
        >>> dist = atom1 - atom2

        Returns
        -------
        dist : float
            distance between two Atom instances.
        """
        return _p.dist(self.x, other.x)

    def dist2(self, other):
        """Calculates the squared distance between two atoms.

        Examples
        --------
        >>> dist = atom1.dist2(atom2)

        Returns
        -------
        dist2 : float
            squared of the distance between two Atom instances
        """
        return _p.dist2(self.x, other.x)

    def translate(self, v):
        """Translates the position of the atom.

        Parameters
        ----------
        v : array
            vector of length 3 containing the values by which to translate the
            atom in the x, y, and z dimensions.
        """
        self.x[0] += v[0]
        self.x[1] += v[1]
        self.x[2] += v[2]

    def nm2a(self):
        """Converts the Atom coordinates from nanometers (nm) to Angstrom (A).
        """
        if self.unity == 'nm':
            self.x[0] *= 10
            self.x[1] *= 10
            self.x[2] *= 10
            self.unity = 'A'

    def a2nm(self):
        """Converts the Atom coordinates from Angstrom (A) to nanometers (nm).
        """
        if self.unity == 'A':
            self.x[0] *= .1
            self.x[1] *= .1
            self.x[2] *= .1
            self.unity = 'nm'

    def angle(self, other1, other2, degree=False):
        """Calcluates the angle between 3 atoms. Returns radians by default.

        Parameters
        ----------
        other1 : Atom
            second atom
        other2 : Atom
            third atom
        degree : bool
            whether to return the angle in degrees rather than radians. Default
            is False.

        Examples
        --------
        >>> atom1.angle(atom2, atom3)
        >>> 1.57079632679
        >>> atom1.angle(atom2, atom3, degree=True)
        >>> 90.0

        Notes
        -----
        Note that the angle being calculated is the one where the Atom instance
        is in the middle. Following the example above: *atom2 - atom1 - atom3*.

        Returns
        -------
        angle : float
            the value of the angle
        """

        # Note: atom1 must be between 2 and 3
        angle = _p.angle(other1.x, self.x, other2.x)
        if degree is True:
            return angle*180.0/pi
        else:
            return angle

    def dihedral(self, other1, other2, other3, degree=False):
        """Calculates the dihedral angle between four atoms. Returns radians
        by default.

        Parameters
        ----------
        other1 : Atom
            second atom
        other2 : Atom
            third atom
        other3 : Atom
            fourth atom
        degree : bool
            whether to return the angle in degrees rather than radians. Default
            is False.

        Examples
        --------
        >>> atom1.dihedral(atom2, atom3, atom4)
        >>> 1.57079632679
        >>> atom1.dihedral(atom2, atom3, atom4, degree=True)
        >>> 90.0

        Returns
        -------
        angle : float
            the value of the dihedral angle
        """
        ang = _p.dihedral(self.x, other1.x, other2.x, other3.x)
        if degree is True:
            return ang*180.0/pi
        else:
            return ang

    def set_resname(self, resname):
        """Change the residue name.
        """
        self.resname = resname

    def set_chain_id(self, chain_id):
        """Change the chain identifier.
        """
        self.chain_id = chain_id
