"""A file to contain different kinds of lattice generation algorithms.
"""

import numba
import numpy as np

from scipy.stats.qmc import LatinHypercube
from scipy.special import comb
from itertools import combinations


@numba.njit()
def fibonacci_sphere(samples: int = 1000) -> np.ndarray:
    """Generate a very even lattice of points on a 3d sphere using the fibonacci sphere
    or fibonacci spiral algorithm.

    Args:
        samples (int, optional): Number of points to be generated. Defaults to 1000.

    Returns:
        np.ndarray: The lattice of points as a 2-D (samples, 3) numpy array.
    """
    points = np.zeros((samples, 3), dtype=np.float_)
    phi = np.pi * (3 - np.sqrt(5))  # golden angle in radians

    for i in range(samples):
        points[i, 1] = 1 - (i / (samples - 1)) * 2  # y goes from 1 to -1
        radius = np.sqrt(1 - points[i, 1] ** 2)  # radius at y

        theta = phi * i  # golden angle increment

        points[i, 0] = np.cos(theta) * radius
        points[i, 2] = np.sin(theta) * radius
    return points


def sphereLHS(num_dimensions: int, num_points: int) -> tuple[np.ndarray, np.ndarray]:
    """Generate an even distribution of points on a spherical surface using the Latin Hypercube Sampling.

    The points are generated by first generating a Latin Hypercube sample in num_dimensions - 1 dimensions. These points
    are treated as the polar coordinates of the points on a sphere of radius 1. The points are then converted to
    Cartesian coordinates. The sphere is centered at the origin, and the points are distributed in the positive orthant.
    Note: To center the points on the nadir point, use the function (X) -> 1 - X.

    Args:
        num_dimensions (int): Number of dimensions of the sphere.
        num_points (int): Number of points to be generated.

    Returns:
        tuple[np.ndarray, np.ndarray]: The lattice of points as a 2-D (num_points, num_dimensions) numpy array. The
            angles of the points in radians as a 2-D (num_points, num_dimensions - 1) numpy array.
    """
    angles = LatinHypercube(d=num_dimensions - 1).random(num_points) * np.pi / 2
    points = np.zeros((num_points, num_dimensions))
    for i in range(num_dimensions):
        if i == num_dimensions - 1:
            points[:, i] = 1
        else:
            points[:, i] = np.cos(angles[:, i])
        for j in range(i):
            points[:, i] *= np.sin(angles[:, j])

    return points, angles


def simplexLatticeDesign(num_dimensions: int, lattice_resolution: int) -> np.ndarray:
    """Create a simplex lattice design. The generated lattice is on the unit hyperplane, and the points are distributed
    in the positive orthant. See more information at:
    https://www.itl.nist.gov/div898/handbook/pri/section5/pri542.htm



    Args:
        num_dimensions (int): Number of dimensions of the simplex.
        lattice_resolution (int): Number of divisions per dimension of the simplex.

    Returns:
        np.ndarray: The lattice of points as a 2-D (num_points, num_dimensions) numpy array.
    """

    number_of_vectors = comb(
        lattice_resolution + num_dimensions - 1, num_dimensions - 1, exact=True
    )

    temp1 = range(1, num_dimensions + lattice_resolution)
    temp1 = np.array(list(combinations(temp1, num_dimensions - 1)))
    temp2 = np.array([range(num_dimensions - 1)] * number_of_vectors)
    temp = temp1 - temp2 - 1
    weight = np.zeros((number_of_vectors, num_dimensions), dtype=int)
    weight[:, 0] = temp[:, 0]
    for i in range(1, num_dimensions - 1):
        weight[:, i] = temp[:, i] - temp[:, i - 1]
    weight[:, -1] = lattice_resolution - temp[:, -1]

    return weight / lattice_resolution


