"""
    Quatfit routines for PDB2PQR

    This module is used to find the coordinates of a new
    atom based on a reference set of
    coordinates and a definition set of coordinates.

    Original Code by David J. Heisterberg
    The Ohio Supercomputer Center
    1224 Kinnear Rd.
    Columbus, OH  43212-1163
    (614)292-6036
    djh@osc.edu    djh@ohstpy.bitnet    ohstpy::djh
    
    Translated to C from fitest.f program and interfaced with
    Xmol program by Jan Labanowski,  jkl@osc.edu   jkl@ohstpy.bitnet
    ohstpy::jkl

    Adapted for PDB2PQR by Jens Erik Nielsen
    UCSD/HHMI

    Ported to Python by Todd Dolinsky (todd@ccb.wustl.edu)
    Washington University in St. Louis

"""

__date__ = "30 September 2003"
__author__ = "David Heisterberg, Jan Labanowski, Jens Erik Nielsen, Todd Dolinsky"

import math
from utilities import *

def findCoordinates(numpoints, refcoords, defcoords, defatomcoords):
    """
        Driver for the quaternion file.  Provide the coordinates as inputs
        and obtain the coordinates for the new atom as output.

        Parameters
            numpoints:     The number of points in each list (int)
            refcoords:     The reference coordinates, a list of lists of form
                           [x,y,z] (list)
            defcoords:     The definition coordinates, a list of lists of form
                           [x,y,z] (list)
            defatomcoords: The definition coordinates for the atom to be
                           placed in the reference frame (list)
        Returns
            newcoords:     The coordinates of the new atom in the
                           reference frame (list)
    """
    refcenter, fitcenter, rotation = qfit(numpoints, refcoords, defcoords)
    newcoords = qtransform(1, defatomcoords, refcenter, fitcenter, rotation)

    # Only return the first coordinates
    return newcoords[0]
    
def qtransform(numpoints, defcoords, refcenter, fitcenter, rotation):
    """
        Transform the set of defcoords using the reference center, the fit
        center, and a rotation matrix.

        Parameters
            numpoints: The number of points in each list (int)
            defcoords: Definition coordinates (list)
            refcenter: The reference center (list)
            defcenter: The definition center (list)
            rotation:  The rotation matrix (list)
        Returns
            newcoords: The coordinates of the new point (list)
    """
    
    if numpoints == 1:
        defcoords = [defcoords]

    fitcoords = translate(numpoints, defcoords, fitcenter, 1)
    rotated = rotmol(numpoints, fitcoords, rotation)
    newcoords = translate(numpoints, rotated, refcenter, 2)

    return newcoords


def qfit(numpoints, refcoords, defcoords):
    """
        Method for getting new atom coordinates from sets of reference
        and definition coordinates.

        Parameters
            numpoints: The number of points in each list (int)
            refcoords: List of reference coordinates, with each set
                       a list of form [x,y,z] (list)
            defcoords: List of definition coordinates, with each set
                       a list of form [x,y,z] (list)
    """
    nrot = 30
    
    refcenter, refcoords = center(numpoints, refcoords)
    defcenter, defcoords = center(numpoints, defcoords)

    q, u = qtrfit(numpoints, defcoords, refcoords, nrot)
    rotated = rotmol(numpoints, defcoords, u)    
    newcoords = translate(numpoints, rotated, refcenter, 2)

    return refcenter, defcenter, u

def qchichange(initcoords, refcoords, angle):
    """
        Change the chiangle of the reference coordinate using the
        initcoords and the given angle

        Parameters
            initcoords: Coordinates based on the point and basis atoms
                        (one dimensional list)
            difchi    : The angle to use (float)
            refcoords : The atoms to analyze (list of many coordinates)
        Returns
            newcoords : The new coordinates of the atoms (list of many coords)
    """
    # Initialize

    L,R = [],[]
    for i in range(3):
        L.append(0.0)
        R.append([0.0,0.0,0.0])

    # Convert to radians and normalize
    
    radangle = math.pi * angle/180.0
    normalized = normalize(initcoords)

    L[0] = normalized[0]
    L[1] = normalized[1]
    L[2] = normalized[2]
 
    # Construct the rotation matrix

    R[0][0] = math.cos(radangle) + L[0]*L[0] * (1.0 - math.cos(radangle))
    R[1][1] = math.cos(radangle) + L[1]*L[1] * (1.0 - math.cos(radangle))
    R[2][2] = math.cos(radangle) + L[2]*L[2] * (1.0 - math.cos(radangle))
    R[1][0] = L[0]*L[1]*(1.0 - math.cos(radangle)) - L[2] * math.sin(radangle)
    R[2][0] = L[0]*L[2]*(1.0 - math.cos(radangle)) + L[1] * math.sin(radangle)
    R[0][1] = L[1]*L[0]*(1.0 - math.cos(radangle)) + L[2] * math.sin(radangle)
    R[2][1] = L[1]*L[2]*(1.0 - math.cos(radangle)) - L[0] * math.sin(radangle)
    R[0][2] = L[2]*L[0]*(1.0 - math.cos(radangle)) - L[1] * math.sin(radangle)
    R[1][2] = L[2]*L[1]*(1.0 - math.cos(radangle)) + L[0] * math.sin(radangle)
    

    numpoints = len(refcoords)
    newcoords = rotmol(numpoints, refcoords, R)

    return newcoords

