// Copyright (C) 2019 Yixuan Qiu <yixuan.qiu@cos.name>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at https://mozilla.org/MPL/2.0/.

#ifndef BK_LDLT_H
#define BK_LDLT_H

#include <Eigen/Core>
#include <vector>
#include <stdexcept>

#include "../Util/CompInfo.h"

namespace Spectra {


// Bunch-Kaufman LDLT decomposition
// References:
// 1. Bunch, J. R., & Kaufman, L. (1977). Some stable methods for calculating inertia and solving symmetric linear systems.
//    Mathematics of computation, 31(137), 163-179.
// 2. Golub, G. H., & Van Loan, C. F. (2012). Matrix computations (Vol. 3). JHU press. Section 4.4.
// 3. Bunch-Parlett diagonal pivoting <http://oz.nthu.edu.tw/~d947207/Chap13_GE3.ppt>
template <typename Scalar = double>
class BKLDLT
{
private:
    typedef Eigen::Index Index;
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> Matrix;
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> Vector;
    typedef Eigen::Map<Vector> MapVec;
    typedef Eigen::Map<const Vector> MapConstVec;

    typedef Eigen::Matrix<Index, Eigen::Dynamic, 1> IntVector;
    typedef Eigen::Ref<Vector> GenericVector;
    typedef Eigen::Ref<Matrix> GenericMatrix;
    typedef const Eigen::Ref<const Matrix> ConstGenericMatrix;
    typedef const Eigen::Ref<const Vector> ConstGenericVector;

    Index m_n;
    Vector m_data;                  // storage for a lower-triangular matrix
    std::vector<Scalar*> m_colptr;  // pointers to columns
    IntVector m_perm;               // [-2, -1, 3, 1, 4, 5]: 0 <-> 2, 1 <-> 1, 2 <-> 3, 3 <-> 1, 4 <-> 4, 5 <-> 5
    std::vector< std::pair<Index, Index> > m_permc;  // compressed version of m_perm: [(0, 2), (2, 3), (3, 1)]

    bool m_computed;
    int  m_info;

    // Access to elements
    // Pointer to the k-th column
    Scalar* col_pointer(Index k) { return m_colptr[k]; }
    // A[i, j] -> m_colptr[j][i - j], i >= j
    Scalar& coeff(Index i, Index j) { return m_colptr[j][i - j]; }
    const Scalar& coeff(Index i, Index j) const { return m_colptr[j][i - j]; }
    // A[i, i] -> m_colptr[i][0]
    Scalar& diag_coeff(Index i) { return m_colptr[i][0]; }
    const Scalar& diag_coeff(Index i) const { return m_colptr[i][0]; }

    // Compute column pointers
    void compute_pointer()
    {
        m_colptr.clear();
        m_colptr.reserve(m_n);
        Scalar* head = m_data.data();

        for(Index i = 0; i < m_n; i++)
        {
            m_colptr.push_back(head);
            head += (m_n - i);
        }
    }

    // Copy mat - shift * I to m_data
    void copy_data(ConstGenericMatrix& mat, int uplo, const Scalar& shift)
    {
        if(uplo == Eigen::Lower)
        {
            for(Index j = 0; j < m_n; j++)
            {
                const Scalar* begin = &mat.coeffRef(j, j);
                const Index len = m_n - j;
                std::copy(begin, begin + len, col_pointer(j));
                diag_coeff(j) -= shift;
            }
        } else {
            Scalar* dest = m_data.data();
            for(Index i = 0; i < m_n; i++)
            {
                for(Index j = i; j < m_n; j++, dest++)
                {
                    *dest = mat.coeff(i, j);
                }
                diag_coeff(i) -= shift;
            }
        }
    }

    // Compute compressed permutations
    void compress_permutation()
    {
        for(Index i = 0; i < m_n; i++)
        {
            // Recover the permutation action
            const Index perm = (m_perm[i] >= 0) ? (m_perm[i]) : (-m_perm[i] - 1);
            if(perm != i)
                m_permc.push_back(std::make_pair(i, perm));
        }
    }

