"""
GEPARD - Gepard-Enabled PARticle Detection
Copyright (C) 2018  Lars Bittrich and Josef Brandt, Leibniz-Institut für
Polymerforschung Dresden e. V. <bittrich-lars@ipfdd.de>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program, see COPYING.
If not, see <https://www.gnu.org/licenses/>.
"""
import numpy as np
from scipy.optimize import least_squares


def getRotMat(angles: np.ndarray) -> np.matrix:
    """
    :param angles: np.array shape(3), rot in x, y, z in radians
    """
    c1, s1 = np.cos(angles[0]), np.sin(angles[0])
    c2, s2 = np.cos(angles[1]), np.sin(angles[1])
    c3, s3 = np.cos(angles[2]), np.sin(angles[2])
    return np.mat([[c1*c3 - s1*c2*s3, -c1*s3 - s1*c2*c3, s1*s2],
                   [s1*c3 + c1*c2*s3, -s1*s3 + c1*c2*c3, -c1*s2],
                   [s1*s3, s2*c3, c2]]).T


def applyTransformToPoints(rotMat: np.matrix, shift: np.ndarray, srcPoints: np.ndarray) -> np.ndarray:
    center: np.ndarray = np.mean(srcPoints, axis=0)
    pointsCentered: np.ndarray = srcPoints-center
    pointsRot: np.matrix = np.dot(pointsCentered, rotMat).A
    return pointsRot + center + shift


def getTransform(srcPoints: np.ndarray, dstPoints: np.ndarray, invertX: bool = False, invertY: bool = False,
                 invertZ: bool = False):
    """
    Calculates transforms to convert from SOURCE to DESTINATION Coordinate System
    :param srcPoints: np.ndarray (Nx3) of N points in 3D (x, y, z) --> SOURCE Coordinate System
    :param dstPoints: np.ndarray (Nx3) of N points in 3D (x, y, z) --> DESTINATION Coordinate System
    :param invertX: bool, wether or not to flip the x-axis
    :param invertY: bool, wether or not to flip the y-axis
    :param invertZ: bool, wether or not to flip the z-axis
    :return rotationMatrix: np.matrix (3x3) for rotation
    :return shift: np.array shape(3), shift in x, y, z
    :return accept: bool: Whether the result was accepted.
    """

    Parity = np.mat(np.diag([-1. if invertX else 1.,
                             -1. if invertY else 1.,
                             -1. if invertZ else 1.]))

    srcCenter: np.ndarray = srcPoints.mean(axis=0)
    dstCenter: np.ndarray = dstPoints.mean(axis=0)
    srcPointsCentered: np.ndarray = srcPoints - srcCenter
    dstPointsCentered: np.ndarray = dstPoints - dstCenter

    def angleErrors(_angles: np.ndarray, _srcCentered: np.ndarray, _dstCentered: np.ndarray) -> np.ndarray:
        """
        :param _angles: shape (3) array of radian angles in x, y, z
        :param _srcCentered: shape (N, 3) array of centered source points
        :param _dstCentered: shape (N, 3) array of centered destination points
        """
        R = (getRotMat(_angles) * Parity).A
        return (np.dot(_srcCentered, R) - _dstCentered).ravel()

    angleOptFun = lambda x: angleErrors(x, srcPointsCentered, dstPointsCentered)
    opt = least_squares(angleOptFun, np.zeros(3), bounds=(np.array([-np.pi/4]*3), np.array([np.pi/4]*3)),
                        method='dogbox')
    optangles = opt.x
    rotMat: np.matrix = getRotMat(optangles) * Parity
    shift: np.ndarray = dstCenter - srcCenter
    errors: np.ndarray = applyTransformToPoints(rotMat, shift, srcPoints) - dstPoints
    d = np.linalg.norm(errors, axis=1)
    accept = True
    if np.any(d > 1.0):
        print(f'Transformation residuals are large:{d}')
        accept = False
    return rotMat, shift, accept