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

/**
 * @file sparse_lu.hpp
 * @brief LU factorization for sparse matrices (SparseLU)
 *
 * Left-looking column LU factorization + partial pivoting.
 * Stores L and U in CSC format, preserving the sparse structure.
 *
 * The API follows Eigen's SparseLU:
 *   SparseLU<T> lu;
 *   lu.compute(A);
 *   Vector<T> x = lu.solve(b);
 */

#ifndef SANGI_SPARSE_LU_HPP
#define SANGI_SPARSE_LU_HPP

#include <vector>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <stdexcept>
#include <limits>

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

namespace sangi {

/**
 * @brief Direct LU factorization for sparse matrices (CSC storage)
 *
 * P * A = L * U (P is a row permutation)
 * L: unit lower triangular (CSC), U: upper triangular (CSC)
 * Manages fill-in via left-looking column factorization.
 */
template<typename T, typename IndexType = std::size_t>
class SparseLU {
public:
    /// @param ordering fill-reducing ordering (default AMD). Since solve is permutation-invariant,
    ///        ordering is enabled by default to prevent fill-in blow-up on real problems.
    explicit SparseLU(SparseOrdering ordering = SparseOrdering::AMD)
        : ordering_(ordering) {}

    /// LU-factorizes a sparse matrix
    void compute(const SparseMatrix<T, IndexType>& A) {
        const auto n = static_cast<std::size_t>(A.rows());
        if (static_cast<std::size_t>(A.cols()) != n)
            throw std::invalid_argument("SparseLU: matrix must be square");

        n_ = n;
        computed_ = false;

        // Build the fill-reducing symmetric ordering P A Pᵀ as a working matrix.
        // (ordering is a permutation only and leaves values unchanged; partial pivoting is done on the working matrix)
        fill_perm_ = (ordering_ == SparseOrdering::AMD && n > 0)
                         ? algorithms::amd_ordering(A)
                         : algorithms::natural_ordering(n);

        SparseMatrix<T, IndexType> B(static_cast<IndexType>(n),
                                     static_cast<IndexType>(n),
                                     SparseStorageFormat::CSC);
        const bool reordered = (ordering_ == SparseOrdering::AMD && n > 0);
        if (reordered) {
            std::vector<std::size_t> inv(n);
            for (std::size_t k = 0; k < n; ++k) inv[fill_perm_[k]] = k;
            const auto& av = A.csc_values();
            const auto& ar = A.csc_row_indices();
            const auto& ac = A.csc_col_ptr();
            for (std::size_t j = 0; j < n; ++j) {
                for (auto p = static_cast<std::size_t>(ac[j]);
                     p < static_cast<std::size_t>(ac[j + 1]); ++p) {
                    const auto i = static_cast<std::size_t>(ar[p]);
                    B.set_coeff(static_cast<IndexType>(inv[i]),
                                static_cast<IndexType>(inv[j]), av[p]);
                }
            }
        }

        // Obtain CSC data (working matrix B when reordered, otherwise A)
        const auto& a_val = reordered ? B.csc_values()      : A.csc_values();
        const auto& a_row = reordered ? B.csc_row_indices()  : A.csc_row_indices();
        const auto& a_col = reordered ? B.csc_col_ptr()      : A.csc_col_ptr();

        // Permutation arrays
        perm_.resize(n);        // perm_[k] = original row number of pivot row k
        inv_perm_.resize(n);    // inv_perm_[original row] = current row number
        std::iota(perm_.begin(), perm_.end(), std::size_t(0));
        std::iota(inv_perm_.begin(), inv_perm_.end(), std::size_t(0));

        // Build L, U column by column (temporary vector of vector)
        std::vector<std::vector<std::size_t>> L_rows(n), U_rows(n);
        std::vector<std::vector<T>> L_vals(n), U_vals(n);

        // Dense work vectors (reused per column)
        std::vector<T> x(n, T(0));
        std::vector<char> x_nz(n, 0);  // nonzero flags (vector<bool> is a proxy and cannot swap)
        std::vector<std::size_t> nz_list;  // list of nonzero positions
        nz_list.reserve(n);

        for (std::size_t k = 0; k < n; ++k) {
            // 1. Scatter column k of A (original row number → current row number)
            nz_list.clear();
            auto col_start = static_cast<std::size_t>(a_col[k]);
            auto col_end   = static_cast<std::size_t>(a_col[k + 1]);
            for (auto p = col_start; p < col_end; ++p) {
                std::size_t orig_row = static_cast<std::size_t>(a_row[p]);
                std::size_t cur_row = inv_perm_[orig_row];
                x[cur_row] = a_val[p];
                if (!x_nz[cur_row]) {
                    x_nz[cur_row] = 1;
                    nz_list.push_back(cur_row);
                }
            }

            // 2. Left-looking solve: for j = 0..k-1, if x[j] ≠ 0 update with L(:,j).
            //    A Gilbert-Peierls style topological DFS was implemented (see commit history),
            //    but benchmarks showed that at typical IPM scales of n=50-500 the per-node
            //    DFS overhead (~50 cycles) exceeds the simple loop's skip cost (~2 cycles),
            //    making it 20-30% slower. GP only wins in ultra-sparse cases with n>>1000
            //    where reach is a few % of n, but in IPM M = A D Aᵀ is not that
            //    sparse, so the simple loop is adopted (details: TODO/PLAN_BENCH_IPM_SPARSE_PHASE2.md).
            for (std::size_t j = 0; j < k; ++j) {
                if (!x_nz[j]) continue;
                if (x[j] == T(0)) continue;

                // Update with each entry of L(:,j)
                for (std::size_t q = 0; q < L_rows[j].size(); ++q) {
                    std::size_t i = L_rows[j][q];
                    T lij = L_vals[j][q];
                    x[i] -= lij * x[j];
                    if (!x_nz[i] && x[i] != T(0)) {
                        x_nz[i] = 1;
                        nz_list.push_back(i);
                    }
                }
            }

            // Sort + deduplicate nz_list for pivot search
            std::sort(nz_list.begin(), nz_list.end());
            nz_list.erase(std::unique(nz_list.begin(), nz_list.end()), nz_list.end());

            // 3. Partial pivoting: find the maximum value in x[k:n-1]
            std::size_t pivot = k;
            T pivot_val = std::abs(x[k]);
            for (auto r : nz_list) {
                if (r < k) continue;
                T v = std::abs(x[r]);
                if (v > pivot_val) {
                    pivot_val = v;
                    pivot = r;
                }
            }

            if (pivot_val < std::numeric_limits<T>::epsilon() * T(100))
                throw std::runtime_error("SparseLU: singular matrix");

            // 4. Row swap (pivot ↔ k)
            if (pivot != k) {
                std::swap(x[k], x[pivot]);
                std::swap(x_nz[k], x_nz[pivot]);

                // Update permutation
                std::size_t ok = perm_[k], op = perm_[pivot];
                std::swap(perm_[k], perm_[pivot]);
                inv_perm_[ok] = pivot;
                inv_perm_[op] = k;

                // Swap rows k and pivot in the existing columns of L (0..k-1)
                for (std::size_t j = 0; j < k; ++j) {
                    T* vk = nullptr;
                    T* vp = nullptr;
                    for (std::size_t q = 0; q < L_rows[j].size(); ++q) {
                        if (L_rows[j][q] == k) vk = &L_vals[j][q];
                        if (L_rows[j][q] == pivot) vp = &L_vals[j][q];
                    }
                    if (vk && vp) {
                        std::swap(*vk, *vp);
                    } else if (vk) {
                        // Present in k but not in pivot
                        // Change the k position to pivot
                        for (std::size_t q = 0; q < L_rows[j].size(); ++q) {
                            if (L_rows[j][q] == k) { L_rows[j][q] = pivot; break; }
                        }
                    } else if (vp) {
                        for (std::size_t q = 0; q < L_rows[j].size(); ++q) {
                            if (L_rows[j][q] == pivot) { L_rows[j][q] = k; break; }
                        }
                    }
                }
            }

            // 5. U column k: the nonzero part of x[0..k]
            for (auto r : nz_list) {
                if (r > k) break;
                if (x[r] != T(0)) {
                    U_rows[k].push_back(r);
                    U_vals[k].push_back(x[r]);
                }
            }
            // Add the diagonal if not already included (numerically near zero)
            T diag = x[k];

            // 6. L column k: the nonzero part of x[k+1..n-1] / diag
            T inv_diag = T(1) / diag;
            for (auto r : nz_list) {
                if (r <= k) continue;
                if (x[r] != T(0)) {
                    L_rows[k].push_back(r);
                    L_vals[k].push_back(x[r] * inv_diag);
                }
            }

            // 7. Clear the work vectors (nonzero positions only)
            for (auto r : nz_list) {
                x[r] = T(0);
                x_nz[r] = 0;
            }
        }

        // Convert to CSC format
        build_csc(L_rows, L_vals, L_col_ptr_, L_row_ind_, L_val_);
        build_csc(U_rows, U_vals, U_col_ptr_, U_row_ind_, U_val_);

        computed_ = true;
    }

