"""Cholesky decomposition functions."""

import numpy as np
from numpy import asarray_chkfinite, asarray, empty_like

# Local imports
from scipy._lib._util import _asarray_validated, _apply_over_batch, _deprecate_dtypes
from ._misc import LinAlgError, _datacopied
from .lapack import (
    get_lapack_funcs, _normalize_lapack_dtype, _ensure_aligned_and_native
)
from . import _batched_linalg

__all__ = ['cholesky', 'cho_factor', 'cho_solve', 'cholesky_banded',
           'cho_solve_banded']


def _check_format_errors_warnings(routine_name, err_lst):
    msg = (
        f"Internal {routine_name} return info = {[e['lapack_info'] for e in err_lst]} "
        f"for slices {[e['num'] for e in err_lst]}."
    )
    raise LinAlgError(msg)


def _cholesky(a, lower=False, overwrite_a=False, clean=True,
              check_finite=True):
    """Common code for cholesky() and cho_factor()."""

    # sanity checks
    a1 = _asarray_validated(a, check_finite=check_finite)
    a1 = np.atleast_2d(a1)
    if a1.shape[-1] != a1.shape[-2]:
        raise ValueError(f"Expected a square matrix or batch thereof, got {a1.shape=}")

    _deprecate_dtypes("linalg.cholesky", a1)

    a1, overwrite_a = _normalize_lapack_dtype(a1, overwrite_a)
    a1, overwrite_a = _ensure_aligned_and_native(a1, overwrite_a)
    overwrite_a = (overwrite_a and (a1.ndim == 2)
                   and (a1.flags["F_CONTIGUOUS"] or a1.flags["C_CONTIGUOUS"]))

    # accommodate empty arrays
    if a1.shape[-1] == 0:
        batch_shape = a1.shape[:-2]
        c = np.zeros(batch_shape + (0, 0), dtype=a1.dtype)
        return c

    # Heavy lifting
    c, err_lst = _batched_linalg._cholesky(a1, lower, overwrite_a, clean)

    if err_lst:
        _check_format_errors_warnings("potrf", err_lst)

    return c


def cholesky(a, lower=False, overwrite_a=False, check_finite=True):
    """
    Compute the Cholesky decomposition of a matrix.

    Returns the Cholesky decomposition, :math:`A = L L^*` or
    :math:`A = U^* U` of a Hermitian positive-definite matrix A.

    Parameters
    ----------
    a : (..., M, M) array_like
        Matrix to be decomposed
    lower : bool, optional
        Whether to compute the upper- or lower-triangular Cholesky
        factorization. During decomposition, only the selected half of the
        matrix is referenced. Default is upper-triangular.
    overwrite_a : bool, optional
        Whether to overwrite data in `a` (may improve performance). Default is False.
        See :ref:`tutorial_linalg_overwrite` for details.
    check_finite : bool, optional
        Whether to check that the entire input matrix contains only finite numbers.
        Disabling may give a performance gain, but may result in problems
        (crashes, non-termination) if the inputs do contain infinities or NaNs.

    Returns
    -------
    c : (..., M, M) ndarray
        Upper- or lower-triangular Cholesky factor of `a`.

    Raises
    ------
    LinAlgError : if decomposition fails.

    See Also
    --------
    cho_factor : Compute the Cholesky decomposition without explicitly setting
                    the other triangle to 0.

    Notes
    -----
    During the finiteness check (if selected), the entire matrix `a` is
    checked. During decomposition, `a` is assumed to be symmetric or Hermitian
    (as applicable), and only the half selected by option `lower` is referenced.
    Consequently, if `a` is asymmetric/non-Hermitian, `cholesky` may still
    succeed if the symmetric/Hermitian matrix represented by the selected half
    is positive definite, yet it may fail if an element in the other half is
    non-finite.

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.linalg import cholesky
    >>> a = np.array([[1,-2j],[2j,5]])
    >>> L = cholesky(a, lower=True)
    >>> L
    array([[ 1.+0.j,  0.+0.j],
           [ 0.+2.j,  1.+0.j]])
    >>> L @ L.T.conj()
    array([[ 1.+0.j,  0.-2.j],
           [ 0.+2.j,  5.+0.j]])

    """
    # `clean = True` represents setting other triangle to 0.
    c = _cholesky(a, lower=lower, overwrite_a=overwrite_a, clean=True,
                check_finite=check_finite)
    return c


