#  Copyright (C) 2020 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

from ..dl_py2f.utils import miscutils

def getRayTriangleIntersect(R0, vec_R, A, B, C):
    '''Return the intersection point's coordinates of a ray (line) and a triangle
       Möller-Trumbore algorithm, see:
       doi:10.1080/10867651.1997.10487468'''

    from numpy import cross, dot

    SMALL = 1.0E-10

    # vectors BA and CA
    vec_BA = B - A
    vec_CA = C - A

    h = cross(vec_R, vec_CA)
    a = dot(vec_BA, h)

    # ray parallel to triangle
    if -SMALL < a < SMALL:
        return None

    s = R0 - A
    u = dot(s, h)/a

    if u < -SMALL or u > 1.0:
        return None

    q = cross(s, vec_BA)
    v = dot(vec_R, q)/a

    if v < -SMALL or u+v > 1.0:
        return None

    # the intersection point's position on the line
    t = dot(vec_CA, q)/a

    # we return the intersection in both directions
    return R0 + vec_R*t

    if t > SMALL:
        return R0 + vec_R*t
    # line intersection rather than a ray
    else:
        return R0 + vec_R*t



def projectOntoPlane(coords, a=0.0, b=0.0, c=0.0, d=0.0):
    '''Project cartesian coordinates onto a plane a*x + b*y + c*z + d = 0'''

    from numpy import array, dot, linalg, multiply

    # vector [a, b, c] is the normal vector of the plane
    normal = array([a, b, c])

    projections = coords - multiply(normal, ((dot(coords, normal)+d)/linalg.norm(normal))[:,None])

    return projections


def getPlane(A, B, C):
    ''''''

    from numpy import cross, dot

    normal = cross((B-A),(C-A))
    d = -dot(A, normal)

    return normal[0], normal[1], normal[2], d


def getDistPoint2Plane(D, a, b, c, d):
    '''Return the distance from a point D to face ABC using given normal'''

    from numpy import dot, linalg

    return (dot(D, [a,b,c])+d)/linalg.norm([a,b,c])


@miscutils.time
def getConvexHull(coords, return_vertices=False, return_edges_and_faces=False):
    '''This is an implementation of the quickhull algorithm: O(n*log(h))
       doi:10.1145/235815.235821
       See: https://en.wikipedia.org/wiki/Quickhull
    '''

    from itertools import combinations
    from numpy     import all, append, arange, array, argmax, argmin, argwhere, cross, delete, dot, empty, isin, linalg, meshgrid, searchsorted, sort, unique, vstack

    SMALL = 1.0E-10

    natoms = coords.shape[0]
    indices = arange(natoms)
        
    dict_normals    = dict()
    dict_remaining  = dict()
    faces = array([], dtype=int).reshape((0,3))