def rotmol(numpoints, x, u):
    """
        Rotate a molecule

        Parameters
            numpoints:  The number of points in the list (int)
            x:          The input coordinates (list)
            u:          The left rotation matrix (list)
        Returns
            out:        The rotated coordinates out=u * x (list)
    """
    out = []
    for i in range(numpoints):
        out.append([])
        out[i].append(u[0][0] *x[i][0] + u[1][0] * x[i][1] + u[2][0] * x[i][2])
        out[i].append(u[0][1] *x[i][0] + u[1][1] * x[i][1] + u[2][1] * x[i][2])
        out[i].append(u[0][2] *x[i][0] + u[1][2] * x[i][1] + u[2][2] * x[i][2])

    return out

def qtrfit(numpoints, defcoords, refcoords, nrot):
    """
        Find the quaternion, q, [and left rotation matrix, u] that minimizes
            | qTXq - Y | ^ 2 [|uX - Y| ^ 2]
        This is equivalent to maximizing Re (qTXTqY)
        The left rotation matrix, u, is obtained from q by
            u = qT1q

        Parameters
            numpoints: The number of points in each list (int)
            defcoords: List of definition coordinates, with each set
                       a list of form [x,y,z] (list)
            refcoords: List of fitted coordinates, with each set
                       a list of form [x,y,z] (list)
            nrot     : The maximum number of jacobi sweeps
        Returns
            q        : The best-fit quaternion
            u        : The best-fit left rotation matrix
    """
    xxyx = 0.0
    xxyy = 0.0
    xxyz = 0.0
    xyyx = 0.0
    xyyy = 0.0
    xyyz = 0.0
    xzyx = 0.0
    xzyy = 0.0
    xzyz = 0.0

    q = []
    c = []

    for i in range(numpoints):
         xxyx = xxyx + defcoords[i][0] * refcoords[i][0]
         xxyy = xxyy + defcoords[i][0] * refcoords[i][1]
         xxyz = xxyz + defcoords[i][0] * refcoords[i][2]
         xyyx = xyyx + defcoords[i][1] * refcoords[i][0]
         xyyy = xyyy + defcoords[i][1] * refcoords[i][1]
         xyyz = xyyz + defcoords[i][1] * refcoords[i][2]
         xzyx = xzyx + defcoords[i][2] * refcoords[i][0]
         xzyy = xzyy + defcoords[i][2] * refcoords[i][1]
         xzyz = xzyz + defcoords[i][2] * refcoords[i][2]

    for i in range(4):
        c.append([])
        for j in range(4): 
            c[i].append(0.0)

    c[0][0] = xxyx + xyyy + xzyz
    
    c[0][1] = xzyy - xyyz
    c[1][1] = xxyx - xyyy - xzyz
    
    c[0][2] = xxyz - xzyx
    c[1][2] = xxyy + xyyx
    c[2][2] = xyyy - xzyz - xxyx
    
    c[0][3] = xyyx - xxyy
    c[1][3] = xzyx + xxyz
    c[2][3] = xyyz + xzyy
    c[3][3] = xzyz - xxyx - xyyy

    d,v = jacobi(c, nrot) # diagonalize c

    for i in range(4):
        q.append(v[i][3])

    u = q2mat(q)

    return q,u