def cho_factor(a, lower=False, overwrite_a=False, check_finite=True):
    """
    Compute the Cholesky decomposition of a matrix, to use in cho_solve.

    Returns a matrix containing the Cholesky decomposition,
    ``A = L L*`` or ``A = U* U`` of a Hermitian positive-definite matrix `a`.
    The return value can be directly used as the first parameter to cho_solve.

    .. warning::
        If `overwrite_a` is disabled, this function matches `cholesky` and
        zeros the other triangle. If it is enabled, the returned matrix may
        contain random data in the entries not used by the Cholesky
        decomposition, saving out on the explicit putting to 0 done by
        `cholesky`.

    Parameters
    ----------
    a : (..., M, M) array_like
        Matrix to be decomposed
    lower : bool, optional
        Whether to compute the upper or lower triangular Cholesky factorization.
        During decomposition, only the selected half of the matrix is referenced.
        (Default: upper-triangular)
    overwrite_a : bool, optional
        Whether to overwrite data in ``a`` (may improve performance). Default is False.
        See :ref:`tutorial_linalg_overwrite` for details.
    check_finite : bool, optional
        Whether to check that the entire input matrix contains only finite numbers.
        Disabling may give a performance gain, but may result in problems
        (crashes, non-termination) if the inputs do contain infinities or NaNs.

    Returns
    -------
    c : (..., M, M) ndarray
        Matrix whose upper or lower triangle contains the Cholesky factor
        of `a`. Other parts of the matrix contain random data.
    lower : ndarray, bool
        Flag indicating whether the factor is in the lower or upper triangle

    Raises
    ------
    LinAlgError
        Raised if decomposition fails.

    See Also
    --------
    cholesky : Compute the Cholesky decomposition and explicitly put the
                other triangle to 0.
    cho_solve : Solve a linear set equations using the Cholesky factorization
                of a matrix.

    Notes
    -----
    During the finiteness check (if selected), the entire matrix `a` is
    checked. During decomposition, `a` is assumed to be symmetric or Hermitian
    (as applicable), and only the half selected by option `lower` is referenced.
    Consequently, if `a` is asymmetric/non-Hermitian, `cholesky` may still
    succeed if the symmetric/Hermitian matrix represented by the selected half
    is positive definite, yet it may fail if an element in the other half is
    non-finite.

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.linalg import cho_factor
    >>> A = np.array([[9, 3, 1, 5], [3, 7, 5, 1], [1, 5, 9, 2], [5, 1, 2, 6]],
    ...             dtype=float, order="F")
    >>> A_ref = np.copy(A)
    >>> c, low = cho_factor(A)
    >>> c
    array([[3.        , 1.        , 0.33333333, 1.66666667],
           [0.        , 2.44948974, 1.90515869, -0.27216553],
           [0.        , 0.        , 2.29330749, 0.8559528 ],
           [0.        , 0.        , 0.        , 1.55418563]])
    >>> np.allclose(c.T @ c - A, np.zeros((4, 4)))
    True
    >>> c, low = cho_factor(A, overwrite_a=True)
    >>> c
    array([[3.        , 1.        , 0.33333333, 1.66666667],
           [3.        , 2.44948974, 1.90515869, -0.27216553],
           [1.        , 5.        , 2.29330749, 0.8559528 ],
           [5.        , 1.        , 2.        , 1.55418563]])
    >>> np.allclose(np.triu(c).T @ np.triu(c) - A_ref, np.zeros((4, 4)))
    True

    """
    # `clean=False` to represent that it is not necessary to set other triangle to 0.
    c = _cholesky(a, lower=lower, overwrite_a=overwrite_a, clean=False,
                    check_finite=check_finite)

    # broadcast `lower` argument for backwards compat
    batch_shape = a.shape[:-2]
    ret_lower = np.tile(lower, reps=batch_shape)

    return c, ret_lower