#    edges    = array([], dtype=int).reshape((0,2))
#    faces    = array([], dtype=int).reshape((0,3))


    def _getA0B0C0(_coords):
        ''''''

        # 6 points with min and max X, Y, and Z, respectively
        mins = argmin(_coords, axis=0)
        maxs = argmax(_coords, axis=0)
    
        index_1D = argmax(linalg.norm(_coords[mins][:,None]-_coords[maxs], axis=2))
        indices_2D = array(meshgrid(range(3), range(3), indexing='ij')).reshape(2,-1).T[index_1D]
    
        # point A and B have the longest distance out of the obove 6 points
        _iA0, _iB0 = mins[indices_2D[0]], maxs[indices_2D[1]]
    
        # unit vector A-B
        vector_AB = (_coords[_iA0]-_coords[_iB0])/linalg.norm(_coords[_iA0]-_coords[_iB0])
    
        # find a point C having the longest distance from the line A-B
        _iC0 = argmax(linalg.norm(cross(_coords-_coords[_iA0], _coords-_coords[_iB0]), axis=1)/linalg.norm(_coords[_iA0]-_coords[_iB0]))

        return _iA0, _iB0, _iC0


    def _getABCCoords(_iA, _iB, _iC):
        '''Return coordinates of three points'''

        nonlocal coords

        return coords[_iA], coords[_iB], coords[_iC]


    def _getSemisphere(_iA, _iB, _iC):
        '''Return a mask array of the points having negative distances from face ABC, using predefined normal'''

        nonlocal coords, dict_normals

        normal_ABC, d = dict_normals[tuple(sort([_iA,_iB,_iC]))]

        # we only want the points with negative distances
        return (dot(coords, normal_ABC)+d)/linalg.norm(normal_ABC) < -SMALL


    def _getEdges(iD, iA, iB, iC, _iI=None, _iJ=None):
        '''Recursively walk through all adjacent faces to find those having negative distances from D'''

        nonlocal _faces, _edges, faces

        dist_D2ABC = _getDistPoint2Plane(iD, dict_normals[tuple(sort([iA,iB,iC]))])
        # we are only interested in distances that are negative
        if dist_D2ABC < -SMALL:
            # YL NB: axis not available utill NumPy v1.13!
            _faces = unique(vstack((_faces, sort([iA,iB,iC]))), axis=0)
            for iI, iJ in combinations([iA,iB,iC], 2):
                for face in faces:
                    if all(isin([iI,iJ], face)) and not all(sort([iA,iB,iC])==face) and not all(sort(face)==_faces, axis=1).sum():
                        _getEdges(iD, *face, _iI=iI, _iJ=iJ)
        else:
            if _iI != _iJ != None:
                # YL NB: axis not available utill NumPy v1.13!
                _edges = unique(vstack((_edges, sort([_iI,_iJ]))), axis=0)
            

    def _delPlanes(_faces):
        '''Remove a face ABC from the list'''

        nonlocal faces

        for _face in _faces:
            ind = argwhere(all(sort([*_face])==sort(faces, axis=1), axis=1))
            if ind.size:
                faces = delete(faces, ind[0], axis=0)


    def _getOrientedNormal(_iA, _iB, _iC):
        '''Calculate the normal of a face and make sure the original simplex have positive distances from the face'''

        nonlocal iA0, iB0, iC0, iD0, coords

        _A, _B, _C = _getABCCoords(_iA, _iB, _iC)

        # plane a*x + b*y + c*z + d = 0
        normal = cross((_B-_A),(_C-_A))
        d = -dot(_A, normal)

        dists = (dot(coords[[iA0,iB0,iC0,iD0]], normal)+d)/linalg.norm(normal)

        # same side as A0/B0/C0/D0 (>=-1E-10 to allow A0/B0/C0/D0 themselves)
        if all(dists>=-SMALL):
            return normal, d
        # different side as A0/B0/C0/D0
        elif all(dists<=1E10):
            return -normal, -d
        else:
            print("\n>>> _getOrientedNormal ERROR: should never happen!\n")
            return


    def _getDistPoint2Plane(_iD, _normal):
        '''Return the distance from a point D to face ABC using given normal'''

        nonlocal coords

        normal_ABC, d = _normal[0], _normal[1]

        return (dot(coords[_iD], normal_ABC)+d)/linalg.norm(normal_ABC)


    def _getDistPoints2Plane(_iPoints, _iA, _iB, _iC):
        ''''''

        nonlocal coords, dict_normals

        normal_ABC, d = _getOrientedNormal(_iA, _iB, _iC)
        dists = (dot(coords[_iPoints], normal_ABC)+d)/linalg.norm(normal_ABC)

        return dists


    def _getNormal0(_iA, _iB, _iC):
        '''Get the unnormalised normal vector of a face according to three points'''

        _A, _B, _C = _getABCCoords(_iA, _iB, _iC)

        # plane a*x + b*y + c*z + d = 0
        normal = cross((_B-_A),(_C-_A))
        d = -dot(_A, normal)

        # also return d
        return normal, d


    def _getFurthest0(_iA, _iB, _iC):
        ''''''

        nonlocal coords, indices

        normal_ABC, d = _getNormal0(_iA, _iB, _iC)
        # distance from D0 should be positive
        dists = (dot(coords, normal_ABC)+d)/linalg.norm(normal_ABC)

        return indices[argmax(dists)]


    def _getFurthest(_iA, _iB, _iC, _masks):
        ''''''

        nonlocal indices

        # all distances from further centres should be negative
        dists = _getDistPoints2Plane(_masks, _iA, _iB, _iC)

        return indices[_masks][argmin(dists)]


    # initialise: find the first tetrahedron simplex
    iA0, iB0, iC0 = indices[[*_getA0B0C0(coords)]]
    iD0 = _getFurthest0(iA0, iB0, iC0)


    # initialise the dictionary with the initial four points
    iA, iB, iC, iD = iA0, iB0, iC0, iD0
    for ijk in combinations([iA,iB,iC,iD], 3):
        # YL NB: axis not available utill NumPy v1.13!
        faces = unique(vstack((faces, sort(ijk))), axis=0)
        # initialise the normals
        normal_ABC, d = _getOrientedNormal(*ijk)
        dict_normals.update({tuple(sort(ijk)):_getOrientedNormal(*ijk)})
        # get the points with negative distances from the current plane
        dict_remaining.update({tuple(sort(ijk)):_getSemisphere(*ijk)})

    while True:

        stop = True

        # loop over remaining faces to expand the simplex until no remaining points left
        for face in faces:

            semi_ABC = dict_remaining[tuple(sort(face))]
            if not semi_ABC.sum():
                continue
            else:
                stop = False

            # find the furthest point (having a negative distance)
            iD = _getFurthest(*face, semi_ABC)

            # recursively find valid edges
            _edges  = array([], dtype=int).reshape((0,2))
            _faces = array([], dtype=int).reshape((0,3))
            _getEdges(iD, *face)

            # remove the recursively walked-through adjacent faces from the list of faces
            _delPlanes(_faces)

            # loop over edges to form new faces using D
            for edge in _edges:
                dij = [iD,*edge]
                # add to the list of faces
                # YL NB: axis not available utill NumPy v1.13!
                faces = unique(vstack((faces, sort(dij))), axis=0)
                # update the normal
                dict_normals.update({tuple(sort(dij)):_getOrientedNormal(*dij)})
                # update candidate points with negative distances
                dict_remaining.update({tuple(sort(dij)):_getSemisphere(*dij)})

        if stop:
            break

    # summarise Euler characteristic
    vertices  = unique(faces.ravel())
    nvertices = vertices.size
    edges = unique(sort(vstack((faces[:,[0,1]],
                                faces[:,[0,2]],
                                faces[:,[1,2]])),
                        axis=1),
                   axis=0)
    nedges = edges.shape[0]
    nfaces = faces.shape[0]
    print("\n Vertices found using quickhull (doi:10.1145/235815.235821):\n")
    print(" nvertices           : {}".format(nvertices))
    print(" nedges              : {}".format(nedges))
    print(" nfaces              : {}".format(nfaces))
    print(" Euler characteristic: {}\n".format(nvertices-nedges+nfaces))

    if return_edges_and_faces:
        if return_vertices:
            # YL 13/01/2022: also return the normals and ds
            normals = empty(shape=(0,3))
            ds      = empty(shape=(0,))
            indices_faces = empty(shape=(0,3), dtype=int)
            list_faces = list(map(tuple, faces))
            for face, (normal, d) in dict_normals.items():
                if face not in list_faces:
                    continue
                normals = vstack((normals, normal))
                ds      = append(ds, d)
                indices_faces = vstack((indices_faces,face))
            # 18/01/2022: for now we don't return indices_faces
            return coords[vertices], vertices, searchsorted(vertices, edges), searchsorted(vertices, faces), normals, ds
        else:
            return coords[vertices], searchsorted(vertices, edges), searchsorted(vertices, faces), normals, ds
    else:
        if return_vertices:
            return coords[vertices], vertices
        else:
            return coords[vertices]


