"""File algebra.py

:author: Michel Bierlaire
:date: Thu Jun 22 15:09:06 2023 

Function related to linear algebra
"""
import logging
import numpy as np
import scipy.linalg as la
from biogeme_optimization.exceptions import OptimizationError

logger = logging.getLogger(__name__)


def schnabel_eskow(
    A,
    tau=np.finfo(np.float64).eps ** 0.3333,
    taubar=np.finfo(np.float64).eps ** 0.6666,
    mu=0.1,
):
    """Modified Cholesky factorization by `Schnabel and Eskow (1999)`_.

    .. _`Schnabel and Eskow (1999)`: https://doi.org/10.1137/s105262349833266x

    If the matrix is 'safely' positive definite, the output is the
    classical Cholesky factor. If not, the diagonal elements are
    inflated in order to make it positive definite. The factor :math:`L`
    is such that :math:`A + E = PLL^TP^T`, where :math:`E` is a diagonal
    matrix contaninig the terms added to the diagonal, :math:`P` is a
    permutation matrix, and :math:`L` is w lower triangular matrix.

    :param A: matrix to factorize. Must be square and symmetric.
    :type A: numpy.array
    :param tau: tolerance factor.
                Default: :math:`\\varepsilon^{\\frac{1}{3}}`.
                See `Schnabel and Eskow (1999)`_
    :type tau: float
    :param taubar: tolerance factor.
                   Default: :math:`\\varepsilon^{\\frac{2}{3}}`.
                   See `Schnabel and Eskow (1999)`_
    :type taubar: float
    :param mu: tolerance factor.
               Default: 0.1.  See `Schnabel and Eskow (1999)`_
    :type mu: float

    :return: tuple :math:`L`, :math:`E`, :math:`P`,
                    where :math:`A + E = PLL^TP^T`.
    :rtype: numpy.array, numpy.array, numpy.array

    :raises biogeme.exceptions.OptimizationError: if the matrix A is not square.
    :raises biogeme.exceptions.OptimizationError: if the matrix A is not symmetric.
    """

    def pivot(j):
        A[j, j] = np.sqrt(A[j, j])
        for i in range(j + 1, dim):
            A[j, i] = A[i, j] = A[i, j] / A[j, j]
            A[i, j + 1 : i + 1] -= A[i, j] * A[j + 1 : i + 1, j]
            A[j + 1 : i + 1, i] = A[i, j + 1 : i + 1]

    def permute(i, j):
        A[[i, j]] = A[[j, i]]
        E[[i, j]] = E[[j, i]]
        A[:, [i, j]] = A[:, [j, i]]
        P[:, [i, j]] = P[:, [j, i]]

    A = A.astype(np.float64)
    dim = A.shape[0]
    if A.shape[1] != dim:
        raise OptimizationError('The matrix must be square')

    if not np.array_equal(A, A.T):
        raise OptimizationError('The matrix must be square and symmetric')

    E = np.zeros(dim, dtype=np.float64)
    P = np.identity(dim)
    phase_one = True
    gamma = abs(A.diagonal()).max()
    j = 0
    while j < dim and phase_one is True:
        a_max = np.max(A.diagonal()[j:])
        a_min = np.min(A.diagonal()[j:])
        if a_max < taubar * gamma or a_min < -mu * a_max:
            phase_one = False
            break

        # Pivot on maximum diagonal of remaining submatrix
        i = j + np.argmax(A.diagonal()[j:])
        if i != j:
            # Switch rows and columns of i and j of A
            permute(i, j)
        if (
            j < dim - 1
            and np.min(A.diagonal()[j + 1 :] - A[j + 1 :, j] ** 2 / A.diagonal()[j])
            < -mu * gamma
        ):
            phase_one = False  # go to phase two
        else:
            # perform jth iteration of factorization
            pivot(j)
            j += 1

    # Phase two, A not positive-definite
    if not phase_one:
        if j == dim - 1:
            E[-1] = delta = -A[-1, -1] + max(
                tau * (-A[-1, -1]) / (1 - tau), taubar * gamma
            )
            A[-1, -1] += delta
            A[-1, -1] = np.sqrt(A[-1, -1])
        else:
            delta_prev = 0.0
            gerschgorin = np.zeros(dim)
            k = j - 1  # k = number of iterations performed in phase one
            # Calculate lower Gerschgorin bounds of A[k+1]
            for i in range(k + 1, dim):
                gerschgorin[i] = (
                    A[i, i]
                    - np.sum(np.abs(A[i, k + 1 : i]))
                    - np.sum(np.abs(A[i + 1 : dim, i]))
                )
            # Modified Cholesky Decomposition
            for j in range(k + 1, dim - 2):
                # Pivot on maximum lower Gerschgorin bound estimate
                i = j + np.argmax(gerschgorin[j:])
                if i != j:
                    # Switch rows and columns of i and j of A
                    permute(i, j)
                # Calculate E[j, j] and add to diagonal
                norm_j = np.sum(np.abs(A[j + 1 : dim, j]))
                E[j] = delta = max(
                    0, -A[j, j] + max(norm_j, taubar * gamma), delta_prev
                )
                if delta > 0:
                    A[j, j] += delta
                    delta_prev = delta  # delta_prev will contain E_inf
                # Update Gerschgorin bound estimates
                if A[j, j] != norm_j:
                    temp = 1.0 - norm_j / A[j, j]
                    gerschgorin[j + 1 :] += np.abs(A[j + 1 :, j]) * temp
                # perform jth iteration of factorization
                pivot(j)

            # Final 2 by 2 submatrix
            eigenvalues = np.linalg.eigvalsh(A[-2:, -2:])
            eigenvalues.sort()
            E[-2] = E[-1] = delta = max(
                0,
                -eigenvalues[0]
                + max(
                    tau * (eigenvalues[1] - eigenvalues[0]) / (1 - tau), taubar * gamma
                ),
                delta_prev,
            )
            if delta > 0:
                A[-2, -2] += delta
                A[-1, -1] += delta
                delta_prev = delta
            A[-2, -2] = np.sqrt(A[-2, -2])  # overwrites A[-2, -2]
            A[-1, -2] = A[-1, -2] / A[-2, -2]  # overwrites A[-1, -2]
            A[-2, -1] = A[-1, -2]
            # overwrites A[-1, -1]
            A[-1, -1] = np.sqrt(A[-1, -1] - A[-1, -2] * A[-1, -2])

    L = np.tril(A)
    return L, np.diag(P @ E), P


def schnabel_eskow_direction(gradient, hessian, check_convexity=False):
    """Calculate a descent direction using the Schnabel-Eskow factorization"""
    if check_convexity:
        L, E, P = schnabel_eskow(hessian)
        corrections = E > 0
        if np.any(corrections):
            raise OptimizationError('The quadratic model is not convex.')
    else:
        L, _, P = schnabel_eskow(hessian)
    y3 = -P.T @ gradient
    y2 = la.solve_triangular(L, y3, lower=True, overwrite_b=True)
    y1 = la.solve_triangular(L.T, y2, lower=False, overwrite_b=True)
    d = P @ y1
    return d