    /// Solves the linear system A*x = b
    [[nodiscard]] Vector<T> solve(const Vector<T>& b) const {
        if (!computed_)
            throw std::runtime_error("SparseLU: compute() not called");
        if (b.size() != n_)
            throw std::invalid_argument("SparseLU::solve: dimension mismatch");

        // Map the fill-reducing symmetric ordering to working coordinates: c = P b
        Vector<T> c(n_);
        for (std::size_t i = 0; i < n_; ++i)
            c[i] = b[fill_perm_[i]];

        // P * c (apply the row-pivot permutation)
        Vector<T> y(n_);
        for (std::size_t i = 0; i < n_; ++i)
            y[i] = c[perm_[i]];

        // Forward substitution: L * z = y (L is unit lower triangular, CSC)
        // Process column k: z[i] -= L(i,k) * z[k]  (i > k)
        for (std::size_t k = 0; k < n_; ++k) {
            if (y[k] == T(0)) continue;
            auto start = L_col_ptr_[k];
            auto end = L_col_ptr_[k + 1];
            for (auto p = start; p < end; ++p)
                y[L_row_ind_[p]] -= L_val_[p] * y[k];
        }

        // Back substitution: U * x = z (U is upper triangular, CSC)
        // Process column k from the right: x[k] = (z[k] - sum U(i,k)*x[i]) / U(k,k)
        for (std::size_t kk = 0; kk < n_; ++kk) {
            std::size_t k = n_ - 1 - kk;
            auto start = U_col_ptr_[k];
            auto end = U_col_ptr_[k + 1];
            // The diagonal of U column k is the last element (row number k)
            T diag = T(0);
            for (auto p = start; p < end; ++p) {
                if (U_row_ind_[p] == k) {
                    diag = U_val_[p];
                } else if (U_row_ind_[p] < k) {
                    // nothing — will handle in column processing from right
                }
            }
            y[k] /= diag;
            // Update with the off-diagonal entries of column k
            for (auto p = start; p < end; ++p) {
                if (U_row_ind_[p] < k)
                    y[U_row_ind_[p]] -= U_val_[p] * y[k];
            }
        }

        // Working coordinates → original coordinates: x[fill_perm_[i]] = y[i]
        Vector<T> x(n_);
        for (std::size_t i = 0; i < n_; ++i)
            x[fill_perm_[i]] = y[i];
        return x;
    }

