"""
Code from https://github.com/ahojnnes/numpy-snippets/blob/master/transformations.py

"""

import math
import warnings

import numpy as np


def make_affine(src, dst):
    """
    Determine parameters of 2D or 3D affine transformation in the order:
        a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3
    where the transformation is defined as:
        X = a0 + a1*x + a2*y[ + a3*z]
        Y = b0 + b1*x + b2*y[ + b3*z]
        [Z = c0 + c1*x + c2*y + c3*z]
    You can determine the over-, well- and under-determined parameters
    with the least-squares method.
    Source and destination coordinates must be Nx2 or Nx3 matrices (x, y, z).
    Explicit parameters are in the order:
        a0, b0, c0, mx, my, mz, alpha [radians], beta [radians], gamma [radians]
    where the 3D transformation is defined as (excluding the :
        X = tx * R3(gamma)*R2(beta)*R1(alpha)*S*x
        with
            X = (X, Y, Z).T
            tx = (a0, b0, c0).T
            R1(alpha) = rotation_matrix(alpha, axis=1)
            R2(beta) = rotation_matrix(beta, axis=2)
            R3(gamma) = rotation_matrix(gamma, axis=3)
            S = diag(mx, my, mz)
            x = (x, y, z).T
    and the simplified 2D transformation as:
        X = a0 + mx*x*cos(alpha) - my*y*sin(alpha+beta)
        Y = b0 + mx*x*sin(alpha) + my*y*cos(alpha+beta)
    In case of 2D coordinates the following parameters are 0:
        a3, b3, c0, c1, c2, c3
    and the explicit parameters
        c0, mz, gamma
    :param src: :class:`numpy.array`
        Nx2 or Nx3 coordinate matrix of source coordinate system
    :param dst: :class:`numpy.array`
        Nx2 or Nx3 coordinate matrix of destination coordinate system
    :returns: params, params_explicit
    """

    xs = src[:, 0]
    ys = src[:, 1]
    rows = src.shape[0]
    A = np.zeros((rows * 3, 12))
    A[:rows, 0] = 1
    A[:rows, 1] = xs
    A[:rows, 2] = ys
    A[rows : rows * 2, 4] = 1
    A[rows : rows * 2, 5] = xs
    A[rows : rows * 2, 6] = ys
    A[rows * 2 :, 8] = 1
    A[rows * 2 :, 9] = xs
    A[rows * 2 :, 10] = ys
    b = np.zeros((rows * 3,))
    b[:rows] = dst[:, 0]
    b[rows : rows * 2] = dst[:, 1]
    if src.shape[1] == 3:
        zs = src[:, 2]
        A[:rows, 3] = zs
        A[rows : rows * 2, 7] = zs
        A[rows * 2 :, 11] = zs
        b[rows * 2 :] = dst[:, 2]
    params = np.linalg.lstsq(A, b, rcond=None)[0]
    #: determine explicit parameters
    a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3 = params
    mx = math.sqrt(a1**2 + b1**2 + c1**2)
    my = math.sqrt(a2**2 + b2**2 + c2**2)
    mz = math.sqrt(a3**2 + b3**2 + c3**2)
    if src.shape[1] == 3 and 0 in (mx, my, mz):
        warnings.warn(
            "One of your scale factors are 0, you should probably "
            "use a 2D instead of a 3D affine transformation.",
            RuntimeWarning,
        )
    alpha = math.atan2(c2, c3)
    beta = math.atan2(-c1, math.sqrt(a1**2 + b1**2))
    gamma = math.atan2(b1, a1)
    params_explicit = np.array([a0, b0, c0, mx, my, mz, alpha, beta, gamma])
    return params, params_explicit


def affine_transform(coords, params, inverse=False):
    """
    Apply 2D or 3D affine transformation.
    :param coords: :class:`numpy.array`
        Nx2 coordinate matrix of source coordinate system
    :param params: :class:`numpy.array`
        parameters returned by `make_tform`
    :param inverse: bool
        apply inverse transformation, default is False
    :returns: :class:`numpy.array`
        transformed coordinates
    """

    a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3 = params
    x = coords[:, 0]
    y = coords[:, 1]
    out = np.zeros(coords.shape)
    if coords.shape[1] == 2:
        z = 0
        if inverse:
            out[:, 0] = (a2 * (y - b0) - b2 * (x - a0)) / (a2 * b1 - a1 * b2)
            out[:, 1] = (b1 * (x - a0) - a1 * (y - b0)) / (a2 * b1 - a1 * b2)
        else:
            out[:, 0] = a0 + a1 * x + a2 * y
            out[:, 1] = b0 + b1 * x + b2 * y
    elif coords.shape[1] == 3:
        z = coords[:, 2]
        if inverse:
            out[:, 0] = (
                (b2 * c3 - b3 * c2) * x
                - (b2 * c3 - b3 * c2) * a0
                + ((z - c0) * b3 - y * c3 + b0 * c3) * a2
                - ((z - c0) * b2 - y * c2 + b0 * c2) * a3
            ) / (
                (b2 * c3 - b3 * c2) * a1
                - (b1 * c3 - b3 * c1) * a2
                + (b1 * c2 - b2 * c1) * a3
            )
            out[:, 1] = -(
                (b1 * c3 - b3 * c1) * x
                - (b1 * c3 - b3 * c1) * a0
                + ((z - c0) * b3 - y * c3 + b0 * c3) * a1
                - ((z - c0) * b1 - y * c1 + b0 * c1) * a3
            ) / (
                (b2 * c3 - b3 * c2) * a1
                - (b1 * c3 - b3 * c1) * a2
                + (b1 * c2 - b2 * c1) * a3
            )
            out[:, 2] = (
                (b1 * c2 - b2 * c1) * x
                - (b1 * c2 - b2 * c1) * a0
                + ((z - c0) * b2 - y * c2 + b0 * c2) * a1
                - ((z - c0) * b1 - y * c1 + b0 * c1) * a2
            ) / (
                (b2 * c3 - b3 * c2) * a1
                - (b1 * c3 - b3 * c1) * a2
                + (b1 * c2 - b2 * c1) * a3
            )
        else:
            out[:, 0] = a0 + a1 * x + a2 * y + a3 * z
            out[:, 1] = b0 + b1 * x + b2 * y + b3 * z
            out[:, 2] = c0 + c1 * x + c2 * y + c3 * z
    return out
