# -*- coding: utf-8 -*-

import numpy
import math
import copy

from metalSite import MetalSite, getCoordinatesOfAtomsInList, multipleCoordination
from poses import PoseWithLinearPatterns, PoseWithPlanarPatterns

from fitting import Fitting

#===============================================================================
"""CONSTANTS"""
#===============================================================================

# Epsilon for testing whether a number is close to zero
_EPS = numpy.finfo(float).eps * 4.0

#===============================================================================
"""CLASS"""
#===============================================================================

class Rotations(object):
    """
    Rotations are defined as all rigid body movements minimizing a value of rmsd between each pair of local patterns
    Rotations are represented by rotation matrices derived from superposition of pars of local patterns
    """  
   
    def __init__(self, site1, site2):
        
        self.site1 = site1
        self.site2 = site2
        
        if multipleCoordination(self.site1, self.site2):
            self.coordination = 'multiple'
        else:
            self.coordination = 'single'

    def getPoses(self):
        
        self.getLocalPatterns()

        poses = []
        
        for query_pattern in self.patterns1:
            query_set = getCoordinatesOfAtomsInList(query_pattern)

            for target_pattern in self.patterns2:
                
                if self.coordination == 'multiple':
                    pose = PoseWithPlanarPatterns()
                else:
                    pose = PoseWithLinearPatterns()
                
                target_set = getCoordinatesOfAtomsInList(target_pattern)
                transformationMatrix = findBestTransformation(query_set, target_set)
                
                pose.rotationMatrix = extractRotationMatrix(transformationMatrix)
                pose.pattern1 = query_set
                pose.pattern2 = target_set
                
                poses.append(pose)
                
        return poses

    def getLocalPatterns(self):
            
        if multipleCoordination(self.site1, self.site2):
            #Both sites have more than one donor atom            
            triangles1 = self.findPlanarLocalPatterns(self.site1)
            triangles2 = self.findPlanarLocalPatterns(self.site2)

            permutations = self.generatePermutations(triangles2)
            
            self.patterns1 = triangles1
            self.patterns2 = triangles2 + permutations
            
        else:
            #At least one site is a single donor site          
            self.patterns1 = self.findLinearLocalPatterns(self.site1)
            self.patterns2 = self.findLinearLocalPatterns(self.site2)

    def getAxes(self):
        'The axes contain local patterns of the first site.'
        
        self.axes = []

        for pattern in self.patterns1:
            for atom in pattern:
                axis = [atom.x, atom.y, atom.z]
                self.axes.append(axis)
        
        return self.axes

    def getNumberOfTargetPatterns(self):
        return len(self.patterns2)

    def generatePermutations(self, triangles):
        '''
        Explicitly sets the permutations for triangles
        Takes a triangle: a list of atoms - > [atomA', atomB']
        Returns the pair of atoms reversed -> [atomB', atomA']
        '''
           
        permutations = []
        for triangle in triangles:
            permutations.append([triangle[1], triangle[0]])
            
        return permutations     
   
    def findPlanarLocalPatterns(self, site):
        '''
        Composes a metal atom and two donor atoms in all possible triangles originated from the metal
        site: object of MetalSite class
        Returns a list of atom pairs -> [[donor_atom1,donor_atom2],[donor_atom1, donor_atom2],...]
        '''    
        
        localPatterns = []
 
        for i in range(0, len(site.donorAtoms)-1):
            for j in range(i+1, len(site.donorAtoms)):
                localPatterns.append([site.donorAtoms[i], site.donorAtoms[j]])
        
        return localPatterns
   
    def findLinearLocalPatterns(self, site):
            
        localPatterns = []
            
        for atom in site.donorAtoms:    
            localPatterns.append([atom])
        
        return localPatterns 

#===============================================================================
"""METHODS"""
#===============================================================================

def extractRotationMatrix(M):
    """
    Extracts a rotation matrix from a tranformation matrix
        
    M: tranformation matrix [4x4]
        
    Returns a rotation matrix [3x3]
    """
    
    rotM = [numpy.take(M[i], [0,1,2]) for i in range(0, 3)]

    
    return rotM    

