// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

/**
 * @file simplicial_cholesky.hpp
 * @brief Simplicial Cholesky factorization for sparse matrices
 *
 * SimplicialLLT<T>  : A = L * L^T  (standard Cholesky)
 * SimplicialLDLT<T> : A = L * D * L^T  (LDL^T, no sqrt needed)
 *
 * Efficient O(flops(L)) implementation using left-looking column Cholesky + row linked lists.
 * Uses only the lower triangle of A (the upper triangle is ignored).
 *
 * Reference: Tim Davis, "Direct Methods for Sparse Linear Systems", SIAM 2006
 */

#ifndef SANGI_SIMPLICIAL_CHOLESKY_HPP
#define SANGI_SIMPLICIAL_CHOLESKY_HPP

#include <vector>
#include <cmath>
#include <stdexcept>
#include <algorithm>
#include <limits>
#include <string>
#include <memory>
#include <functional>
#include <utility>
#include <cassert>

#include "../concepts/algebraic_concepts.hpp"
#include "../core/sparse_matrix.hpp"
#include "../core/vector.hpp"
#include "sparse_ordering.hpp"

namespace sangi {

namespace detail {
namespace schol {

/// Extracts the lower triangular part of A into CSC arrays
template<typename T, typename IndexType>
void extract_lower_csc(
    const SparseMatrix<T, IndexType>& A,
    IndexType n,
    std::vector<IndexType>& col_ptr,
    std::vector<IndexType>& row_ind,
    std::vector<T>& values)
{
    // Collect the lower triangular entries per column
    std::vector<std::vector<std::pair<IndexType, T>>> cols(n);

    for (IndexType j = 0; j < n; ++j) {
        for (IndexType i = j; i < n; ++i) {
            T v = A.coeff(i, j);
            if (v != T{0}) {
                cols[j].emplace_back(i, v);
            }
        }
    }

    // Build the CSC pointers
    col_ptr.resize(n + 1);
    col_ptr[0] = IndexType{0};
    for (IndexType j = 0; j < n; ++j) {
        col_ptr[j + 1] = col_ptr[j] + static_cast<IndexType>(cols[j].size());
    }

    auto nnz = col_ptr[n];
    row_ind.resize(nnz);
    values.resize(nnz);

    for (IndexType j = 0; j < n; ++j) {
        auto base = col_ptr[j];
        for (std::size_t q = 0; q < cols[j].size(); ++q) {
            row_ind[base + q] = cols[j][q].first;
            values[base + q] = cols[j][q].second;
        }
    }
}

} // namespace schol
} // namespace detail

// ============================================================================
// SimplicialLLT — A = L * L^T
// ============================================================================

/**
 * @brief Simplicial LLT factorization for sparse matrices
 *
 * Factorizes a symmetric positive definite sparse matrix A into a lower triangular matrix L (A = L L^T).
 * Efficiently tracks the columns to update using left-looking column Cholesky + row linked lists.
 *
 * Usage example:
 * @code
 *   SimplicialLLT<double> chol(A);
 *   Vector<double> x = chol.solve(b);  // solve Ax = b
 *   auto L = chol.matrixL();           // obtain L
 * @endcode
 */
template<typename T, typename IndexType = std::size_t>
    requires concepts::OrderedField<T> && std::integral<IndexType>
class SimplicialLLT {
public:
    SimplicialLLT() = default;

    /// Constructs by factorizing A (default is Natural: preserves the meaning of matrixL().
    /// Specify SparseOrdering::AMD for fill reduction)
    explicit SimplicialLLT(const SparseMatrix<T, IndexType>& A,
                           SparseOrdering ordering = SparseOrdering::Natural) {
        compute(A, ordering);
    }