def getKabschMatrix(coords1, coords2):
    '''An implementation of the Kabsch algorithm for calculating the optimal rotation matrix that minimizes the RMSD between two paired sets of points, see:
           Kabsch, Wolfgang (1976). "A solution for the best rotation to relate two sets of vectors". Acta Crystallographica. A32 (5): 922. DOI:10.1107/S0567739476001873
           Kabsch, Wolfgang (1978). "A discussion of the solution for the best rotation to relate two sets of vectors". Acta Crystallographica. A34 (5): 827–828. DOI:10.1107/S0567739478001680.
    '''

    from numpy import dot, linalg, mean, sign

    # step 1: a translation to align the centoids to the origin
    centroid1 = mean(coords1, axis=0)
    centroid2 = mean(coords2, axis=0)
    coords1 -= centroid1
    coords2 -= centroid2

    # step 2: the covariance matrix (real)
    H = coords1.T @ coords2

    # step 3: the optimal rotation matrix
    # 3a: calculate the SVD of H because the latter may not have an inverse
    #     u @ np.diag(s) @ vh = (u * s) @ vh
    u, s, vh = linalg.svd(H)

    # ensure a right-handed coordinate system
    d = sign(linalg.det(vh.T @ u.T))

    # the rotation matrix
    R = vh.T @ [[1,0,0],[0,1,0],[0,0,d]] @ u.T

    return R