def cho_solve(c_and_lower, b, overwrite_b=False, check_finite=True):
    """Solve the linear equations A x = b, given the Cholesky factorization of A.

    The documentation is written assuming array arguments are of specified
    "core" shapes. However, array argument(s) of this function may have additional
    "batch" dimensions prepended to the core shape. In this case, the array is treated
    as a batch of lower-dimensional slices; see :ref:`linalg_batch` for details.

    Parameters
    ----------
    c_and_lower : tuple, (array, bool)
        Cholesky factorization of an array, as given by cho_factor.
        The first element of the
        tuple is the matrix containing the Cholesky factor, and the second element is a
        boolean flag indicating whether the factor is in the lower or upper triangle.
    b : array
        Right-hand side
    overwrite_b : bool, optional
        Whether to overwrite data in b (may improve performance)
        See :ref:`tutorial_linalg_overwrite` for details.
    check_finite : bool, optional
        Whether to check that the input matrices contain only finite numbers.
        Disabling may give a performance gain, but may result in problems
        (crashes, non-termination) if the inputs do contain infinities or NaNs.

    Returns
    -------
    x : array
        The solution to the system A x = b

    See Also
    --------
    cho_factor : Cholesky factorization of a matrix

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.linalg import cho_factor, cho_solve
    >>> A = np.array([[9, 3, 1, 5], [3, 7, 5, 1], [1, 5, 9, 2], [5, 1, 2, 6]])
    >>> c, low = cho_factor(A)
    >>> x = cho_solve((c, low), [1, 1, 1, 1])
    >>> np.allclose(A @ x - [1, 1, 1, 1], np.zeros(4))
    True

    """
    c, lower = c_and_lower
    return _cho_solve(c, b, lower, overwrite_b=overwrite_b, check_finite=check_finite)


@_apply_over_batch(('c', 2), ('b', '1|2'))
def _cho_solve(c, b, lower, overwrite_b, check_finite):
    if check_finite:
        b1 = asarray_chkfinite(b)
        c = asarray_chkfinite(c)
    else:
        b1 = asarray(b)
        c = asarray(c)

    if c.ndim != 2 or c.shape[0] != c.shape[1]:
        raise ValueError("The factored matrix c is not square.")
    if c.shape[1] != b1.shape[0]:
        raise ValueError(f"incompatible dimensions ({c.shape} and {b1.shape})")

    # accommodate empty arrays
    if b1.size == 0:
        dt = cho_solve((np.eye(2, dtype=b1.dtype), True),
                        np.ones(2, dtype=c.dtype)).dtype
        return empty_like(b1, dtype=dt)

    overwrite_b = overwrite_b or _datacopied(b1, b)

    potrs, = get_lapack_funcs(('potrs',), (c, b1))
    x, info = potrs(c, b1, lower=lower, overwrite_b=overwrite_b)
    if info != 0:
        raise ValueError(f'illegal value in {-info}th argument of internal potrs')
    return x


@_apply_over_batch(("ab", 2))
def cholesky_banded(ab, overwrite_ab=False, lower=False, check_finite=True):
    """
    Cholesky decompose a banded Hermitian positive-definite matrix.

    The matrix a is stored in ab either in lower-diagonal or upper-
    diagonal ordered form::

        ab[u + i - j, j] == a[i,j]        (if upper form; i <= j)
        ab[    i - j, j] == a[i,j]        (if lower form; i >= j)

    Example of ab (shape of a is (6,6), u=2)::

        upper form:
        *   *   a02 a13 a24 a35
        *   a01 a12 a23 a34 a45
        a00 a11 a22 a33 a44 a55

        lower form:
        a00 a11 a22 a33 a44 a55
        a10 a21 a32 a43 a54 *
        a20 a31 a42 a53 *   *

    Parameters
    ----------
    ab : (u + 1, M) array_like
        Banded matrix
    overwrite_ab : bool, optional
        Discard data in ab (may enhance performance)
        See :ref:`tutorial_linalg_overwrite` for details.
    lower : bool, optional
        Is the matrix in the lower form. (Default is upper form)
    check_finite : bool, optional
        Whether to check that the input matrix contains only finite numbers.
        Disabling may give a performance gain, but may result in problems
        (crashes, non-termination) if the inputs do contain infinities or NaNs.

    Returns
    -------
    c : (u + 1, M) ndarray
        Cholesky factorization of a, in the same banded format as ab

    See Also
    --------
    cho_solve_banded :
        Solve a linear set equations, given the Cholesky factorization
        of a banded Hermitian.

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.linalg import cholesky_banded
    >>> from numpy import allclose, zeros, diag
    >>> Ab = np.array([[0, 0, 1j, 2, 3j], [0, -1, -2, 3, 4], [9, 8, 7, 6, 9]])
    >>> A = np.diag(Ab[0,2:], k=2) + np.diag(Ab[1,1:], k=1)
    >>> A = A + A.conj().T + np.diag(Ab[2, :])
    >>> c = cholesky_banded(Ab)
    >>> C = np.diag(c[0, 2:], k=2) + np.diag(c[1, 1:], k=1) + np.diag(c[2, :])
    >>> np.allclose(C.conj().T @ C - A, np.zeros((5, 5)))
    True

    """
    if check_finite:
        ab = asarray_chkfinite(ab)
    else:
        ab = asarray(ab)

    # accommodate square empty matrices
    if ab.size == 0:
        dt = cholesky_banded(np.array([[0, 0], [1, 1]], dtype=ab.dtype)).dtype
        return empty_like(ab, dtype=dt)

    pbtrf, = get_lapack_funcs(('pbtrf',), (ab,))
    c, info = pbtrf(ab, lower=lower, overwrite_ab=overwrite_ab)
    if info > 0:
        raise LinAlgError(f"{info}-th leading minor not positive definite")
    if info < 0:
        raise ValueError(f'illegal value in {info}-th argument of internal pbtrf')
    return c