def applyRotationMatrices(site, rotationalMatrices):
    '''
    Rotates a site using all rotation matrices derived from superposing all possible triangles         
    site: object of metalSite class;
    rotMatrices: list of rotation matrices
        
    Returns a list of MetalSite objects
    '''                      
        
    rotatedSites = [applyRotationMatrix(site, matrix) for matrix in rotMatrices]                     
               
    return rotatedSites

def applyRotationMatrix(site, rotM):
    """
    Rotates a site using a rotation matrix M
    
    site: MetalSite object
    rotM: rotation matrix [3x3]
    
    Returns a site with changed coordinates (creates a new site object)
    """
    
    rotSite = MetalSite()
    
    name = site.name
    
    metals = []
    for atom in site.metals:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        rot_v = numpy.dot(rotM,v)
        tempAtom = copy.copy(atom)
        tempAtom.x = rot_v[0]
        tempAtom.y = rot_v[1]
        tempAtom.z = rot_v[2]    
        metals.append(tempAtom)

    donorAtoms = []
    for atom in site.donorAtoms:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        rot_v = numpy.dot(rotM,v)
        tempAtom = copy.copy(atom)
        tempAtom.x = rot_v[0]
        tempAtom.y = rot_v[1]
        tempAtom.z = rot_v[2]    
        donorAtoms.append(tempAtom)            

    approxSphereAtoms = []
    for atom in site.approxSphereAtoms:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        rot_v = numpy.dot(rotM,v)
        tempAtom = copy.copy(atom)
        tempAtom.x = rot_v[0]
        tempAtom.y = rot_v[1]
        tempAtom.z = rot_v[2]    
        approxSphereAtoms.append(tempAtom)
        
   
    rotSite.name = name
    rotSite.metals = metals
    rotSite.donorAtoms = donorAtoms
    rotSite.approxSphereAtoms = approxSphereAtoms

    rotSite.getAtomsForMatching()
    rotSite.getBackboneAndSidechain()   
    
    return rotSite    

def applyRotationMatrix2(site, tempSite, rotM):
    """
    Rotates a site using a rotation matrix M
    
    site: MetalSite object
    rotM: rotation matrix [3x3]
    
    Returns a site with changed coordinates (changes tempSite)
    """
    
    v = numpy.array([0, 0, 0], dtype=numpy.float64, copy=True)
    i=0
    for atom in site.metals:                    
        v[0] = atom.x
        v[1] = atom.y
        v[2] = atom.z
        
        rot_v = numpy.dot(rotM,v)
        tempSite.metals[i].x = rot_v[0]
        tempSite.metals[i].y = rot_v[1]
        tempSite.metals[i].z = rot_v[2]    
        i=i+1

    i=0
    for atom in site.donorAtoms:                    
        v[0] = atom.x
        v[1] = atom.y
        v[2] = atom.z
        rot_v = numpy.dot(rotM,v)
        tempSite.donorAtoms[i].x = rot_v[0]
        tempSite.donorAtoms[i].y = rot_v[1]
        tempSite.donorAtoms[i].z = rot_v[2]    
        i=i+1
        
    i=0
    for atom in site.approxSphereAtoms:                    
        v[0] = atom.x
        v[1] = atom.y
        v[2] = atom.z
        rot_v = numpy.dot(rotM,v)
        tempSite.approxSphereAtoms[i].x = rot_v[0]
        tempSite.approxSphereAtoms[i].y = rot_v[1]
        tempSite.approxSphereAtoms[i].z = rot_v[2]    
        i=i+1   
        
    return tempSite