def jacobi(a, nrot):
    """
        Jacobi diagonalizer with sorted output, only good for 4x4 matrices

        Parameters
            a:    Matrix to diagonalize (4x4 list)
            nrot: Maximum number of sweeps
        Returns
            d:    Eigenvalues
            v:    Eigenvectors
    """
    v = []
    d = []
    for j in range(4):
        d.append(0)
        v.append([])
        for i in range(4):
            v[j].append(0.0)
        v[j][j] = 1.0
        d[j] = a[j][j]

    for l in range(nrot):
        dnorm = 0.0
        onorm = 0.0
        for j in range(4):
            dnorm = dnorm + abs(d[j])
            for i in range(j):
                onorm = onorm + abs(a[i][j])

        if dnorm != 0:
            if onorm/dnorm <= 1e-12: break
            
        for j in range(1,4):
            for i in range(j):
                b = a[i][j]
                if abs(b) > 0.0:
                    dma = d[j] - d[i]
                    if abs(dma) + abs(b) <= abs(dma):
                        t = b / dma
                    else:
                        q = 0.5 * dma/b
                        t = 1.0/(abs(q) + math.sqrt(1 + q*q))
                        if q < 0:
                            t = t * -1
                    c = 1.0/math.sqrt(t*t + 1)
                    s = t*c
                    a[i][j] = 0.0
                    for k in range(i):
                        atemp = c * a[k][i] - s * a[k][j]
                        a[k][j] = s * a[k][i] + c * a[k][j]
                        a[k][i] = atemp
                    for k in range(i+1 ,j):
                        atemp = c * a[i][k] - s * a[k][j]
                        a[k][j] = s * a[i][k] + c * a[k][j]
                        a[i][k] = atemp
                    for k in range(j+1, 4):
                        atemp = c * a[i][k] - s * a[j][k]
                        a[j][k] = s * a[i][k] + c * a[j][k]
                        a[i][k] = atemp
                    for k in range(4):
                        vtemp = c * v[k][i] - s * v[k][j]
                        v[k][j] = s * v[k][i] + c * v[k][j]
                        v[k][i] = vtemp
                            
                    dtemp = c*c*d[i] + s*s*d[j] - 2.0*c*s*b
                    d[j] = s*s*d[i] + c*c*d[j] +  2.0*c*s*b
                    d[i] = dtemp

    nrot = l
    for j in range(3):
        k = j
        dtemp = d[k]
        for i in range(j+1,4):
            if d[i] < dtemp:
                k = i
                dtemp = d[k]
        if k > j:
            d[k] = d[j]
            d[j] = dtemp
            for i in range(4):
                dtemp = v[i][k]
                v[i][k] = v[i][j]
                v[i][j] = dtemp
                        
    return d,v

def q2mat(q):
    """
        Generate a left rotation matrix from a normalized quaternion

        Parameters
            q:  The normalized quaternion (list)
        Returns
            u:  The rotation matrix (2-dimensional list)
    """
    u = []
    for i in range(3):
        u.append([])
        for j in range(3):
            u[i].append(0.0)

    u[0][0] = q[0]*q[0] + q[1]*q[1] - q[2]*q[2] - q[3]*q[3]
    u[0][1] = 2.0 * (q[1] * q[2] - q[0] * q[3])
    u[0][2] = 2.0 * (q[1] * q[3] + q[0] * q[2])
    
    u[1][0] = 2.0 * (q[2] * q[1] + q[0] * q[3])
    u[1][1] = q[0]*q[0] - q[1]*q[1] + q[2]*q[2] - q[3]*q[3]
    u[1][2] = 2.0 * (q[2] * q[3] - q[0] * q[1])
    
    u[2][0] = 2.0 *(q[3] * q[1] - q[0] * q[2])
    u[2][1] = 2.0 * (q[3] * q[2] + q[0] * q[1])
    u[2][2] = q[0]*q[0] - q[1]*q[1] - q[2]*q[2] + q[3]*q[3]

    return u
    
def center(numpoints, refcoords):
    """
        Center a molecule using equally weighted points

        Parameters
            numpoints: Number of points
            refcoords: List of reference coordinates, with each set
                       a list of form [x,y,z] (list)
        Returns
            refcenter: Center of the set of points (list)
            relcoords: Moved refcoords relative to refcenter (list)
    """
    refcenter = []
    relcoords = []
    
    for i in range(3):
        refcenter.append(0.0)

    for i in range(numpoints):
        refcenter[0] += refcoords[i][0]
        refcenter[1] += refcoords[i][1]
        refcenter[2] += refcoords[i][2]

    for i in range(3):
        refcenter[i] = refcenter[i] / numpoints
        
    for i in range(numpoints):
        relcoords.append([])
        relcoords[i].append(refcoords[i][0] - refcenter[0])
        relcoords[i].append(refcoords[i][1] - refcenter[1])
        relcoords[i].append(refcoords[i][2] - refcenter[2])
        
    return refcenter, relcoords


def translate(numpoints, refcoords, center, mode):
    """
        Translate a molecule using equally weighted points

        Parameters
            numpoints: Number of points
            refcoords: List of reference coordinates, with each set
                       a list of form [x,y,z] (list)
            center:    Center of the system(list)
            mode:      If 1, center will be subtracted from refcoords
                       If 2, center will be added to refcoords
        Returns
            relcoords: Moved refcoords relative to refcenter (list)
    """
    relcoords = []
    
    if mode == 1:
        modif = -1
    elif mode == 2:
        modif = 1

    for i in range(numpoints):
        relcoords.append([])
        relcoords[i].append(refcoords[i][0] + modif * center[0])
        relcoords[i].append(refcoords[i][1] + modif * center[1])
        relcoords[i].append(refcoords[i][2] + modif * center[2])

    return relcoords