def cho_solve_banded(cb_and_lower, b, overwrite_b=False, check_finite=True):
    """
    Solve the linear equations ``A x = b``, given the Cholesky factorization of
    the banded Hermitian ``A``.

    The documentation is written assuming array arguments are of specified
    "core" shapes. However, array argument(s) of this function may have additional
    "batch" dimensions prepended to the core shape. In this case, the array is treated
    as a batch of lower-dimensional slices; see :ref:`linalg_batch` for details.

    Parameters
    ----------
    cb_and_lower : tuple, (ndarray, bool)
        The first element of the tuple is the matrix containing the Cholesky factor,
        as provided by `cholesky_banded`. The second element is a boolean flag
        indicating whether the factor is in the lower triangle, as provided *to*
        `cholesky_banded`.
    b : array_like
        Right-hand side
    overwrite_b : bool, optional
        If True, the function will overwrite the values in `b`.
        See :ref:`tutorial_linalg_overwrite` for details.
    check_finite : bool, optional
        Whether to check that the input matrices contain only finite numbers.
        Disabling may give a performance gain, but may result in problems
        (crashes, non-termination) if the inputs do contain infinities or NaNs.

    Returns
    -------
    x : array
        The solution to the system A x = b

    See Also
    --------
    cholesky_banded : Cholesky factorization of a banded matrix

    Notes
    -----

    .. versionadded:: 0.8.0

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.linalg import cholesky_banded, cho_solve_banded
    >>> Ab = np.array([[0, 0, 1j, 2, 3j], [0, -1, -2, 3, 4], [9, 8, 7, 6, 9]])
    >>> A = np.diag(Ab[0,2:], k=2) + np.diag(Ab[1,1:], k=1)
    >>> A = A + A.conj().T + np.diag(Ab[2, :])
    >>> c = cholesky_banded(Ab)
    >>> x = cho_solve_banded((c, False), np.ones(5))
    >>> np.allclose(A @ x - np.ones(5), np.zeros(5))
    True

    """
    (cb, lower) = cb_and_lower
    return _cho_solve_banded(cb, b, lower, overwrite_b=overwrite_b,
                             check_finite=check_finite)


@_apply_over_batch(('cb', 2), ('b', '1|2'))
def _cho_solve_banded(cb, b, lower, overwrite_b, check_finite):
    if check_finite:
        cb = asarray_chkfinite(cb)
        b = asarray_chkfinite(b)
    else:
        cb = asarray(cb)
        b = asarray(b)

    # Validate shapes.
    if cb.shape[-1] != b.shape[0]:
        raise ValueError("shapes of cb and b are not compatible.")

    # accommodate empty arrays
    if b.size == 0:
        m = cholesky_banded(np.array([[0, 0], [1, 1]], dtype=cb.dtype))
        dt = cho_solve_banded((m, True), np.ones(2, dtype=b.dtype)).dtype
        return empty_like(b, dtype=dt)

    pbtrs, = get_lapack_funcs(('pbtrs',), (cb, b))
    x, info = pbtrs(cb, b, lower=lower, overwrite_b=overwrite_b)
    if info > 0:
        raise LinAlgError(f"{info}th leading minor not positive definite")
    if info < 0:
        raise ValueError(f'illegal value in {-info}th argument of internal pbtrs')
    return x
