# Copyright (C) 2004, Thomas Hamelryck (thamelry@binf.ku.dk)
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

"""Vector class, including rotation-related functions."""

from __future__ import print_function

import numpy


def m2rotaxis(m):
    """
    Return angles, axis pair that corresponds to rotation matrix m.
    """
    # Angle always between 0 and pi
    # Sense of rotation is defined by axis orientation
    t=0.5*(numpy.trace(m)-1)
    t=max(-1, t)
    t=min(1, t)
    angle=numpy.arccos(t)
    if angle<1e-15:
        # Angle is 0
        return 0.0, Vector(1, 0, 0)
    elif angle<numpy.pi:
        # Angle is smaller than pi
        x=m[2, 1]-m[1, 2]
        y=m[0, 2]-m[2, 0]
        z=m[1, 0]-m[0, 1]
        axis=Vector(x, y, z)
        axis.normalize()
        return angle, axis
    else:
        # Angle is pi - special case!
        m00=m[0, 0]
        m11=m[1, 1]
        m22=m[2, 2]
        if m00>m11 and m00>m22:
            x=numpy.sqrt(m00-m11-m22+0.5)
            y=m[0, 1]/(2*x)
            z=m[0, 2]/(2*x)
        elif m11>m00 and m11>m22:
            y=numpy.sqrt(m11-m00-m22+0.5)
            x=m[0, 1]/(2*y)
            z=m[1, 2]/(2*y)
        else:
            z=numpy.sqrt(m22-m00-m11+0.5)
            x=m[0, 2]/(2*z)
            y=m[1, 2]/(2*z)
        axis=Vector(x, y, z)
        axis.normalize()
        return numpy.pi, axis


def vector_to_axis(line, point):
    """
    Returns the vector between a point and
    the closest point on a line (ie. the perpendicular
    projection of the point on the line).

    @type line: L{Vector}
    @param line: vector defining a line

    @type point: L{Vector}
    @param point: vector defining the point
    """
    line=line.normalized()
    np=point.norm()
    angle=line.angle(point)
    return point-line**(np*numpy.cos(angle))


def rotaxis2m(theta, vector):
    """
    Calculate a left multiplying rotation matrix that rotates
    theta rad around vector.

    Example:

        >>> m=rotaxis(pi, Vector(1, 0, 0))
        >>> rotated_vector=any_vector.left_multiply(m)

    @type theta: float
    @param theta: the rotation angle


    @type vector: L{Vector}
    @param vector: the rotation axis

    @return: The rotation matrix, a 3x3 Numeric array.
    """
    vector=vector.copy()
    vector.normalize()
    c=numpy.cos(theta)
    s=numpy.sin(theta)
    t=1-c
    x, y, z=vector.get_array()
    rot=numpy.zeros((3, 3))
    # 1st row
    rot[0, 0]=t*x*x+c
    rot[0, 1]=t*x*y-s*z
    rot[0, 2]=t*x*z+s*y
    # 2nd row
    rot[1, 0]=t*x*y+s*z
    rot[1, 1]=t*y*y+c
    rot[1, 2]=t*y*z-s*x
    # 3rd row
    rot[2, 0]=t*x*z-s*y
    rot[2, 1]=t*y*z+s*x
    rot[2, 2]=t*z*z+c
    return rot

rotaxis=rotaxis2m


def refmat(p, q):
    """
    Return a (left multiplying) matrix that mirrors p onto q.

    Example:
        >>> mirror=refmat(p, q)
        >>> qq=p.left_multiply(mirror)
        >>> print(q)
        >>> print(qq) # q and qq should be the same

    @type p,q: L{Vector}
    @return: The mirror operation, a 3x3 Numeric array.
    """
    p.normalize()
    q.normalize()
    if (p-q).norm()<1e-5:
        return numpy.identity(3)
    pq=p-q
    pq.normalize()
    b=pq.get_array()
    b.shape=(3, 1)
    i=numpy.identity(3)
    ref=i-2*numpy.dot(b, numpy.transpose(b))
    return ref


def rotmat(p, q):
    """
    Return a (left multiplying) matrix that rotates p onto q.

    Example:
        >>> r=rotmat(p, q)
        >>> print(q)
        >>> print(p.left_multiply(r))

    @param p: moving vector
    @type p: L{Vector}

    @param q: fixed vector
    @type q: L{Vector}

    @return: rotation matrix that rotates p onto q
    @rtype: 3x3 Numeric array
    """
    rot=numpy.dot(refmat(q, -p), refmat(p, -p))
    return rot


def calc_angle(v1, v2, v3):
    """
    Calculate the angle between 3 vectors
    representing 3 connected points.

    @param v1, v2, v3: the tree points that define the angle
    @type v1, v2, v3: L{Vector}

    @return: angle
    @rtype: float
    """
    v1=v1-v2
    v3=v3-v2
    return v1.angle(v3)