    // Working on the A[k:end, k:end] submatrix
    // Exchange k <-> r
    // Assume r >= k
    void pivoting_1x1(Index k, Index r)
    {
        // No permutation
        if(k == r)
        {
            m_perm[k] = r;
            return;
        }

        // A[k, k] <-> A[r, r]
        std::swap(diag_coeff(k), diag_coeff(r));

        // A[(r+1):end, k] <-> A[(r+1):end, r]
        std::swap_ranges(&coeff(r + 1, k), col_pointer(k + 1), &coeff(r + 1, r));

        // A[(k+1):(r-1), k] <-> A[r, (k+1):(r-1)]
        Scalar* src = &coeff(k + 1, k);
        for(Index j = k + 1; j < r; j++, src++)
        {
            std::swap(*src, coeff(r, j));
        }

        m_perm[k] = r;
    }

    // Working on the A[k:end, k:end] submatrix
    // Exchange [k+1, k] <-> [r, p]
    // Assume p >= k, r >= k+1
    void pivoting_2x2(Index k, Index r, Index p)
    {
        pivoting_1x1(k, p);
        pivoting_1x1(k + 1, r);

        // A[k+1, k] <-> A[r, k]
        std::swap(coeff(k + 1, k), coeff(r, k));

        // Use negative signs to indicate a 2x2 block
        // Also minus one to distinguish a negative zero from a positive zero
        m_perm[k] = -m_perm[k] - 1;
        m_perm[k + 1] = -m_perm[k + 1] - 1;
    }

    // A[r1, c1:c2] <-> A[r2, c1:c2]
    // Assume r2 >= r1 > c2 >= c1
    void interchange_rows(Index r1, Index r2, Index c1, Index c2)
    {
        if(r1 == r2)
            return;

        for(Index j = c1; j <= c2; j++)
        {
            std::swap(coeff(r1, j), coeff(r2, j));
        }
    }

    // lambda = |A[r, k]| = max{|A[k+1, k]|, ..., |A[end, k]|}
    // Assume k < end
    Scalar find_lambda(Index k, Index& r)
    {
        using std::abs;

        const Scalar* head = col_pointer(k);
        const Scalar* end = col_pointer(k + 1);
        Scalar lambda = abs(head[1]);
        r = k + 1;
        for(const Scalar* ptr = head + 2; ptr < end; ptr++)
        {
            const Scalar abs_elem = abs(*ptr);
            if(lambda < abs_elem)
            {
                lambda = abs_elem;
                r = k + (ptr - head);
            }
        }

        return lambda;
    }

    // sigma = |A[p, r]| = max {|A[k, r]|, ..., |A[end, r]|} \ {A[r, r]}
    // Assume k < r < end
    Scalar find_sigma(Index k, Index r, Index& p)
    {
        using std::abs;

        // First search A[r+1, r], ...,  A[end, r], which has the same task as find_lambda()
        // If r == end, we skip this search
        Scalar sigma = Scalar(-1);
        if(r < m_n - 1)
            sigma = find_lambda(r, p);

        // Then search A[k, r], ..., A[r-1, r], which maps to A[r, k], ..., A[r, r-1]
        for(Index j = k; j < r; j++)
        {
            const Scalar abs_elem = abs(coeff(r, j));
            if(sigma < abs_elem)
            {
                sigma = abs_elem;
                p = j;
            }
        }

        return sigma;
    }

    // Generate permutations and apply to A
    // Return true if the resulting pivoting is 1x1, and false if 2x2
    bool permutate_mat(Index k, const Scalar& alpha)
    {
        using std::abs;

        Index r = k, p = k;
        const Scalar lambda = find_lambda(k, r);

        // If lambda=0, no need to interchange
        if(lambda > Scalar(0))
        {
            const Scalar abs_akk = abs(diag_coeff(k));
            // If |A[k, k]| >= alpha * lambda, no need to interchange
            if(abs_akk < alpha * lambda)
            {
                const Scalar sigma = find_sigma(k, r, p);

                // If sigma * |A[k, k]| >= alpha * lambda^2, no need to interchange
                if(sigma * abs_akk < alpha * lambda * lambda)
                {
                    if(abs_akk >= alpha * sigma)
                    {
                        // Permutation on A
                        pivoting_1x1(k, r);

                        // Permutation on L
                        interchange_rows(k, r, 0, k - 1);
                        return true;
                    } else {
                        // Permutation on A
                        // [r, p] and [p, r] are symmetric, but we need to make sure
                        // p >= k and r >= k+1, so it is safe to always make r > p
                        // One exception is when min{r,p} == k+1, in which case we make
                        // r = k+1, so that only one permutation needs to be performed
                        const Index rp_min = std::min(r, p);
                        const Index rp_max = std::max(r, p);
                        if(rp_min == k + 1)
                        {
                            r = rp_min; p = rp_max;
                        } else {
                            r = rp_max; p = rp_min;
                        }
                        pivoting_2x2(k, r, p);
                        // Permutation on L
                        interchange_rows(k, p, 0, k - 1);
                        interchange_rows(k + 1, r, 0, k - 1);
                        return false;
                    }
                }
            }
        }

        return true;
    }

