"""
Module for random matrix factorisations.

This module provides functions for computing low-rank matrix factorisations
using randomized algorithms. These methods are particularly effective for large
matrices where deterministic methods would be too slow. The implementations
are based on the work of Halko, Martinsson, and Tropp (2011).
"""

import numpy as np
from scipy.linalg import (
    cho_factor,
    solve_triangular,
    eigh,
    svd,
    qr,
)
from scipy.sparse.linalg import LinearOperator as ScipyLinOp
from typing import Tuple, Union

# A type for objects that act like matrices (numpy arrays or SciPy LinearOperators)
MatrixLike = Union[np.ndarray, ScipyLinOp]


def fixed_rank_random_range(
    matrix: MatrixLike, rank: int, power: int = 0
) -> np.ndarray:
    """
    Computes an orthonormal basis for a fixed-rank approximation to the
    range of a matrix using a randomized method.

    This is a two-stage algorithm that finds a low-dimensional subspace that
    captures most of the action of the matrix.

    Args:
        matrix (matrix-like): An (m, n) matrix or LinearOperator whose range
            is to be approximated.
        rank (int): The desired rank for the approximation. Must be
            greater than 1.
        power (int): The exponent for power iterations, used to improve the
            accuracy of the approximation.

    Returns:
        numpy.ndarray: An (m, rank) matrix with orthonormal columns whose
            span approximates the range of the input matrix.

    Notes:
        If the input matrix is a scipy LinearOperator, it must have the
        `matmat` and `rmatmat` methods implemented.

        This method is based on Algorithm 4.4 in Halko et al. 2011.
    """

    m, n = matrix.shape
    random_matrix = np.random.randn(n, rank)

    product_matrix = matrix @ random_matrix
    qr_factor, _ = qr(product_matrix, overwrite_a=True, mode="economic")

    for _ in range(power):
        tilde_product_matrix = matrix.T @ qr_factor
        tilde_qr_factor, _ = qr(tilde_product_matrix, overwrite_a=True, mode="economic")
        product_matrix = matrix @ tilde_qr_factor
        qr_factor, _ = qr(product_matrix, overwrite_a=True, mode="economic")

    return qr_factor


def variable_rank_random_range(
    matrix: MatrixLike, rank: int, /, *, power: int = 0, rtol: float = 1e-6
) -> np.ndarray:
    """
    Computes an orthonormal basis for a variable-rank approximation to the
    range of a matrix using a randomized method.

    The algorithm adaptively determines the rank required to meet a given
    error tolerance.

    Args:
        matrix (matrix-like): An (m, n) matrix or LinearOperator whose range
            is to be approximated.
        rank (int): The maximum rank for the approximation. The algorithm
            may return a basis with a smaller rank.
        power (int): Exponent for power iterations. Note: This parameter is
            reserved for future functionality and is currently unused.
        rtol (float): The relative tolerance for the approximation error, used
            to determine the output rank.

    Returns:
        numpy.ndarray: An (m, k) matrix with orthonormal columns, where k <= rank.
            Its span approximates the range of the input matrix.

    Notes:
        If the input matrix is a scipy LinearOperator, it must have the
        `matvec` method implemented.

        This method is based on Algorithm 4.5 in Halko et al. 2011.
    """

    m, n = matrix.shape

    random_vectors = [np.random.randn(n) for _ in range(rank)]
    ys = [matrix @ x for x in random_vectors]
    basis_vectors = []

    def projection(xs: list, y: np.ndarray) -> np.ndarray:
        ps = [np.dot(x, y) for x in xs]
        for p, x in zip(ps, xs):
            y -= p * x
        return y

    norm = max(np.linalg.norm(y) for y in ys)

    tol = rtol * norm / (10 * np.sqrt(2 / np.pi))
    error = 2 * tol
    j = -1
    while error > tol:
        j += 1

        ys[j] = projection(basis_vectors, ys[j])
        ys[j] /= np.linalg.norm(ys[j])
        basis_vectors.append(ys[j])

        y = matrix @ np.random.randn(n)
        y = projection(basis_vectors, y)
        ys.append(y)

        for i in range(j + 1, j + rank):
            p = np.dot(basis_vectors[j], ys[i])
            ys[i] -= p * basis_vectors[j]

        error = max(np.linalg.norm(ys[i]) for i in range(j + 1, j + rank + 1))

        if j > min(n, m):
            raise RuntimeError("Convergence has failed")

    qr_factor = np.column_stack(basis_vectors)

    return qr_factor


def random_svd(
    matrix: MatrixLike, qr_factor: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Computes an approximate Singular Value Decomposition (SVD) from a
    low-rank range approximation.

    Args:
        matrix (matrix-like): The original (m, n) matrix or LinearOperator.
        qr_factor (numpy.ndarray): An (m, k) orthonormal basis for the
            approximate range of the matrix, typically from a
            `random_range` function.

    Returns:
        (numpy.ndarray, numpy.ndarray, numpy.ndarray): A tuple (U, S, Vh)
            containing the approximate SVD factors, such that A ~= U @ S @ Vh.
            S is a 1D array of singular values.

    Notes:
        Based on Algorithm 5.1 of Halko et al. 2011.
    """
    small_matrix = qr_factor.T @ matrix
    left_factor, diagonal_factor, right_factor_transposed = svd(
        small_matrix, full_matrices=False, overwrite_a=True
    )
    return (
        qr_factor @ left_factor,
        diagonal_factor,
        right_factor_transposed,
    )


def random_eig(
    matrix: MatrixLike, qr_factor: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Computes an approximate eigendecomposition for a symmetric matrix from a
    low-rank range approximation.

    Args:
        matrix (matrix-like): The original symmetric (n, n) matrix or
            LinearOperator.
        qr_factor (numpy.ndarray): An (n, k) orthonormal basis for the
            approximate range of the matrix.

    Returns:
        (numpy.ndarray, numpy.ndarray): A tuple (U, S) containing the
            approximate eigenvectors and eigenvalues, such that A ~= U @ S @ U.T.
            S is a 1D array of eigenvalues.

    Notes:
        Based on Algorithm 5.3 of Halko et al. 2011.
    """
    m, n = matrix.shape
    assert m == n
    small_matrix = qr_factor.T @ matrix @ qr_factor
    eigenvalues, eigenvectors = eigh(small_matrix, overwrite_a=True)
    return qr_factor @ eigenvectors, eigenvalues


def random_cholesky(matrix: MatrixLike, qr_factor: np.ndarray) -> np.ndarray:
    """
    Computes an approximate Cholesky factorisation for a symmetric positive-
    definite matrix from a low-rank range approximation.

    Args:
        matrix (matrix-like): The original symmetric positive-definite (n, n)
            matrix or LinearOperator.
        qr_factor (numpy.ndarray): An (n, k) orthonormal basis for the
            approximate range of the matrix.

    Returns:
        numpy.ndarray: The approximate Cholesky factor F, such that A ~= F @ F.T.

    Notes:
        Based on Algorithm 5.5 of Halko et al. 2011.
    """
    small_matrix_1 = matrix @ qr_factor
    small_matrix_2 = qr_factor.T @ small_matrix_1
    factor, lower = cho_factor(small_matrix_2, overwrite_a=True)
    identity_operator = np.identity(factor.shape[0])
    inverse_factor = solve_triangular(
        factor, identity_operator, overwrite_b=True, lower=lower
    )
    return small_matrix_1 @ inverse_factor
