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

// balance_sparse.hpp
//
// Equilibration extension for SparseMatrix (= sparse version of LAPACK dgeequ).
//
// Useful as a preconditioning step when solving ill-conditioned sparse systems
// with iterative solvers (CG / BiCGSTAB / GMRES / MINRES / FGMRES). It reuses
// the same EquilibrateInfo<T> + equilibrate_rhs + unequilibrate_solution as the
// dense equilibrate_matrix.
//
// Recommended pattern:
//   auto [A_eq, info] = make_equilibrated_sparse(A);
//   auto b_eq = b;
//   equilibrate_rhs(b_eq, info);
//   auto result = bicgstab_logged(A_eq, b_eq, criteria, log);
//   unequilibrate_solution(result.x, info);
//
// Symmetric systems (CG / MINRES) require D_l = D_r so that D_l A D_r stays
// symmetric. This implementation scales the left and right sides independently
// (= non-symmetric), so it is not suitable for symmetric systems. For those,
// use jacobi_preconditioner (= the existing D = 1/sqrt(diag(A))).

#ifndef SANGI_BALANCE_SPARSE_HPP
#define SANGI_BALANCE_SPARSE_HPP

#include <math/core/sparse_matrix.hpp>
#include <math/linalg/balance.hpp>

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <limits>
#include <utility>

namespace sangi {
    namespace algorithms {

        /**
         * @brief Independent left/right scaling of a sparse matrix
         *        (sparse equivalent of LAPACK dgeequ + dlaqge)
         *
         * @param A in const sparse matrix
         * @return {A_eq, info} pair. A_eq = D_l A D_r, with the sparsity pattern
         *         preserved (an entry is dropped only if it becomes effectively zero).
         *
         * @note Algorithm:
         *   1. Take max |A(i,j)| over each row, then row_scale[i] = 1 / row_max[i]
         *   2. Take max |A(i,j)| over each column after row scaling, then
         *      col_scale[j] = 1 / col_max[j]
         *   3. Build a new SparseMatrix with A_eq(i,j) = row_scale[i] * A(i,j) * col_scale[j]
         *
         * @note Use with an iterative solver:
         *     auto [A_eq, info] = make_equilibrated_sparse(A);
         *     equilibrate_rhs(b, info);  // b ← D_l b
         *     auto result = solver(A_eq, b, ...);
         *     unequilibrate_solution(result.x, info);  // x ← D_r y
         */
        template<typename T, typename IndexType>
        std::pair<SparseMatrix<T, IndexType>, EquilibrateInfo<T>>
        make_equilibrated_sparse(const SparseMatrix<T, IndexType>& A) {
            using IT = IndexType;
            const IT m = A.rows();
            const IT n = A.cols();

            EquilibrateInfo<T> info;
            info.row_scale.assign(m, T(1));
            info.col_scale.assign(n, T(1));

            // Empty matrix → no-op
            SparseMatrix<T, IndexType> A_eq(m, n, SparseStorageFormat::COO);
            if (m == 0 || n == 0) return {std::move(A_eq), info};

            // SparseMatrix exposes only CSC accessors (CSR is not public).
            // Scan all nnz in CSC order, accumulating row_max, and resolve
            // col_max one column at a time.
            const auto& csc_vals = A.csc_values();
            const auto& csc_rows = A.csc_row_indices();
            const auto& csc_ptr  = A.csc_col_ptr();

            std::vector<T> row_max(static_cast<std::size_t>(m), T(0));
            T amax = T(0);

            for (IT j = 0; j < n; ++j) {
                const std::size_t s = static_cast<std::size_t>(csc_ptr[j]);
                const std::size_t e = static_cast<std::size_t>(csc_ptr[j + 1]);
                for (std::size_t k = s; k < e; ++k) {
                    const IT i = csc_rows[k];
                    const T v = std::abs(csc_vals[k]);
                    if (v > row_max[i]) row_max[i] = v;
                    if (v > amax) amax = v;
                }
            }
            info.amax = amax;
            if (amax == T(0)) return {std::move(A_eq), info};

            // row_scale[i] = 1 / row_max[i]
            T r_max = T(0), r_min = std::numeric_limits<T>::max();
            for (IT i = 0; i < m; ++i) {
                if (row_max[i] > T(0)) {
                    info.row_scale[i] = T(1) / row_max[i];
                    if (row_max[i] > r_max) r_max = row_max[i];
                    if (row_max[i] < r_min) r_min = row_max[i];
                }
            }
            info.row_cond = (r_max > T(0)) ? r_min / r_max : T(1);
            info.row_equilibrated = true;

            // Compute col_max after row scaling (= max_i |row_scale[i] · A(i, j)|).
            T c_max = T(0), c_min = std::numeric_limits<T>::max();
            for (IT j = 0; j < n; ++j) {
                T col_amax = T(0);
                const std::size_t s = static_cast<std::size_t>(csc_ptr[j]);
                const std::size_t e = static_cast<std::size_t>(csc_ptr[j + 1]);
                for (std::size_t k = s; k < e; ++k) {
                    const IT i = csc_rows[k];
                    const T v = std::abs(csc_vals[k]) * info.row_scale[i];
                    if (v > col_amax) col_amax = v;
                }
                if (col_amax > T(0)) {
                    info.col_scale[j] = T(1) / col_amax;
                    if (col_amax > c_max) c_max = col_amax;
                    if (col_amax < c_min) c_min = col_amax;
                }
            }
            info.col_cond = (c_max > T(0)) ? c_min / c_max : T(1);
            info.col_equilibrated = true;

            // Build the new sparse matrix: A_eq(i, j) = row_scale[i] * A(i, j) * col_scale[j]
            for (IT j = 0; j < n; ++j) {
                const T cs = info.col_scale[j];
                const std::size_t s = static_cast<std::size_t>(csc_ptr[j]);
                const std::size_t e = static_cast<std::size_t>(csc_ptr[j + 1]);
                for (std::size_t k = s; k < e; ++k) {
                    const IT i = csc_rows[k];
                    const T new_val = info.row_scale[i] * csc_vals[k] * cs;
                    A_eq.set_coeff(i, j, new_val);
                }
            }

            return {std::move(A_eq), info};
        }

    } // namespace algorithms
} // namespace sangi

#endif // SANGI_BALANCE_SPARSE_HPP