# -*- coding: utf-8 -*-
"""
Linear algebra operations and helpers.

Inspired by Christoph Gohlke's transformation.py <http://www.lfd.uci.edu/~gohlke/>

This module is not directly exported by the `crystals` library. Use it with caution.
"""

import math
from typing import Tuple

import numpy as np
from numpy.typing import ArrayLike

Basis = Tuple[ArrayLike, ArrayLike, ArrayLike]

# standard basis
e1, e2, e3 = np.eye(3)


def affine_map(array: ArrayLike) -> np.ndarray:
    """
    Extends 3x3 transform matrices to 4x4, i.e. general affine transforms.

    Parameters
    ----------
    array : array-like, shape {(3,3), (4,4)}
        Transformation matrix. If shape = (4,4), returned intact.

    Returns
    -------
    extended : ndarray, shape (4,4)
        Extended array

    Raises
    ------
    ValueError : If the transformation matrix is neither 3x3 or 4x4
    """
    array = np.asarray(array)
    if array.shape == (4, 4):  # Already the right shape
        return array
    elif array.shape == (3, 3):
        extended_matrix = np.zeros(shape=(4, 4), dtype=array.dtype)
        extended_matrix[-1, -1] = 1
        extended_matrix[:3, :3] = array
        return extended_matrix
    else:
        raise ValueError(
            "Array shape not 3x3 or 4x4, and thus is not a transformation matrix."
        )


def transform(matrix: ArrayLike, array: ArrayLike) -> np.ndarray:
    """
    Applies a matrix transform on an array.

    Parameters
    ----------
    matrix : array-like, shape {(3,3), (4,4)}
        Transformation matrix.
    array : array-like, shape {(3,), (3,3), (4,4)}
        Array to be transformed. Either a 1x3 vector, or a transformation
        matrix in 3x3 or 4x4 shape.

    Returns
    -------
    transformed : ndarray
        Transformed array, either a 1D vector or a 4x4 transformation matrix

    Raises
    ------
    ValueError : If the transformation matrix is neither 3x3 or 4x4
    """
    array = np.asarray(array)
    matrix = np.asarray(matrix)

    if matrix.shape not in [(3, 3), (4, 4)]:
        raise ValueError(
            f"Input matrix is neither a 3x3 or 4x4 matrix, but \
						  rather of shape {matrix.shape}."
        )

    matrix = affine_map(matrix)
    # Case of a vector (e.g. position vector):
    if array.ndim == 1:
        extended_vector = np.array([0, 0, 0, 1], dtype=array.dtype)
        extended_vector[:3] = array
        return np.dot(matrix, extended_vector)[:3]
    else:
        array = affine_map(array)
    return np.dot(matrix, array)


def translation_matrix(direction: ArrayLike) -> np.ndarray:
    """
    Return matrix to translate by direction vector.

    Parameters
    ----------
    direction : array_like, shape (3,)

    Returns
    -------
    translation : `~numpy.ndarray`, shape (4,4)
        4x4 translation matrix.
    """
    matrix = np.eye(4)
    matrix[:3, 3] = np.asarray(direction)[:3]
    return matrix


def change_of_basis(basis1: Basis, basis2: Basis = (e1, e2, e3)) -> np.ndarray:
    """
    Returns the matrix transforms vectors expressed in one basis,
    to vectors expressed in another basis.

    Parameters
    ----------
    basis1 : list of array_like, shape (3,)
        First basis
    basis2 : list of array_like, shape (3,), optional
        Second basis. By default, this is the standard basis

    Returns
    -------
    cob : `~numpy.ndarray`, shape (3,3)
        Change-of-basis matrix.
    """
    # Calculate the transform that goes from basis 1 to standard basis
    basis1 = [np.asarray(vector).reshape(3, 1) for vector in basis1]
    basis1_to_standard = np.hstack(tuple(basis1))

    # Calculate the transform that goes from standard basis to basis2
    basis2 = [np.asarray(vector).reshape(3, 1) for vector in basis2]
    standard_to_basis2 = np.linalg.inv(np.hstack(tuple(basis2)))

    return np.dot(standard_to_basis2, basis1_to_standard)


def is_basis(basis: Basis) -> bool:
    """
    Returns true if the set of vectors forms a basis. This is done by checking
    whether basis vectors are independent via an eigenvalue calculation.

    Parameters
    ----------
    basis : list of array-like, shape (3,)

    Returns
    -------
    out : bool
        Whether or not the basis is valid.
    """
    return 0 not in np.linalg.eigvals(np.asarray(basis))