    // E = [e11, e12]
    //     [e21, e22]
    // Overwrite E with inv(E)
    void inverse_inplace_2x2(Scalar& e11, Scalar& e21, Scalar& e22) const
    {
        // inv(E) = [d11, d12], d11 = e22/delta, d21 = -e21/delta, d22 = e11/delta
        //          [d21, d22]
        const Scalar delta = e11 * e22 - e21 * e21;
        std::swap(e11, e22);
        e11 /= delta;
        e22 /= delta;
        e21 = -e21 / delta;
    }

    // Return value is the status, SUCCESSFUL/NUMERICAL_ISSUE
    int gaussian_elimination_1x1(Index k)
    {
        // D = 1 / A[k, k]
        const Scalar akk = diag_coeff(k);
        // Return NUMERICAL_ISSUE if not invertible
        if(akk == Scalar(0))
            return NUMERICAL_ISSUE;

        diag_coeff(k) = 1.0 / akk;

        // B -= l * l' / A[k, k], B := A[(k+1):end, (k+1):end], l := L[(k+1):end, k]
        Scalar* lptr = col_pointer(k) + 1;
        const Index ldim = m_n - k - 1;
        MapVec l(lptr, ldim);
        for(Index j = 0; j < ldim; j++)
        {
            MapVec(col_pointer(j + k + 1), ldim - j).noalias() -= (lptr[j] / akk) * l.tail(ldim - j);
        }

        // l /= A[k, k]
        l /= akk;

        return SUCCESSFUL;
    }

    // Return value is the status, SUCCESSFUL/NUMERICAL_ISSUE
    int gaussian_elimination_2x2(Index k)
    {
        // D = inv(E)
        Scalar& e11 = diag_coeff(k);
        Scalar& e21 = coeff(k + 1, k);
        Scalar& e22 = diag_coeff(k + 1);
        // Return NUMERICAL_ISSUE if not invertible
        if(e11 * e22 - e21 * e21 == Scalar(0))
            return NUMERICAL_ISSUE;

        inverse_inplace_2x2(e11, e21, e22);

        // X = l * inv(E), l := L[(k+2):end, k:(k+1)]
        Scalar* l1ptr = &coeff(k + 2, k);
        Scalar* l2ptr = &coeff(k + 2, k + 1);
        const Index ldim = m_n - k - 2;
        MapVec l1(l1ptr, ldim), l2(l2ptr, ldim);

        Eigen::Matrix<Scalar, Eigen::Dynamic, 2> X(ldim, 2);
        X.col(0).noalias() = l1 * e11 + l2 * e21;
        X.col(1).noalias() = l1 * e21 + l2 * e22;

        // B -= l * inv(E) * l' = X * l', B = A[(k+2):end, (k+2):end]
        for(Index j = 0; j < ldim; j++)
        {
            MapVec(col_pointer(j + k + 2), ldim - j).noalias() -= (X.col(0).tail(ldim - j) * l1ptr[j] + X.col(1).tail(ldim - j) * l2ptr[j]);
        }

        // l = X
        l1.noalias() = X.col(0);
        l2.noalias() = X.col(1);

        return SUCCESSFUL;
    }

public:
    BKLDLT() :
        m_n(0), m_computed(false), m_info(NOT_COMPUTED)
    {}

    // Factorize mat - shift * I
    BKLDLT(ConstGenericMatrix& mat, int uplo = Eigen::Lower, const Scalar& shift = Scalar(0)) :
        m_n(mat.rows()), m_computed(false), m_info(NOT_COMPUTED)
    {
        compute(mat, uplo, shift);
    }

