#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

'''
Extensions to the scipy.linalg module
'''

import numpy


# Numpy/scipy does not seem to have a convenient interface for
# pivoted Cholesky factorization. Newer versions of scipy (>=1.4) provide
# access to the raw lapack function, which is wrapped around here.
# With older versions of scipy, we use our own implementation instead.
try:
    from scipy.linalg.lapack import dpstrf as _dpstrf
except ImportError:
    def _pivoted_cholesky_wrapper(A, tol, lower):
        return pivoted_cholesky_python(A, tol=tol, lower=lower)
else:
    def _pivoted_cholesky_wrapper(A, tol, lower):
        N = A.shape[0]
        assert(A.shape == (N, N))
        L, piv, rank, info = _dpstrf(A, tol=tol, lower=lower)
        if info < 0:
            raise RuntimeError('Pivoted Cholesky factorization failed.')
        if lower:
            L[numpy.triu_indices(N, k=1)] = 0
            L[:, rank:] = 0
        else:
            L[numpy.tril_indices(N, k=-1)] = 0
            L[rank:, :] = 0
        return L, piv-1, rank


def pivoted_cholesky(A, tol=-1.0, lower=False):
    '''
    Performs a Cholesky factorization of A with full pivoting.
    A can be a (singular) positive semidefinite matrix.

    P.T * A * P = L * L.T   if   lower is True
    P.T * A * P = U.T * U   if   lower if False

    Use regular Cholesky factorization for positive definite matrices instead.

    Args:
        A : the matrix to be factorized
        tol : the stopping tolerance (see LAPACK documentation for dpstrf)
        lower : return lower triangular matrix L if true
                return upper triangular matrix U if false

    Returns:
        the factor L or U, the pivot vector (starting with 0), the rank
    '''
    return _pivoted_cholesky_wrapper(A, tol=tol, lower=lower)


def pivoted_cholesky_python(A, tol=-1.0, lower=False):
    '''
    Pedestrian implementation of Cholesky factorization with full column pivoting.
    The LAPACK version should be used instead whenever possible!

    Args:
        A : the positive semidefinite matrix to be factorized
        tol : stopping tolerance
        lower : return the lower or upper diagonal factorization

    Returns:
        the factor, the permutation vector, the rank
    '''
    N = A.shape[0]
    assert(A.shape == (N, N))

    D = numpy.diag(A).copy()
    if tol < 0:
        machine_epsilon = numpy.finfo(numpy.float).eps
        tol = N * machine_epsilon * numpy.amax(numpy.diag(A))

    L = numpy.zeros((N, N))
    piv = numpy.arange(N)
    rank = 0
    for k in range(N):
        s = k + numpy.argmax(D[k:])
        piv[k], piv[s] = piv[s], piv[k]
        D[k], D[s] = D[s], D[k]
        L[[k, s], :] = L[[s, k], :]
        if D[k] <= tol:
            break
        rank += 1
        L[k, k] = numpy.sqrt(D[k])
        L[k+1:, k] = (A[piv[k+1:], piv[k]] - numpy.dot(L[k+1:, :k], L[k, :k])) / L[k, k]
        D[k+1:] -= L[k+1:, k] ** 2

    if lower:
        return L, piv, rank
    else:
        return L.T, piv, rank