    /// Computes A = L * L^T
    void compute(const SparseMatrix<T, IndexType>& A,
                 SparseOrdering ordering = SparseOrdering::Natural) {
        const auto n = static_cast<IndexType>(A.rows());
        if (static_cast<IndexType>(A.cols()) != n) {
            assert(false && "DimensionError: SimplicialLLT: matrix must be square");
            throw DimensionError("SimplicialLLT: matrix must be square");
        }
        n_ = n;
        ordering_ = ordering;
        factored_ = false;

        if (n == 0) {
            fill_perm_.clear();
            lp_.assign(1, IndexType{0});
            factored_ = true;
            return;
        }

        const bool reordered = (ordering_ == SparseOrdering::AMD);
        fill_perm_ = reordered ? algorithms::amd_ordering(A)
                               : algorithms::natural_ordering(static_cast<std::size_t>(n));

        std::vector<IndexType> a_cp, a_ri;
        std::vector<T> a_val;
        if (reordered) {
            auto B = algorithms::permute_symmetric(A, fill_perm_);
            detail::schol::extract_lower_csc(B, n, a_cp, a_ri, a_val);
        } else {
            detail::schol::extract_lower_csc(A, n, a_cp, a_ri, a_val);
        }

        factorize(n, a_cp, a_ri, a_val);
        factored_ = true;
    }

    /// Solves Ax = b (L y = b → L^T x = y)
    Vector<T> solve(const Vector<T>& b) const {
        if (!factored_) throw MathError("SimplicialLLT: not factored");
        if (b.size() != n_) {
            assert(false && "DimensionError: SimplicialLLT::solve: dimension mismatch");
            throw DimensionError("SimplicialLLT::solve: dimension mismatch");
        }
        if (ordering_ == SparseOrdering::Natural)
            return backward_solve(forward_solve(b));
        // When reordered: c = P b → solve z → x[fill_perm_[i]] = z[i]
        Vector<T> c(n_);
        for (std::size_t i = 0; i < static_cast<std::size_t>(n_); ++i)
            c[i] = b[fill_perm_[i]];
        Vector<T> z = backward_solve(forward_solve(c));
        Vector<T> x(n_);
        for (std::size_t i = 0; i < static_cast<std::size_t>(n_); ++i)
            x[fill_perm_[i]] = z[i];
        return x;
    }

    /// fill-reducing permutation (nontrivial only when AMD is specified; matrixL() is L in this permutation)
    const std::vector<std::size_t>& permutation() const { return fill_perm_; }

    /// Returns the lower triangular matrix L as a SparseMatrix
    SparseMatrix<T, IndexType> matrixL() const {
        SparseMatrix<T, IndexType> L(n_, n_, SparseStorageFormat::COO);
        for (IndexType j = 0; j < n_; ++j) {
            for (auto p = lp_[j]; p < lp_[j + 1]; ++p) {
                L.set_coeff(li_[p], j, lx_[p]);
            }
        }
        return L;
    }

    bool success() const { return factored_; }
    IndexType size() const { return n_; }

private:
    IndexType n_ = 0;
    std::vector<T> lx_;
    std::vector<IndexType> li_;
    std::vector<IndexType> lp_;
    SparseOrdering ordering_ = SparseOrdering::Natural;
    std::vector<std::size_t> fill_perm_;
    bool factored_ = false;

