"""
Implements various means of generating rotation matrices.

Copyright (c) 2023-2025 European Molecular Biology Laboratory

Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
"""

import yaml
import warnings
from typing import Tuple
from os.path import join, dirname

import numpy as np
from scipy.spatial.transform import Rotation

from .types import NDArray

__all__ = [
    "get_cone_rotations",
    "align_vectors",
    "euler_to_rotationmatrix",
    "euler_from_rotationmatrix",
    "get_rotation_matrices",
    "align_to_axis",
]


def _sample_cone(
    angle: float, sampling: float, axis: Tuple[float] = (1, 0, 0)
) -> NDArray:
    """
    Sample points on a cone surface around cone_axis using golden spiral distribution.

    Parameters
    ----------
    angle : float
        The half-angle of the cone in degrees.
    sampling : float
        Angular increment used for sampling points in degrees.
    axis : tuple of floats
        Vector to align the cone with.

    Returns
    -------
    NDArray
        Array of points around axis with shape n,3.

    References
    ----------
    .. [1] https://stackoverflow.com/questions/9600801/evenly-distributing-n-points-on-a-sphere
    """
    theta = np.linspace(0, angle, round(angle / sampling) + 1)
    number_of_points = np.ceil(
        360 * np.divide(np.sin(np.radians(theta)), sampling),
    )
    number_of_points = int(np.sum(number_of_points + 1) + 2)

    indices = np.arange(0, number_of_points, dtype=float) + 0.5
    radius = np.radians(angle * np.sqrt(indices / number_of_points))
    theta = np.pi * (1 + np.sqrt(5)) * indices

    points = np.stack(
        [
            np.cos(radius),
            np.cos(theta) * np.sin(radius),
            np.sin(theta) * np.sin(radius),
        ],
        axis=1,
    )

    rotation = Rotation.from_euler(
        angles=align_vectors((1, 0, 0), axis, seq="zyz"),
        seq="zyz",
        degrees=True,
    )
    return rotation.apply(points)


def get_cone_rotations(
    cone_angle: float,
    cone_sampling: float,
    axis: Tuple[float] = (1, 0, 0),
    axis_angle: float = 360.0,
    axis_sampling: float = None,
    reference: Tuple[float] = (1, 0, 0),
    n_symmetry: int = 1,
    seq: str = None,
) -> NDArray:
    """
    Generate rotations describing the possible placements of a vector in a cone.

    Parameters
    ----------
    cone_angle : float
        The half-angle of the cone in degrees.
    cone_sampling : float
        Angular increment used for sampling points on the cone in degrees.
    axis : Tuple[float], optional
        Base-vector of the cone.
    axis_angle : float, optional
        The total angle of rotation around the cone axis in degrees (default is 360.0).
    axis_sampling : float, optional
        Angular increment used for sampling points around the cone axis in degrees.
        If None, it takes the value of cone_sampling.
    reference : Tuple[float], optional
        Returned rotations will map this point onto the cone. In practice, this is
        the principal axis of the template.
    n_symmetry : int, optional
        Number of symmetry axis around the vector axis.
    seq : str, optional
        Convention for angles. By default returns rotation matrices.

    Returns
    -------
    NDArray
        An arary of rotations represented as stack of rotation matrices when convention
        None (n, 3, 3) or an array of Euler angles (n, 3) for available conventions.
    """
    if axis_sampling is None:
        axis_sampling = cone_sampling

    points = _sample_cone(angle=cone_angle, sampling=cone_sampling, axis=axis)
    reference = np.asarray(reference).astype(np.float32)
    reference /= np.linalg.norm(reference)

    axis_angle /= n_symmetry
    phi_steps = np.maximum(np.round(axis_angle / axis_sampling), 1).astype(int)
    phi = np.linspace(0, axis_angle, phi_steps + 1)[:-1]

    axis_rotation = Rotation.from_rotvec(axis * np.radians(phi)[:, None])
    all_rotations = [
        axis_rotation * Rotation.from_matrix(align_vectors(reference, x))
        for x in points
    ]

    rotations = Rotation.concatenate(all_rotations)
    if seq is None:
        return rotations.as_matrix()

    return rotations.as_euler(seq=seq, degrees=True)


def align_vectors(base: NDArray, target: NDArray = (0, 0, 1), seq: str = None):
    """
    Compute the rotation matrix or Euler angles required to align an initial vector with a target vector.

    Parameters
    ----------
    base : NDArray
        The basis vector.
    target : NDArray, optional
        The vector to map base to, defaults to (0,0,1).
    seq : str, optional
        Euler angle convention, None returns a rotation matrix instead.

    Returns
    -------
    NDArray
        Rotation matrix if seq is None otherwise Euler angles in desired convention
    """
    base = np.asarray(base, dtype=np.float32)
    target = np.asarray(target, dtype=np.float32)

    rotation, error = Rotation.align_vectors(target, base)
    if seq is None:
        return rotation.as_matrix()
    return rotation.as_euler(seq=seq, degrees=True)