def findBestTransformation(v0, v1):
    '''
    Takes as input a pair of triangles and finds superpositions between the two, optimizing a value of RMSD
    v0: vector of coordinates
    v1: vector of coordinates
        
    Returns a transformation matrix [4x4] 
    '''

    v0 = numpy.array(v0, dtype=numpy.float64, copy=True)
    v1 = numpy.array(v1, dtype=numpy.float64, copy=True)
      
    #[A'-A:[x'-x, y'-y, z'-z], B'-B:[x'-x, y'-y, z'-z]]    
    delta_m = v1-v0
    #[A'+A:[x'+x, y'+y, z'+z], B'+B:[x'+x, y'+y, z'+z]]
    delta_p = v1+v0
         
    #[(A'-A)^2:[(x'-x)^2, (y'-y)^2, (z'-z)^2], (B'-B)^2:[(x'-x)^2, (y'-y)^2, (z'-z)^2]]
    sq_delta_m = delta_m* delta_m
    #[(A'+A)^2:[(x'+x)^2, (y'+y)^2, (z'+z)^2], (B'+B)^2:[(x'+x)^2, (y'+y)^2, (z'+z)^2]]
    sq_delta_p = delta_p* delta_p
    
    xx_m, yy_m, zz_m = numpy.sum( sq_delta_m, axis=0)
    xx_p, yy_p, zz_p = numpy.sum( sq_delta_p, axis=0)
       
    xm_ym, ym_zm, zm_xm = numpy.sum(delta_m * numpy.roll(delta_m, -1, axis=1), axis=0)
    xp_yp, yp_zp, zp_xp = numpy.sum(delta_p * numpy.roll(delta_p, -1, axis=1), axis=0)
      
    xm_yp, ym_zp, zm_xp = numpy.sum(delta_m * numpy.roll(delta_p, -1, axis=1), axis=0)
    xm_zp, ym_xp, zm_yp = numpy.sum(delta_m * numpy.roll(delta_p, -2, axis=1), axis=0)
        
    # Computes symmetric matrix N
    N = [[ xx_m+yy_m+zz_m, zm_yp - ym_zp , xm_zp - zm_xp , ym_xp - xm_yp ],
         [ zm_yp - ym_zp ,yy_p+zz_p+xx_m , xm_ym - xp_yp , zm_xm - zp_xp ],
         [ xm_zp - zm_xp , xm_ym - xp_yp ,xx_p+zz_p+yy_m , ym_zm - yp_zp ],
         [ ym_xp - xm_yp , zm_xm - zp_xp , ym_zm - yp_zp ,xx_p+yy_p+zz_m ]]
        
    # Finds eignvalues and eignvectors    
    w, V = numpy.linalg.eigh(N)
        
    # Gives an eigenvector corresponding to the smallest positive eigenvalue
    q = V[:, numpy.argmin(w)]
        
    # Unit quaternion
    q /= vectorNorm(q)
        
    # Homogeneous transformation matrix
    M = matrixOutOfQuaternion(q)   
        
    return M

def vectorNorm(data, axis=None, out=None):
    """
    Returns length, i.e. eucledian norm, of ndarray along axis.
    """
        
    data = numpy.array(data, dtype=numpy.float64, copy=True)
    if out is None:
        if data.ndim == 1:
            return numpy.sqrt(numpy.dot(data, data))
        data *= data
        out = numpy.atleast_1d(numpy.sum(data, axis=axis))
        numpy.sqrt(out, out)
        return out
    else:
        data *= data
        numpy.sum(data, axis=axis, out=out)
        numpy.sqrt(out, out)

def quaternionAboutAxis(angle, axis):
    """
    Returns quaternion for rotation about axis.

    >>> q = quaternion_about_axis(0.123, [1, 0, 0])
    >>> numpy.allclose(q, [0.99810947, 0.06146124, 0, 0])
    True
    """
    
    q = numpy.array([0.0, axis[0], axis[1], axis[2]])
    qlen = vectorNorm(q)
    if qlen > _EPS:
        q *= numpy.sin(angle/2.0) / qlen
    q[0] = numpy.cos(angle/2.0)
    
    
    return q

def matrixOutOfQuaternion(quaternion):
    """
    Returns a homogeneous transformation matrix from quaternion.
        
    quaternion: quaternion vector [1x4]
        
    Returns a transformation matrix [4x4]
    """
        
    q = numpy.array(quaternion, dtype=numpy.float64, copy=True)
    n = numpy.dot(q, q)
    if n < _EPS:
        return numpy.identity(4)
    q *= numpy.sqrt(2.0 / n)
    q = numpy.outer(q, q)
        
    return numpy.array([
            [1.0-q[2, 2]-q[3, 3],     q[1, 2]-q[3, 0],     q[1, 3]+q[2, 0], 0.0],
            [    q[1, 2]+q[3, 0], 1.0-q[1, 1]-q[3, 3],     q[2, 3]-q[1, 0], 0.0],
            [    q[1, 3]-q[2, 0],     q[2, 3]+q[1, 0], 1.0-q[1, 1]-q[2, 2], 0.0],
            [                0.0,                 0.0,                 0.0, 1.0]])