def is_rotation_matrix(matrix: ArrayLike) -> bool:
    """
    Checks whether a matrix is orthogonal with unit determinant (1 or -1), properties
    of rotation matrices.

    Parameters
    ----------
    matrix : array-like, shape {(3,3), (4,4)}
        Rotation matrix candidate. If (4,4) matrix is provided,
        only the top-left block matrix of (3,) is checked

    Returns
    -------
    result : bool
        If True, input could be a rotation matrix.
    """
    # TODO: is this necessary? should a composite transformation
    # 		of translation and rotation return True?
    # if matrix.shape == (4,4):
    #    matrix = matrix[:3,:3]
    matrix = np.asarray(matrix)
    is_orthogonal = np.allclose(np.linalg.inv(matrix), np.transpose(matrix))
    unit_determinant = np.allclose(abs(np.linalg.det(matrix)), 1)
    return is_orthogonal and unit_determinant


def rotation_matrix(angle: float, axis: ArrayLike = (0, 0, 1)) -> np.ndarray:
    """
    Return matrix to rotate about axis defined by direction around the origin [0,0,0].

    Parameters
    ----------
    angle : float
        Rotation angle [rad]
    axis : array-like of length 3
        Axis about which to rotate

    Returns
    -------
    matrix : `~numpy.ndarray`, shape (3,3)
        Rotation matrix.

    See also
    --------
    translation_rotation_matrix

    Notes
    -----
    To combine rotation and translations, see
    http://www.euclideanspace.com/maths/geometry/affine/matrix4x4/index.htm
    """
    sina, cosa = math.sin(angle), math.cos(angle)

    # Make sure direction is a numpy vector of unit length
    direction = np.asarray(axis)
    direction = direction / np.linalg.norm(direction)

    # rotation matrix around unit vector
    R = np.diag([cosa, cosa, cosa])
    R += np.outer(direction, direction) * (1.0 - cosa)
    direction *= sina
    R += np.array(
        [
            [0.0, -direction[2], direction[1]],
            [direction[2], 0.0, -direction[0]],
            [-direction[1], direction[0], 0.0],
        ]
    )

    return R


def translation_rotation_matrix(
    angle: float, axis: ArrayLike, translation: ArrayLike
) -> np.ndarray:
    """
    Returns a 4x4 matrix that includes a rotation and a translation.

    Parameters
    ----------
    angle : float
        Rotation angle [rad]
    axis : array-like of length 3
        Axis about which to rotate
    translation : array_like, shape (3,)
        Translation vector

    Returns
    -------
    matrix : `~numpy.ndarray`, shape (4,4)
        Affine transform matrix.
    """
    rmat = affine_map(rotation_matrix(angle=angle, axis=axis))
    rmat[:3, 3] = np.asarray(translation)
    return rmat


def change_basis_mesh(
    xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, basis1: Basis, basis2: Basis
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Changes the basis of meshgrid arrays.

    Parameters
    ----------
    xx, yy, zz : ndarrays
        Arrays of equal shape, such as produced by numpy.meshgrid.
    basis1 : list of ndarrays, shape(3,)
        Basis of the mesh
    basis2 : list of ndarrays, shape(3,)
        Basis in which to express the mesh

    Returns
    -------
    XX, YY, ZZ : `~numpy.ndarray`
    """
    # Build coordinate array row-wise
    changed = np.empty(shape=(3, xx.size), dtype=float)
    linearized = np.empty(shape=(3, xx.size), dtype=float)
    linearized[0, :] = xx.ravel()
    linearized[1, :] = yy.ravel()
    linearized[2, :] = zz.ravel()

    # Change the basis at each row
    COB = change_of_basis(basis1, basis2)
    np.dot(COB, linearized, out=changed)
    return (
        changed[0, :].reshape(xx.shape),
        changed[1, :].reshape(yy.shape),
        changed[2, :].reshape(zz.shape),
    )


def minimum_image_distance(
    xx: np.ndarray, yy: np.ndarray, zz: np.ndarray, lattice: Basis
) -> np.ndarray:
    """
    Returns a periodic array according to the minimum image convention.

    Parameters
    ----------
    xx, yy, zz : ndarrays
        Arrays of equal shape, such as produced by numpy.meshgrid.
    lattice : list of ndarrays, shape(3,)
        Basis of the mesh

    Returns
    -------
    r : `~numpy.ndarray`
            Minimum image distance over the lattice
    """
    COB = change_of_basis(np.eye(3), lattice)
    linearized = np.empty(shape=(3, xx.size), dtype=float)  # In the standard basis
    ulinearized = np.empty_like(linearized)  # In the unitcell basis

    linearized[0, :] = xx.ravel()
    linearized[1, :] = yy.ravel()
    linearized[2, :] = zz.ravel()

    # Go to unitcell basis, where the cell is cubic of side length 1
    np.dot(COB, linearized, out=ulinearized)
    ulinearized -= np.rint(ulinearized)
    np.dot(np.linalg.inv(COB), ulinearized, out=linearized)

    return np.reshape(np.linalg.norm(linearized, axis=0), xx.shape)