    /**
     * @brief Left-looking column Cholesky
     *
     * Efficiently tracks the columns to update via row linked lists:
     *   head[i] = head of the list of columns k whose next entry is in row i
     *   nxt[k]  = next column in the same list
     *   cpos[k] = current position of column k (index within L(:,k))
     *
     * When processing column j, the list at head[j] gives all
     * columns k with L(j, k) != 0.
     */
    void factorize(
        IndexType n,
        const std::vector<IndexType>& a_cp,
        const std::vector<IndexType>& a_ri,
        const std::vector<T>& a_val)
    {
        const auto NONE = static_cast<IndexType>(-1);

        // Column data of L (built dynamically)
        std::vector<std::vector<IndexType>> cr(n);  // row indices of column j
        std::vector<std::vector<T>> cv(n);           // values of column j

        // Row linked lists
        std::vector<IndexType> head(n, NONE);
        std::vector<IndexType> nxt(n, NONE);
        std::vector<std::size_t> cpos(n, 0);

        // Dense work & change tracking
        std::vector<T> x(n, T{0});
        std::vector<IndexType> touched;
        touched.reserve(n);

        for (IndexType j = 0; j < n; ++j) {
            // Scatter A(:,j)
            for (auto p = a_cp[j]; p < a_cp[j + 1]; ++p) {
                x[a_ri[p]] = a_val[p];
                touched.push_back(a_ri[p]);
            }

            // cmod: traverse the linked list at head[j]
            IndexType k = head[j];
            head[j] = NONE;
            while (k != NONE) {
                IndexType nk = nxt[k];
                nxt[k] = NONE;

                T ljk = cv[k][cpos[k]];

                // x(i) -= L(j,k) * L(i,k) for i >= j
                for (auto q = cpos[k]; q < cr[k].size(); ++q) {
                    IndexType i = cr[k][q];
                    if (x[i] == T{0} && ljk * cv[k][q] != T{0}) {
                        touched.push_back(i);
                    }
                    x[i] -= ljk * cv[k][q];
                }

                // Advance to the next entry and relink
                cpos[k]++;
                if (cpos[k] < cr[k].size()) {
                    IndexType nr = cr[k][cpos[k]];
                    nxt[k] = head[nr];
                    head[nr] = k;
                }

                k = nk;
            }

            // Diagonal check
            if (x[j] <= T{0}) {
                throw MathError(
                    "SimplicialLLT: matrix is not positive definite "
                    "(non-positive diagonal at column " + std::to_string(j) + ")");
            }
            T ljj = std::sqrt(x[j]);

            // Store L(:,j) (diagonal + sorted off-diagonal)
            cr[j].push_back(j);
            cv[j].push_back(ljj);

            std::sort(touched.begin(), touched.end());
            touched.erase(std::unique(touched.begin(), touched.end()), touched.end());

            for (auto i : touched) {
                if (i > j && x[i] != T{0}) {
                    cr[j].push_back(i);
                    cv[j].push_back(x[i] / ljj);
                }
            }

            // Link column j (to its first off-diagonal row)
            if (cr[j].size() > 1) {
                cpos[j] = 1;
                IndexType nr = cr[j][1];
                nxt[j] = head[nr];
                head[nr] = j;
            }

            // Clear the workspace
            for (auto i : touched) x[i] = T{0};
            touched.clear();
        }

        pack_csc(n, cr, cv);
    }

    void pack_csc(
        IndexType n,
        const std::vector<std::vector<IndexType>>& cr,
        const std::vector<std::vector<T>>& cv)
    {
        lp_.resize(n + 1);
        lp_[0] = 0;
        for (IndexType j = 0; j < n; ++j) {
            lp_[j + 1] = lp_[j] + static_cast<IndexType>(cr[j].size());
        }
        auto total = lp_[n];
        li_.resize(total);
        lx_.resize(total);
        for (IndexType j = 0; j < n; ++j) {
            auto base = lp_[j];
            for (std::size_t q = 0; q < cr[j].size(); ++q) {
                li_[base + q] = cr[j][q];
                lx_[base + q] = cv[j][q];
            }
        }
    }

    /// Forward substitution: L y = b
    Vector<T> forward_solve(const Vector<T>& b) const {
        Vector<T> y = b;
        for (IndexType j = 0; j < n_; ++j) {
            auto p0 = lp_[j];
            y[j] /= lx_[p0];
            for (auto p = p0 + 1; p < lp_[j + 1]; ++p) {
                y[li_[p]] -= lx_[p] * y[j];
            }
        }
        return y;
    }