    void compute(ConstGenericMatrix& mat, int uplo = Eigen::Lower, const Scalar& shift = Scalar(0))
    {
        using std::abs;

        m_n = mat.rows();
        if(m_n != mat.cols())
            throw std::invalid_argument("BKLDLT: matrix must be square");

        m_perm.setLinSpaced(m_n, 0, m_n - 1);
        m_permc.clear();

        // Copy data
        m_data.resize((m_n * (m_n + 1)) / 2);
        compute_pointer();
        copy_data(mat, uplo, shift);

        const Scalar alpha = (1.0 + std::sqrt(17.0)) / 8.0;
        Index k = 0;
        for(k = 0; k < m_n - 1; k++)
        {
            // 1. Interchange rows and columns of A, and save the result to m_perm
            bool is_1x1 = permutate_mat(k, alpha);

            // 2. Gaussian elimination
            if(is_1x1)
            {
                m_info = gaussian_elimination_1x1(k);
            } else {
                m_info = gaussian_elimination_2x2(k);
                k++;
            }

            // 3. Check status
            if(m_info != SUCCESSFUL)
                break;
        }
        // Invert the last 1x1 block if it exists
        if(k == m_n - 1)
        {
            const Scalar akk = diag_coeff(k);
            if(akk == Scalar(0))
                m_info = NUMERICAL_ISSUE;

            diag_coeff(k) = Scalar(1) / diag_coeff(k);
        }

        compress_permutation();

        m_computed = true;
    }

    // Solve Ax=b
    void solve_inplace(GenericVector b) const
    {
        if(!m_computed)
            throw std::logic_error("BKLDLT: need to call compute() first");

        // PAP' = LDL'
        // 1. b -> Pb
        Scalar* x = b.data();
        MapVec res(x, m_n);
        Index npermc = m_permc.size();
        for(Index i = 0; i < npermc; i++)
        {
            std::swap(x[m_permc[i].first], x[m_permc[i].second]);
        }

        // 2. Lz = Pb
        // If m_perm[end] < 0, then end with m_n - 3, otherwise end with m_n - 2
        const Index end = (m_perm[m_n - 1] < 0) ? (m_n - 3) : (m_n - 2);
        for(Index i = 0; i <= end; i++)
        {
            const Index b1size = m_n - i - 1;
            const Index b2size = b1size - 1;
            if(m_perm[i] >= 0)
            {
                MapConstVec l(&coeff(i + 1, i), b1size);
                res.segment(i + 1, b1size).noalias() -= l * x[i];
            } else {
                MapConstVec l1(&coeff(i + 2, i), b2size);
                MapConstVec l2(&coeff(i + 2, i + 1), b2size);
                res.segment(i + 2, b2size).noalias() -= (l1 * x[i] + l2 * x[i + 1]);
                i++;
            }
        }

        // 3. Dw = z
        for(Index i = 0; i < m_n; i++)
        {
            const Scalar e11 = diag_coeff(i);
            if(m_perm[i] >= 0)
            {
                x[i] *= e11;
            } else {
                const Scalar e21 = coeff(i + 1, i), e22 = diag_coeff(i + 1);
                const Scalar wi = x[i] * e11 + x[i + 1] * e21;
                x[i + 1] = x[i] * e21 + x[i + 1] * e22;
                x[i] = wi;
                i++;
            }
        }

        // 4. L'y = w
        // If m_perm[end] < 0, then start with m_n - 3, otherwise start with m_n - 2
        Index i = (m_perm[m_n - 1] < 0) ? (m_n - 3) : (m_n - 2);
        for(; i >= 0; i--)
        {
            const Index ldim = m_n - i - 1;
            MapConstVec l(&coeff(i + 1, i), ldim);
            x[i] -= res.segment(i + 1, ldim).dot(l);

            if(m_perm[i] < 0)
            {
                MapConstVec l2(&coeff(i + 1, i - 1), ldim);
                x[i - 1] -= res.segment(i + 1, ldim).dot(l2);
                i--;
            }
        }

        // 5. x = P'y
        for(Index i = npermc - 1; i >= 0; i--)
        {
            std::swap(x[m_permc[i].first], x[m_permc[i].second]);
        }
    }

    Vector solve(ConstGenericVector& b) const
    {
        Vector res = b;
        solve_inplace(res);
        return res;
    }

    int info() const { return m_info; }
};


} // namespace Spectra

#endif // BK_LDLT_H