def simplexLatticefromNumPoints(
    num_dimensions: int, num_points: int, atleast: bool = False
) -> np.ndarray:
    """Create a simplex lattice design with the given number of points.  The generated lattice is on the unit
    hyperplane, and the points are distributed in the positive orthant. See more information at:
    https://www.itl.nist.gov/div898/handbook/pri/section5/pri542.htm

    The lattice is generated by first finding the largest lattice resolution that can generate a lattice with the
    given number of points. Then the simplexLatticeDesign function is called to generate the lattice.

    Args:
        num_dimensions (int): Number of dimensions of the simplex.
        num_points (int):  Approximate number of points to be generated. The actual number of points generated will be
            defined using the simplex generating algorithm.
        atleast (bool, optional): If True, the number of points in the generated lattice will be at least the given
            number of points. If False, the number of points in the generated lattice will be at most the given number
            of points. Defaults to False.

    Returns:
        np.ndarray: The lattice of points as a 2-D (num_points, num_dimensions) numpy array.
    """
    temp_lattice_resolution = 0
    while True:
        temp_lattice_resolution += 1
        number_of_vectors = comb(
            temp_lattice_resolution + num_dimensions - 1, num_dimensions - 1, exact=True
        )
        if number_of_vectors > num_points:
            break
    if atleast:
        return simplexLatticeDesign(num_dimensions, temp_lattice_resolution)
    return simplexLatticeDesign(num_dimensions, temp_lattice_resolution - 1)


def normalize(vectors):
    """
    Normalize a set of vectors.

    The length of the returned vectors will be unity.

    Parameters
    ----------
    vectors : np.ndarray
        Set of vectors of any length, except zero.

    """
    if len(np.asarray(vectors).shape) == 1:
        return vectors / np.linalg.norm(vectors)
    norm = np.linalg.norm(vectors, axis=1)
    return vectors / norm[:, np.newaxis]


def rotate(initial_vector, rotated_vector, other_vectors):
    """Calculate the rotation matrix that rotates the initial_vector to the
    rotated_vector. Apply that rotation on other_vectors and return.
    Uses Householder reflections twice to achieve this."""

    init_vec_norm = normalize(initial_vector)
    rot_vec_norm = normalize(np.asarray(rotated_vector))
    middle_vec_norm = normalize(init_vec_norm + rot_vec_norm)
    first_reflector = init_vec_norm - middle_vec_norm
    second_reflector = middle_vec_norm - rot_vec_norm
    Q1 = householder(first_reflector)
    Q2 = householder(second_reflector)
    reflection_matrix = np.matmul(Q2, Q1)
    rotated_vectors = np.matmul(other_vectors, np.transpose(reflection_matrix))
    return rotated_vectors


def householder(vector):
    """Return reflection matrix via householder transformation."""
    identity_mat = np.eye(len(vector))
    v = vector[np.newaxis]
    denominator = np.matmul(v, v.T)
    numerator = np.matmul(v.T, v)
    rot_mat = identity_mat - (2 * numerator / denominator)
    return rot_mat


def rotate_toward(initial_vector, final_vector, other_vectors, degrees: float = 5):
    """
    Rotate other_vectors (with the centre at initial_vector) towards final_vector
    by an angle degrees.

    Parameters
    ----------
    initial_vector : np.ndarray
        Centre of the vectors to be rotated.
    final_vector : np.ndarray
        The final position of the center of other_vectors.
    other_vectors : np.ndarray
        The array of vectors to be rotated
    degrees : float, optional
        The amount of rotation (the default is 5)

    Returns
    -------
    rotated_vectors : np.ndarray
        The rotated vectors
    reached: bool
        True if final_vector has been reached
    """
    final_vector = normalize(final_vector)
    initial_vector = normalize(initial_vector)
    cos_phi = np.dot(initial_vector, final_vector)
    theta = degrees * np.pi / 180
    cos_theta = np.cos(theta)
    phi = np.arccos(cos_phi)
    if phi < theta:
        return (rotate(initial_vector, final_vector, other_vectors), True)
    cos_phi_theta = np.cos(phi - theta)
    A = np.asarray([[cos_phi, 1], [1, cos_phi]])
    B = np.asarray([cos_phi_theta, cos_theta])
    x = np.linalg.solve(A, B)
    rotated_vector = x[0] * initial_vector + x[1] * final_vector
    return (rotate(initial_vector, rotated_vector, other_vectors), False)