    /// Back substitution: L^T x = y
    Vector<T> backward_solve(const Vector<T>& y) const {
        Vector<T> x = y;
        for (IndexType jj = 0; jj < n_; ++jj) {
            IndexType j = n_ - 1 - jj;
            for (auto p = lp_[j] + 1; p < lp_[j + 1]; ++p) {
                x[j] -= lx_[p] * x[li_[p]];
            }
            x[j] /= lx_[lp_[j]];
        }
        return x;
    }
};

// ============================================================================
// SimplicialLDLT — A = L * D * L^T
// ============================================================================

/**
 * @brief Simplicial LDLT factorization for sparse matrices
 *
 * L is unit lower triangular (diagonal=1) and D is a diagonal matrix. No sqrt needed and more stable than LLT.
 *
 * Usage example:
 * @code
 *   SimplicialLDLT<double> ldlt(A);
 *   Vector<double> x = ldlt.solve(b);
 *   auto D = ldlt.vectorD();
 * @endcode
 */
template<typename T, typename IndexType = std::size_t>
    requires concepts::OrderedField<T> && std::integral<IndexType>
class SimplicialLDLT {
public:
    SimplicialLDLT() = default;

    explicit SimplicialLDLT(const SparseMatrix<T, IndexType>& A,
                            SparseOrdering ordering = SparseOrdering::Natural) {
        compute(A, ordering);
    }

    void compute(const SparseMatrix<T, IndexType>& A,
                 SparseOrdering ordering = SparseOrdering::Natural) {
        const auto n = static_cast<IndexType>(A.rows());
        if (static_cast<IndexType>(A.cols()) != n) {
            assert(false && "DimensionError: SimplicialLDLT: matrix must be square");
            throw DimensionError("SimplicialLDLT: matrix must be square");
        }
        n_ = n;
        ordering_ = ordering;
        factored_ = false;

        if (n == 0) {
            fill_perm_.clear();
            lp_.assign(1, IndexType{0});
            factored_ = true;
            return;
        }

        const bool reordered = (ordering_ == SparseOrdering::AMD);
        fill_perm_ = reordered ? algorithms::amd_ordering(A)
                               : algorithms::natural_ordering(static_cast<std::size_t>(n));

        std::vector<IndexType> a_cp, a_ri;
        std::vector<T> a_val;
        if (reordered) {
            auto B = algorithms::permute_symmetric(A, fill_perm_);
            detail::schol::extract_lower_csc(B, n, a_cp, a_ri, a_val);
        } else {
            detail::schol::extract_lower_csc(A, n, a_cp, a_ri, a_val);
        }

        factorize(n, a_cp, a_ri, a_val);
        factored_ = true;
    }

    /// Solves Ax = b (Ly = b → Dz = y → L^T x = z)
    Vector<T> solve(const Vector<T>& b) const {
        if (!factored_) throw MathError("SimplicialLDLT: not factored");
        if (b.size() != n_) {
            assert(false && "DimensionError: SimplicialLDLT::solve: dimension mismatch");
            throw DimensionError("SimplicialLDLT::solve: dimension mismatch");
        }

        const bool reordered = (ordering_ != SparseOrdering::Natural);
        Vector<T> c(n_);
        if (reordered)
            for (std::size_t i = 0; i < static_cast<std::size_t>(n_); ++i)
                c[i] = b[fill_perm_[i]];
        else
            c = b;

        Vector<T> y = forward_solve(c);
        for (IndexType i = 0; i < n_; ++i) {
            if (std::abs(d_[i]) < std::numeric_limits<T>::epsilon()) {
                throw MathError(
                    "SimplicialLDLT::solve: zero diagonal in D at index "
                    + std::to_string(i));
            }
            y[i] /= d_[i];
        }
        Vector<T> z = backward_solve(y);
        if (!reordered) return z;
        Vector<T> x(n_);
        for (std::size_t i = 0; i < static_cast<std::size_t>(n_); ++i)
            x[fill_perm_[i]] = z[i];
        return x;
    }