def calc_dihedral(v1, v2, v3, v4):
    """
    Calculate the dihedral angle between 4 vectors
    representing 4 connected points. The angle is in
    ]-pi, pi].

    @param v1, v2, v3, v4: the four points that define the dihedral angle
    @type v1, v2, v3, v4: L{Vector}
    """
    ab=v1-v2
    cb=v3-v2
    db=v4-v3
    u=ab**cb
    v=db**cb
    w=u**v
    angle=u.angle(v)
    # Determine sign of angle
    try:
        if cb.angle(w)>0.001:
            angle=-angle
    except ZeroDivisionError:
        # dihedral=pi
        pass
    return angle


class Vector(object):
    "3D vector"

    def __init__(self, x, y=None, z=None):
        if y is None and z is None:
            # Array, list, tuple...
            if len(x)!=3:
                raise ValueError("Vector: x is not a "
                                 "list/tuple/array of 3 numbers")
            self._ar=numpy.array(x, 'd')
        else:
            # Three numbers
            self._ar=numpy.array((x, y, z), 'd')

    def __repr__(self):
        x, y, z=self._ar
        return "<Vector %.2f, %.2f, %.2f>" % (x, y, z)

    def __neg__(self):
        "Return Vector(-x, -y, -z)"
        a=-self._ar
        return Vector(a)

    def __add__(self, other):
        "Return Vector+other Vector or scalar"
        if isinstance(other, Vector):
            a=self._ar+other._ar
        else:
            a=self._ar+numpy.array(other)
        return Vector(a)

    def __sub__(self, other):
        "Return Vector-other Vector or scalar"
        if isinstance(other, Vector):
            a=self._ar-other._ar
        else:
            a=self._ar-numpy.array(other)
        return Vector(a)

    def __mul__(self, other):
        "Return Vector.Vector (dot product)"
        return sum(self._ar*other._ar)

    def __div__(self, x):
        "Return Vector(coords/a)"
        a=self._ar/numpy.array(x)
        return Vector(a)

    def __pow__(self, other):
        "Return VectorxVector (cross product) or Vectorxscalar"
        if isinstance(other, Vector):
            a, b, c=self._ar
            d, e, f=other._ar
            c1=numpy.linalg.det(numpy.array(((b, c), (e, f))))
            c2=-numpy.linalg.det(numpy.array(((a, c), (d, f))))
            c3=numpy.linalg.det(numpy.array(((a, b), (d, e))))
            return Vector(c1, c2, c3)
        else:
            a=self._ar*numpy.array(other)
            return Vector(a)

    def __getitem__(self, i):
        return self._ar[i]

    def __setitem__(self, i, value):
        self._ar[i]=value

    def __contains__(self, i):
        return (i in self._ar)

    def norm(self):
        "Return vector norm"
        return numpy.sqrt(sum(self._ar*self._ar))

    def normsq(self):
        "Return square of vector norm"
        return abs(sum(self._ar*self._ar))

    def normalize(self):
        "Normalize the Vector"
        self._ar=self._ar/self.norm()

    def normalized(self):
        "Return a normalized copy of the Vector"
        v=self.copy()
        v.normalize()
        return v

    def angle(self, other):
        "Return angle between two vectors"
        n1=self.norm()
        n2=other.norm()
        c=(self*other)/(n1*n2)
        # Take care of roundoff errors
        c=min(c, 1)
        c=max(-1, c)
        return numpy.arccos(c)

    def get_array(self):
        "Return (a copy of) the array of coordinates"
        return numpy.array(self._ar)

    def left_multiply(self, matrix):
        "Return Vector=Matrix x Vector"
        a=numpy.dot(matrix, self._ar)
        return Vector(a)

    def right_multiply(self, matrix):
        "Return Vector=Vector x Matrix"
        a=numpy.dot(self._ar, matrix)
        return Vector(a)

    def copy(self):
        "Return a deep copy of the Vector"
        return Vector(self._ar)

if __name__=="__main__":

    from numpy.random import random

    v1=Vector(0, 0, 1)
    v2=Vector(0, 0, 0)
    v3=Vector(0, 1, 0)
    v4=Vector(1, 1, 0)

    v4.normalize()

    print(v4)

    print(calc_angle(v1, v2, v3))
    dih=calc_dihedral(v1, v2, v3, v4)
    # Test dihedral sign
    assert(dih>0)
    print("DIHEDRAL %f" % dih)

    ref=refmat(v1, v3)
    rot=rotmat(v1, v3)

    print(v3)
    print(v1.left_multiply(ref))
    print(v1.left_multiply(rot))
    print(v1.right_multiply(numpy.transpose(rot)))

    # -
    print(v1-v2)
    print(v1-1)
    print(v1+(1, 2, 3))
    # +
    print(v1+v2)
    print(v1+3)
    print(v1-(1, 2, 3))
    # *
    print(v1*v2)
    # /
    print(v1/2)
    print(v1/(1, 2, 3))
    # **
    print(v1**v2)
    print(v1**2)
    print(v1**(1, 2, 3))
    # norm
    print(v1.norm())
    # norm squared
    print(v1.normsq())
    # setitem
    v1[2]=10
    print(v1)
    # getitem
    print(v1[2])

    print(numpy.array(v1))

    print("ROT")

    angle=random()*numpy.pi
    axis=Vector(random(3)-random(3))
    axis.normalize()

    m=rotaxis(angle, axis)

    cangle, caxis=m2rotaxis(m)

    print(angle-cangle)
    print(axis-caxis)
    print("")