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

/**
 * @file simplicial_cholesky.hpp
 * @brief 疎行列用 Simplicial Cholesky 分解
 *
 * SimplicialLLT<T>  : A = L * L^T  (標準 Cholesky)
 * SimplicialLDLT<T> : A = L * D * L^T  (LDL^T, sqrt 不要)
 *
 * 左向き列 Cholesky + 行リンクリストで O(flops(L)) の効率的な実装。
 * A の下三角のみ使用 (上三角は無視)。
 *
 * 参考: Tim Davis, "Direct Methods for Sparse Linear Systems", SIAM 2006
 */

#ifndef CALX_SIMPLICIAL_CHOLESKY_HPP
#define CALX_SIMPLICIAL_CHOLESKY_HPP

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

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

namespace calx {

namespace detail {
namespace schol {

/// A の下三角部分を CSC 配列に抽出する
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)
{
    // 列ごとの下三角エントリを収集
    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);
            }
        }
    }

    // CSC ポインタ構築
    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 分解
 *
 * 対称正定値疎行列 A を下三角行列 L に分解する (A = L L^T)。
 * 左向き列 Cholesky + 行リンクリストで効率的に更新列を追跡。
 *
 * 使用例:
 * @code
 *   SimplicialLLT<double> chol(A);
 *   Vector<double> x = chol.solve(b);  // Ax = b を解く
 *   auto L = chol.matrixL();           // L を取得
 * @endcode
 */
template<typename T, typename IndexType = std::size_t>
    requires concepts::OrderedField<T> && std::integral<IndexType>
class SimplicialLLT {
public:
    SimplicialLLT() = default;

    /// A を分解して構築
    explicit SimplicialLLT(const SparseMatrix<T, IndexType>& A) {
        compute(A);
    }

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

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

        std::vector<IndexType> a_cp, a_ri;
        std::vector<T> a_val;
        detail::schol::extract_lower_csc(A, n, a_cp, a_ri, a_val);

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

    /// 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_) throw DimensionError("SimplicialLLT::solve: dimension mismatch");
        return backward_solve(forward_solve(b));
    }

    /// 下三角行列 L を 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_;
    bool factored_ = false;

    /**
     * @brief 左向き列 Cholesky
     *
     * 行リンクリストで更新列を効率追跡:
     *   head[i] = 行 i に次のエントリを持つ列 k のリスト先頭
     *   nxt[k]  = 同リスト内の次の列
     *   cpos[k] = 列 k の現在位置 (L(:,k) 内のインデックス)
     *
     * 列 j を処理するとき head[j] のリストが
     * L(j, k) != 0 となる全ての列 k を与える。
     */
    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);

        // L の列データ (動的構築)
        std::vector<std::vector<IndexType>> cr(n);  // 列 j の行インデックス
        std::vector<std::vector<T>> cv(n);           // 列 j の値

        // 行リンクリスト
        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 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: 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];
                }

                // 次のエントリへ進めてリンク
                cpos[k]++;
                if (cpos[k] < cr[k].size()) {
                    IndexType nr = cr[k][cpos[k]];
                    nxt[k] = head[nr];
                    head[nr] = k;
                }

                k = nk;
            }

            // 対角チェック
            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]);

            // L(:,j) を格納 (対角 + ソート済み非対角)
            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);
                }
            }

            // 列 j をリンク (最初の非対角行へ)
            if (cr[j].size() > 1) {
                cpos[j] = 1;
                IndexType nr = cr[j][1];
                nxt[j] = head[nr];
                head[nr] = j;
            }

            // ワークスペースをクリア
            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];
            }
        }
    }

    /// 前進代入: 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;
    }

    /// 後退代入: 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 分解
 *
 * L は単位下三角 (対角=1)、D は対角行列。sqrt 不要で LLT より安定。
 *
 * 使用例:
 * @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) {
        compute(A);
    }

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

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

        std::vector<IndexType> a_cp, a_ri;
        std::vector<T> a_val;
        detail::schol::extract_lower_csc(A, n, a_cp, a_ri, a_val);

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

    /// 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_) throw DimensionError("SimplicialLDLT::solve: dimension mismatch");

        Vector<T> y = forward_solve(b);
        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];
        }
        return backward_solve(y);
    }

    /// 単位下三角行列 L (対角 = 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;
    }

    /// 対角行列 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_;         // L の値 (対角なし, CSC)
    std::vector<IndexType> li_;  // L の行インデックス (CSC)
    std::vector<IndexType> lp_;  // L の列ポインタ (CSC, size n+1)
    std::vector<T> d_;           // D の対角
    bool factored_ = false;

    /**
     * @brief LDLT 左向き列分解
     *
     * cmod: x(i) -= L(j,k) * D(k) * L(i,k)
     * D(j) = x(j), L(i,j) = x(i) / D(j)
     *
     * L は対角なし (単位対角)。リンクリストの開始位置は
     * col_rows[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);

        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)  [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;
            }

            // 対角 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));
            }

            // 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]);
                }
            }

            // 列 j をリンク
            if (!cr[j].empty()) {
                cpos[j] = 0;
                IndexType nr = cr[j][0];
                nxt[j] = head[nr];
                head[nr] = j;
            }

            // クリア
            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];
            }
        }
    }

    /// 前進代入: L y = b (L は単位下三角)
    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;
    }

    /// 後退代入: L^T x = y (L は単位下三角)
    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;
    }
};

// ============================================================================
// 便利関数
// ============================================================================

/// 疎行列の Cholesky 分解 (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();
}

/// 疎行列の Cholesky で Ax = b を解く
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)
{
    SimplicialLLT<T, IndexType> chol(A);
    return chol.solve(b);
}

/// SimplicialLLT を前処理器として返す
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 calx

#endif // CALX_SIMPLICIAL_CHOLESKY_HPP