    /// fill-reducing permutation (nontrivial only when AMD is specified)
    const std::vector<std::size_t>& permutation() const { return fill_perm_; }

    /// Returns the unit lower triangular matrix L (diagonal = 1)
    SparseMatrix<T, IndexType> matrixL() const {
        SparseMatrix<T, IndexType> L(n_, n_, SparseStorageFormat::COO);
        for (IndexType j = 0; j < n_; ++j) {
            L.set_coeff(j, j, T{1});
            for (auto p = lp_[j]; p < lp_[j + 1]; ++p) {
                L.set_coeff(li_[p], j, lx_[p]);
            }
        }
        return L;
    }

    /// Vector of the diagonal matrix D
    const std::vector<T>& vectorD() const { return d_; }

    bool success() const { return factored_; }
    IndexType size() const { return n_; }

private:
    IndexType n_ = 0;
    std::vector<T> lx_;         // values of L (no diagonal, CSC)
    std::vector<IndexType> li_;  // row indices of L (CSC)
    std::vector<IndexType> lp_;  // column pointers of L (CSC, size n+1)
    std::vector<T> d_;           // diagonal of D
    SparseOrdering ordering_ = SparseOrdering::Natural;
    std::vector<std::size_t> fill_perm_;
    bool factored_ = false;

    /**
     * @brief LDLT left-looking column factorization
     *
     * cmod: x(i) -= L(j,k) * D(k) * L(i,k)
     * D(j) = x(j), L(i,j) = x(i) / D(j)
     *
     * L has no diagonal (unit diagonal). The linked-list start position is
     * col_rows[k][0] (the first off-diagonal entry).
     */
    void factorize(
        IndexType n,
        const std::vector<IndexType>& a_cp,
        const std::vector<IndexType>& a_ri,
        const std::vector<T>& a_val)
    {
        const auto NONE = static_cast<IndexType>(-1);

        std::vector<std::vector<IndexType>> cr(n);
        std::vector<std::vector<T>> cv(n);
        d_.resize(n, T{0});

        std::vector<IndexType> head(n, NONE);
        std::vector<IndexType> nxt(n, NONE);
        std::vector<std::size_t> cpos(n, 0);

        std::vector<T> x(n, T{0});
        std::vector<IndexType> touched;
        touched.reserve(n);

        for (IndexType j = 0; j < n; ++j) {
            // Scatter
            for (auto p = a_cp[j]; p < a_cp[j + 1]; ++p) {
                x[a_ri[p]] = a_val[p];
                touched.push_back(a_ri[p]);
            }

            // cmod
            IndexType k = head[j];
            head[j] = NONE;
            while (k != NONE) {
                IndexType nk = nxt[k];
                nxt[k] = NONE;

                T ljk = cv[k][cpos[k]];
                T ljk_dk = ljk * d_[k];

                // x(j) -= L(j,k)^2 * D(k)  [entry at cpos[k], row = j]
                if (x[j] == T{0} && ljk_dk * ljk != T{0}) {
                    touched.push_back(j);
                }
                x[j] -= ljk * ljk_dk;

                // x(i) -= L(i,k) * L(j,k) * D(k) for i > j
                for (auto q = cpos[k] + 1; q < cr[k].size(); ++q) {
                    IndexType i = cr[k][q];
                    if (x[i] == T{0} && ljk_dk * cv[k][q] != T{0}) {
                        touched.push_back(i);
                    }
                    x[i] -= cv[k][q] * ljk_dk;
                }

                cpos[k]++;
                if (cpos[k] < cr[k].size()) {
                    IndexType nr = cr[k][cpos[k]];
                    nxt[k] = head[nr];
                    head[nr] = k;
                }

                k = nk;
            }

            // Diagonal D(j)
            d_[j] = x[j];
            if (std::abs(d_[j]) < std::numeric_limits<T>::epsilon()) {
                throw MathError(
                    "SimplicialLDLT: zero or near-zero pivot at column "
                    + std::to_string(j));
            }

            // Store the off-diagonal of L(:,j)
            std::sort(touched.begin(), touched.end());
            touched.erase(std::unique(touched.begin(), touched.end()), touched.end());

            for (auto i : touched) {
                if (i > j && x[i] != T{0}) {
                    cr[j].push_back(i);
                    cv[j].push_back(x[i] / d_[j]);
                }
            }

            // Link column j
            if (!cr[j].empty()) {
                cpos[j] = 0;
                IndexType nr = cr[j][0];
                nxt[j] = head[nr];
                head[nr] = j;
            }

            // Clear
            for (auto i : touched) x[i] = T{0};
            touched.clear();
        }

        pack_csc(n, cr, cv);
    }