def setAxis(lable):
    
    if lable is 'x':
        axis = [1, 0, 0]
    elif lable is 'y':
        axis = [0, 1, 0]
    elif lable is 'z':
        axis = [0, 0, 1]
    
    return axis

def rotateSiteAboutAxis(site, axis, angleGrad):
    """
    Rotates a given site about a given axis by a given angle
    
    site: a MatalSite object    
    axis: an axis for rotation to be about
    angleGrad: an angle of rotation (degree scale)
    
    Returns a rotated site
    """
    
    angleRad = angleGrad*(2*math.pi/360)

    q = quaternionAboutAxis(angleRad, axis)
    tranformMatrix = matrixOutOfQuaternion(q)
    rotationMatrix = extractRotationMatrix(tranformMatrix)
    rotatedSite = applyRotationMatrix(site, rotationMatrix)    
    
    return rotatedSite, rotationMatrix 

def rotateSiteAboutAxis2(site, tempSite, axis, angleGrad):
    """
    Rotates a site by a given angle
    
    site: a MetalSite object
    axis: an axis for rotation to be about
    angleGrad: an angle of rotation (degree scale)      

    Returns a pertrubated site
    """
    
    angleRad = angleGrad*(2*math.pi/360)
    
    q = quaternionAboutAxis(angleRad, axis)
    tranformMatrix = matrixOutOfQuaternion(q)
    rotationMatrix = extractRotationMatrix(tranformMatrix)
    rotatedSite = applyRotationMatrix2(site, tempSite, rotationMatrix)    

    return rotatedSite,rotationMatrix
 
def translateSiteAlongAxis(site, axis, delta):
    
    transSite = MetalSite()
    
    name = site.name
    
    if axis is 'x':
        metals = []
        for atom in site.metals:                    
            tempAtom = copy.copy(atom)
            tempAtom.x = atom.x + delta   
            metals.append(tempAtom)
        donorAtoms = []
        for atom in site.donorAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.x = atom.x + delta
            donorAtoms.append(tempAtom)            
        approxSphereAtoms = []
        for atom in site.approxSphereAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.x = atom.x + delta  
            approxSphereAtoms.append(tempAtom)
    
    if axis is 'y':
        metals = []
        for atom in site.metals:                    
            tempAtom = copy.copy(atom)
            tempAtom.y = atom.y + delta   
            metals.append(tempAtom)
        donorAtoms = []
        for atom in site.donorAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.y = atom.y + delta
            donorAtoms.append(tempAtom)            
        approxSphereAtoms = []
        for atom in site.approxSphereAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.y = atom.y + delta  
            approxSphereAtoms.append(tempAtom)
    
    if axis is 'z':
        metals = []
        for atom in site.metals:                    
            tempAtom = copy.copy(atom)
            tempAtom.z = atom.z + delta   
            metals.append(tempAtom)
        donorAtoms = []
        for atom in site.donorAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.z = atom.z + delta
            donorAtoms.append(tempAtom)            
        approxSphereAtoms = []
        for atom in site.approxSphereAtoms:                    
            tempAtom = copy.copy(atom)
            tempAtom.z = atom.z + delta  
            approxSphereAtoms.append(tempAtom)
            
    transSite.name = name
    transSite.metals = metals
    transSite.donorAtoms = donorAtoms
    transSite.approxSphereAtoms = approxSphereAtoms
    
    transSite.getAtomsForMatching()
    transSite.getBackboneAndSidechain()   
    
    return transSite