    [[nodiscard]] bool computed() const { return computed_; }
    [[nodiscard]] std::size_t rows() const { return n_; }
    [[nodiscard]] std::size_t cols() const { return n_; }
    /// fill-reducing permutation (perm[k] = original index eliminated at step k)
    [[nodiscard]] const std::vector<std::size_t>& permutation() const { return fill_perm_; }

private:
    SparseOrdering ordering_ = SparseOrdering::AMD;
    std::size_t n_ = 0;
    bool computed_ = false;
    std::vector<std::size_t> fill_perm_;  // fill-reducing symmetric ordering
    std::vector<std::size_t> perm_;
    std::vector<std::size_t> inv_perm_;

    // Store L (unit lower triangular, diagonal implicitly 1) in CSC
    std::vector<std::size_t> L_col_ptr_;
    std::vector<std::size_t> L_row_ind_;
    std::vector<T> L_val_;

    // Store U (upper triangular) in CSC
    std::vector<std::size_t> U_col_ptr_;
    std::vector<std::size_t> U_row_ind_;
    std::vector<T> U_val_;

    /// vector of vector → CSC conversion
    static void build_csc(const std::vector<std::vector<std::size_t>>& rows,
                           const std::vector<std::vector<T>>& vals,
                           std::vector<std::size_t>& col_ptr,
                           std::vector<std::size_t>& row_ind,
                           std::vector<T>& val) {
        const auto n = rows.size();
        col_ptr.resize(n + 1);
        col_ptr[0] = 0;
        std::size_t total = 0;
        for (std::size_t j = 0; j < n; ++j) {
            total += rows[j].size();
            col_ptr[j + 1] = total;
        }
        row_ind.resize(total);
        val.resize(total);
        std::size_t pos = 0;
        for (std::size_t j = 0; j < n; ++j) {
            for (std::size_t q = 0; q < rows[j].size(); ++q) {
                row_ind[pos] = rows[j][q];
                val[pos] = vals[j][q];
                ++pos;
            }
        }
    }
};

} // namespace sangi

#endif // SANGI_SPARSE_LU_HPP