def euler_to_rotationmatrix(angles: Tuple[float], seq: str = "zyz") -> NDArray:
    """
    Convert Euler angles to a rotation matrix.

    Parameters
    ----------
    angles : tuple
        A tuple representing the Euler angles in degrees.
    seq : str, optional
        Euler angle convention.

    Returns
    -------
    NDArray
        The generated rotation matrix.
    """
    angles = np.asarray(angles)

    n_angles = len(angles)
    if angles.ndim == 2:
        n_angles = angles.shape[1]

    rotation_matrix = Rotation.from_euler(
        seq=seq[:n_angles], angles=angles, degrees=True
    )
    return rotation_matrix.as_matrix().astype(np.float32)


def euler_from_rotationmatrix(rotation_matrix: NDArray, seq: str = "zyz") -> Tuple:
    """
    Convert a rotation matrix to euler angles.

    Parameters
    ----------
    rotation_matrix : NDArray
        A 2 x 2 or 3 x 3 rotation matrix.
    seq : str, optional
        Euler angle convention, zyz by default.

    Returns
    -------
    Tuple
        The generate euler angles in degrees
    """
    if rotation_matrix.shape[0] == 2:
        temp_matrix = np.eye(3)
        temp_matrix[:2, :2] = rotation_matrix
        rotation_matrix = temp_matrix

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        rotation = Rotation.from_matrix(rotation_matrix)
        angles = rotation.as_euler(seq=seq, degrees=True).astype(np.float32)
    return angles


def get_rotation_matrices(
    angular_sampling: float, dim: int = 3, use_optimized_set: bool = True
) -> NDArray:
    """
    Returns rotation matrices with desired ``angular_sampling`` rate.

    Parameters
    ----------
    angular_sampling : float
        The desired angular sampling in degrees.
    dim : int, optional
        Dimension of the rotation matrices.
    use_optimized_set : bool, optional
        Use optimized rotational sets, True by default and available for dim=3.

    Notes
    -----
    For dim = 3 optimized sets are used, otherwise QR-decomposition.

    Returns
    -------
    NDArray
        Array of shape (n, d, d) containing n rotation matrices.
    """
    if dim == 3 and use_optimized_set:
        quaternions, *_ = _load_quaternions_by_angle(angular_sampling)
        ret = Rotation.from_quat(quaternions, scalar_first=True).as_matrix()
    else:
        num_rotations = dim * (dim - 1) // 2
        k = int((360 / angular_sampling) ** num_rotations)
        As = np.random.randn(k, dim, dim)
        ret, _ = np.linalg.qr(As)
        dets = np.linalg.det(ret)
        neg_dets = dets < 0
        ret[neg_dets, :, -1] *= -1
        ret[0] = np.eye(dim, dtype=ret.dtype)
    return ret


def _load_quaternions_by_angle(
    angular_sampling: float,
) -> Tuple[NDArray, NDArray, float]:
    """
    Get orientations and weights proportional to the given angular_sampling.

    Parameters
    ----------
    angular_sampling : float
        Requested angular sampling.

    Returns
    -------
    Tuple[NDArray, NDArray, float]
        Quaternions, associated weights and realized angular sampling rate.
    """
    # Metadata contains (N orientations, rotational sampling, coverage as values)
    with open(join(dirname(__file__), "data", "metadata.yaml"), "r") as infile:
        metadata = yaml.full_load(infile)

    set_diffs = {
        setname: abs(angular_sampling - set_angle)
        for setname, (_, set_angle, _) in metadata.items()
    }
    fname = min(set_diffs, key=set_diffs.get)

    infile = join(dirname(__file__), "data", fname)
    quat_weights = np.load(infile)

    quat = quat_weights[:, :4]
    weights = quat_weights[:, -1]
    angle = metadata[fname][0]

    return quat, weights, angle


def align_to_axis(
    coordinates: NDArray,
    weights: NDArray = None,
    axis: int = 2,
    flip: bool = False,
    eigenvector_index: int = 0,
) -> NDArray:
    """
    Calculate a rotation matrix that aligns the principal axis of a point cloud
    with a specified coordinate axis.

    Parameters
    ----------
    coordinates : NDArray
        Array of 3D coordinates with shape (n, 3) representing the point cloud.
    weights : NDArray
        Coordinate weighting factors with shape (n,).
    axis : int, optional
        The target axis to align with, defaults to 2 (z-axis).
    flip : bool, optional
        Whether to align with the negative direction of the axis, default is False.
    eigenvector_index : int, optional
        Index of eigenvector to select, sorted by descending eigenvalues.
        0 = largest eigenvalue (most variance), 1 = second largest, etc.
        Default is 0 (primary principal component).

    Returns
    -------
    NDArray
        3x3 rotation matrix that aligns the principal component of the
        coordinates with the specified axis.
    """
    axis = int(axis)
    coordinates = np.asarray(coordinates)
    alignment_axis = np.array(
        [0 if i != axis else 1 for i in range(coordinates.shape[1])]
    )
    if flip:
        alignment_axis *= -1

    ndim = coordinates.shape[1]
    if eigenvector_index >= ndim:
        raise ValueError(f"eigenvector_index has to be less than {ndim}.")

    avg = np.average(coordinates, axis=0, weights=weights)
    coordinates = coordinates - avg
    cov_matrix = np.cov(coordinates.T, aweights=weights)

    # Eigenvalues are already sorted in ascending order
    eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
    eigenvector = eigenvectors[:, -(eigenvector_index + 1)]
    return align_vectors(eigenvector, alignment_axis)