def translateSiteAlongAxis2(site, tempSite, axis, delta):

    if axis is 'x':
        i = 0
        for atom in site.metals:                    
            tempSite.metals[i].x = atom.x + delta
            i=i+1   
        i = 0
        for atom in site.donorAtoms:                    
            tempSite.donorAtoms[i].x = atom.x + delta
            i=i+1 
        i = 0      
        for atom in site.approxSphereAtoms:                    
            tempSite.approxSphereAtoms[i].x = atom.x + delta
            i=i+1   
    
    if axis is 'y':
        i = 0
        for atom in site.metals:                    
            tempSite.metals[i].y = atom.y + delta
            i=i+1   
        i = 0
        for atom in site.donorAtoms:                    
            tempSite.donorAtoms[i].y = atom.y + delta
            i=i+1 
        i = 0      
        for atom in site.approxSphereAtoms:                    
            tempSite.approxSphereAtoms[i].y = atom.y + delta
            i=i+1 
    
    if axis is 'z':
        i = 0
        for atom in site.metals:                    
            tempSite.metals[i].z = atom.z + delta
            i=i+1   
        i = 0
        for atom in site.donorAtoms:                    
            tempSite.donorAtoms[i].z = atom.z + delta
            i=i+1 
        i = 0      
        for atom in site.approxSphereAtoms:                    
            tempSite.approxSphereAtoms[i].z = atom.z + delta
            i=i+1 
    
    return tempSite

def applyTransformationToPdb(pdb, rotM, tranM):
    
    for atom in pdb.atoms:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        transf_v = numpy.dot(v, rotM) + tranM
        
        atom.x = float(transf_v[0])
        atom.y = float(transf_v[1])
        atom.z = float(transf_v[2])
    
    return pdb

def applyTransformationToSite(site, rotM, tranM):
    
    for atom in site.metals:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        transf_v = numpy.dot(v, rotM) + tranM
        
        atom.x = transf_v[0]
        atom.y = transf_v[1]
        atom.z = transf_v[2]    

    for atom in site.donorAtoms:
        
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        transf_v = numpy.dot(v, rotM) + tranM
        
        atom.x = transf_v[0]
        atom.y = transf_v[1]
        atom.z = transf_v[2]           

    for atom in site.approxSphereAtoms:                    
        v = [atom.x, atom.y, atom.z]
        v = numpy.array(v, dtype=numpy.float64, copy=True)
        transf_v = numpy.dot(v, rotM) + tranM
        
        atom.x = transf_v[0]
        atom.y = transf_v[1]
        atom.z = transf_v[2]

    site.getAtomsForMatching()
    site.getBackboneAndSidechain()   
    
    return site

def getMutualCoordinateSets(site, pdb):

    site_coords = []
    pdb_coords = []
            
    for site_atom_item in site.backbone:
                
        site_atom = site_atom_item[1]
        pdb_atom_q = pdb.query(' chain %s resname %s and resid = %s atype CA ' % (site_atom.chain, site_atom.aa, site_atom.resid))
        
        if pdb_atom_q:
            pdb_atom = pdb_atom_q[0]
            site_coords.append([site_atom.x, site_atom.y, site_atom.z])
            pdb_coords.append([pdb_atom.x, pdb_atom.y, pdb_atom.z])
        
    site_set = numpy.array(site_coords, 'f')
    pdb_set = numpy.array(pdb_coords, 'f')
    
    return site_set, pdb_set

def superimposeSiteOnPdb(site, pdb):
    
    site_moved = None
    site_set, pdb_set = getMutualCoordinateSets(site, pdb)

    if (site_set.any() and pdb_set.any()):
        sup = Fitting()
        sup.set(pdb_set, site_set)
        sup.run()
        rotM, tranM = sup.get_rotran()
            
        site_moved = applyTransformationToSite(site, rotM, tranM)
        
        return site_moved
    
def superimposePdbOnSite(site, pdb):
    
    pdb_moved = None
    site_set, pdb_set = getMutualCoordinateSets(site, pdb)
    
    if (site_set.any() and pdb_set.any()):
        
        sup = Fitting()

        sup.set(site_set, pdb_set)
        sup.run()
        rotM, tranM = sup.get_rotran()
        pdb_moved = applyTransformationToPdb(pdb, rotM, tranM)  
        
        return pdb_moved, rotM, tranM