    void pack_csc(
        IndexType n,
        const std::vector<std::vector<IndexType>>& cr,
        const std::vector<std::vector<T>>& cv)
    {
        lp_.resize(n + 1);
        lp_[0] = 0;
        for (IndexType j = 0; j < n; ++j) {
            lp_[j + 1] = lp_[j] + static_cast<IndexType>(cr[j].size());
        }
        auto total = lp_[n];
        li_.resize(total);
        lx_.resize(total);
        for (IndexType j = 0; j < n; ++j) {
            auto base = lp_[j];
            for (std::size_t q = 0; q < cr[j].size(); ++q) {
                li_[base + q] = cr[j][q];
                lx_[base + q] = cv[j][q];
            }
        }
    }

    /// Forward substitution: L y = b (L is unit lower triangular)
    Vector<T> forward_solve(const Vector<T>& b) const {
        Vector<T> y = b;
        for (IndexType j = 0; j < n_; ++j) {
            for (auto p = lp_[j]; p < lp_[j + 1]; ++p) {
                y[li_[p]] -= lx_[p] * y[j];
            }
        }
        return y;
    }

    /// Back substitution: L^T x = y (L is unit lower triangular)
    Vector<T> backward_solve(const Vector<T>& y) const {
        Vector<T> x = y;
        for (IndexType jj = 0; jj < n_; ++jj) {
            IndexType j = n_ - 1 - jj;
            for (auto p = lp_[j]; p < lp_[j + 1]; ++p) {
                x[j] -= lx_[p] * x[li_[p]];
            }
        }
        return x;
    }
};

// ============================================================================
// Convenience functions
// ============================================================================

/// Cholesky factorization of a sparse matrix (returns L)
template<typename T, typename IndexType>
    requires concepts::OrderedField<T> && std::integral<IndexType>
SparseMatrix<T, IndexType> simplicial_cholesky(
    const SparseMatrix<T, IndexType>& A)
{
    SimplicialLLT<T, IndexType> chol(A);
    return chol.matrixL();
}

/// Solves Ax = b via Cholesky of a sparse matrix (applies fill-reducing AMD by default since only the solution is returned)
template<typename T, typename IndexType>
    requires concepts::OrderedField<T> && std::integral<IndexType>
Vector<T> simplicial_cholesky_solve(
    const SparseMatrix<T, IndexType>& A,
    const Vector<T>& b,
    SparseOrdering ordering = SparseOrdering::AMD)
{
    SimplicialLLT<T, IndexType> chol(A, ordering);
    return chol.solve(b);
}

/// Returns a SimplicialLLT as a preconditioner
template<typename T, typename IndexType>
    requires concepts::OrderedField<T> && std::integral<IndexType>
std::function<Vector<T>(const Vector<T>&)>
simplicial_cholesky_preconditioner(
    const SparseMatrix<T, IndexType>& A)
{
    auto chol = std::make_shared<SimplicialLLT<T, IndexType>>(A);
    return [chol](const Vector<T>& r) -> Vector<T> {
        return chol->solve(r);
    };
}

} // namespace sangi

#endif // SANGI_SIMPLICIAL_CHOLESKY_HPP