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

// balance.hpp
//
// Matrix balancing / equilibration preprocessing
//
// This file provides 2 kinds of diagonal scaling:
//
//  (1) balance_matrix — balancing by a similarity transform (= equivalent to LAPACK dgebal)
//      A → P^T D^{-1} A D P (D = diag(2^k), P = permutation). Preserves eigenvalues and, by
//      equalizing the off-diagonal norms of each row and column, improves the numerical stability of Hessenberg / Schur / QR iterations.
//      Used as preprocessing for eigenvalue computations such as eigen() / power_method / Riccati Schur.
//
//      Steps:
//        Step A (permutation): push isolated eigenvalues (= rows/columns that reference only themselves)
//                              to the ends, extracting the nontrivial block [ilo, ihi].
//        Step B (scaling)    : equalize the row norms and column norms over the
//                              [ilo, ihi] interval with the iterative power-of-2 scaling of Parlett-Reinsch (1969).
//                              Powers of 2 are exact in floating-point arithmetic, with zero error.
//
//      Inverse transforms (= equivalent to dgebak):
//        unbalance_eigenvectors(V, info)  — return columns to A's eigenvector space.
//        unbalance_vector(v, info)        — inverse transform of a single vector.
//        rebalance_invariant_subspace(Z, info) — for invariant subspaces (LQR Laub form).
//
//  (2) equilibrate_matrix — independent left/right scaling for Ax = b (= equivalent to LAPACK dgeequ)
//      A → D_l A D_r, b → D_l b. The solution is recovered by x = D_r y. Effective as preprocessing
//      for iterative solvers (CG / GMRES / BiCGSTAB) (improves diagonal dominance and the convergence rate).
//
// References:
//   Parlett, B. N. & Reinsch, C. (1969) "Balancing a matrix for calculation of
//     eigenvalues and eigenvectors", Numer. Math. 13, 293-304.
//   LAPACK Users' Guide 3rd ed. §4.10 (balancing) / §4.11 (equilibration).
//   Anderson et al. 1999, LAPACK source dgebal.f / dgebak.f / dgeequ.f / dlaqge.f.

#ifndef SANGI_BALANCE_HPP
#define SANGI_BALANCE_HPP

#include <math/core/common.hpp>
#include <math/core/matrix.hpp>
#include <math/core/vector.hpp>
#include <math/core/Complex.hpp>

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

namespace sangi {
    namespace algorithms {

        //=====================================================================
        // BalanceInfo — output of balance_matrix (used by dgebak)
        //=====================================================================

        /**
         * @brief Output information of balance_matrix
         *
         * Both the permutation and the scaling are encoded into a single array perm using the same sign
         * convention as LAPACK (= dgebal compatible):
         *   - i < ilo                : Step A left-side push-out permutation partner = perm[i] (size_t)
         *   - ihi < i                : Step A right-side push-out permutation partner = perm[i]
         *   - ilo <= i <= ihi        : Step B scaling factor D(i,i) (perm[i] is unused, set to 0)
         *
         * scale[i] holds D(i, i) = scale[i] over the [ilo, ihi] interval. Outside it is 1.
         * Both are powers of 2 (exact in floating-point arithmetic, with zero error).
         */
        template<typename T>
        struct BalanceInfo {
            std::size_t n   = 0;     ///< size of the original matrix
            std::size_t ilo = 0;     ///< start of the nontrivial block (0-based)
            std::size_t ihi = 0;     ///< end of the nontrivial block (0-based, inclusive)
            std::vector<std::size_t> perm; ///< Step A permutation partners (size n, meaningful only outside)
            std::vector<T> scale;          ///< Step B diagonal scaling (size n, nontrivial only over [ilo,ihi])

            /// Build a no-op BalanceInfo corresponding to the identity matrix for all elements.
            static BalanceInfo<T> identity(std::size_t n) {
                BalanceInfo<T> info;
                info.n   = n;
                info.ilo = 0;
                info.ihi = (n == 0) ? 0 : n - 1;
                info.perm.assign(n, 0);
                info.scale.assign(n, T(1));
                return info;
            }
        };

        //=====================================================================
        // balance_matrix — A → P^T D^{-1} A D P (similarity transform)
        //=====================================================================

        /**
         * @brief Balancing of a general n×n matrix (in-place, equivalent to LAPACK dgebal)
         *
         * @param A in/out square matrix. Rewritten to P^T D^{-1} A D P.
         * @return BalanceInfo (Step A permutation + Step B scaling)
         *
         * @note Eigenvalues are preserved because this is a similarity transform.
         * @note Step A (permutation), as in LAPACK, pushes to the ends the columns/rows for which
         *       "row i has no off-diagonal elements" or "column i has no off-diagonal elements".
         *       This excludes diagonally dominant isolated eigenvalues from the QR iteration, improving the operation count and numerical stability.
         * @note Step B (scaling) iterates the Parlett-Reinsch (1969) procedure with threshold 0.95 until no further improvement is seen.
         */
        template<typename T>
        BalanceInfo<T> balance_matrix(BaseMatrix<T>& A) {
            const std::size_t n = A.rows();
            if (A.cols() != n) {
                throw DimensionError("balance_matrix: A must be square");
            }

            BalanceInfo<T> info;
            info.n = n;
            info.perm.assign(n, std::size_t{0});
            info.scale.assign(n, T(1));

            if (n == 0) { info.ilo = 0; info.ihi = 0; return info; }
            if (n == 1) { info.ilo = 0; info.ihi = 0; return info; }

            std::size_t ilo = 0;
            std::size_t ihi = n - 1;

            //---------------------------------------------------------------
            // Step A: push isolated eigenvalues outside [ilo, ihi] via permutation
            //
            //   Right-side scan: if "all off-diagonal elements are 0 in column j (ilo<=j<=ihi)"
            //                  → swap column j to the ihi position, ihi--
            //   Left-side scan:  if "all off-diagonal elements are 0 in row i (ilo<=i<=ihi)"
            //                  → swap row i to the ilo position, ilo++
            //   If either is found, re-scan (= the same deflation loop as LAPACK dgebal)
            //---------------------------------------------------------------

            // Helper: swap row i with row j, and column i with column j
            auto swap_rows_cols = [&](std::size_t i, std::size_t j) {
                if (i == j) return;
                for (std::size_t k = 0; k < n; ++k) {
                    std::swap(A(i, k), A(j, k));
                }
                for (std::size_t k = 0; k < n; ++k) {
                    std::swap(A(k, i), A(k, j));
                }
            };

            // Right-side scan: does the column have no off-diagonal elements?
            auto col_isolated = [&](std::size_t col) {
                for (std::size_t r = ilo; r <= ihi; ++r) {
                    if (r != col && A(r, col) != T(0)) return false;
                }
                return true;
            };
            // Left-side scan: does the row have no off-diagonal elements?
            auto row_isolated = [&](std::size_t row) {
                for (std::size_t c = ilo; c <= ihi; ++c) {
                    if (c != row && A(row, c) != T(0)) return false;
                }
                return true;
            };

            // Right-side deflation: push isolated columns toward the ihi side
            bool found = true;
            while (found && ihi > ilo) {
                found = false;
                // Scan from the back (decrement ihi as soon as one is found)
                for (std::size_t j = ihi + 1; j-- > ilo; ) {
                    if (col_isolated(j)) {
                        // Swap column j with the ihi position, and record the permutation partner in perm
                        swap_rows_cols(j, ihi);
                        info.perm[ihi] = j;  // "the original position of ihi was j"
                        if (ihi == 0) break;
                        --ihi;
                        found = true;
                        break;
                    }
                    if (j == ilo) break;  // loop terminator
                }
            }

            // Left-side deflation: push isolated rows toward the ilo side
            found = true;
            while (found && ilo < ihi) {
                found = false;
                for (std::size_t i = ilo; i <= ihi; ++i) {
                    if (row_isolated(i)) {
                        swap_rows_cols(i, ilo);
                        info.perm[ilo] = i;  // "the original position of ilo was i"
                        ++ilo;
                        found = true;
                        break;
                    }
                }
            }

            info.ilo = ilo;
            info.ihi = ihi;

            //---------------------------------------------------------------
            // Step B: Parlett-Reinsch scaling
            //
            //   For each row i over the [ilo, ihi] interval
            //     c = Σ_{j != i, j in [ilo,ihi]} |A(j, i)|   (off-diagonal norm of column i)
            //     r = Σ_{j != i, j in [ilo,ihi]} |A(i, j)|   (off-diagonal norm of row i)
            //   If c and r are strongly asymmetric, find a scaling factor f = 2^k and
            //     multiply row i of A by f^{-1} and column i by f (= the similarity transform D^{-1} A D)
            //   Iterate with threshold 0.95 until improvement stops (= converges in O(n)).
            //---------------------------------------------------------------

            if (ihi > ilo) {
                const T B   = T(2);
                const T BSQ = B * B;
                bool changed = true;
                std::size_t max_iter = (ihi - ilo + 1) * 4 + 32;
                while (changed && max_iter-- > 0) {
                    changed = false;
                    for (std::size_t i = ilo; i <= ihi; ++i) {
                        T c = T(0), r = T(0);
                        for (std::size_t j = ilo; j <= ihi; ++j) {
                            if (j == i) continue;
                            c += std::abs(A(j, i));
                            r += std::abs(A(i, j));
                        }
                        if (c == T(0) || r == T(0)) continue;
                        const T s = c + r;
                        T f = T(1);
                        T g = r / B;
                        while (c < g) { c *= BSQ; f *= B; }
                        g = r * B;
                        while (c >= g) { c /= BSQ; f /= B; }
                        if ((c + r) / f < T(0.95) * s) {
                            const T inv_f = T(1) / f;
                            info.scale[i] *= f;
                            for (std::size_t j = 0; j < n; ++j) {
                                A(i, j) *= inv_f;
                                A(j, i) *= f;
                            }
                            changed = true;
                        }
                    }
                }
            }

            return info;
        }

        //=====================================================================
        // unbalance_* — recover A's eigenvectors / invariant subspaces (equivalent to dgebak)
        //=====================================================================

        namespace detail_balance {
            // Inverse-transform column j by Step B (scaling): V(i, j) *= D(i, i)  (i ∈ [ilo, ihi])
            template<typename V, typename T>
            void apply_scaling_to_column(V& M, std::size_t j,
                                         const BalanceInfo<T>& info) {
                for (std::size_t i = info.ilo; i <= info.ihi; ++i) {
                    const T s = info.scale[i];
                    if (s == T(1)) continue;
                    M(i, j) = M(i, j) * s;
                }
            }

            // Inverse-transform by Step A (permutation): for column j, swap row i (i outside) with perm[i]
            template<typename V, typename T>
            void apply_permutation_to_column(V& M, std::size_t j,
                                             const BalanceInfo<T>& info) {
                // Same order as dgebak: ilo side (left, reverse order from ilo-1) → ihi side (right, from ihi+1)
                // namely a full reverse replay of the order in which balance_matrix recorded the swaps
                // (the left side was recorded in ilo++ order → reverse replay is ilo--)
                const std::size_t n = info.n;
                if (info.ilo > 0) {
                    for (std::size_t k = info.ilo; k-- > 0; ) {
                        const std::size_t p = info.perm[k];
                        if (p != k) {
                            auto tmp = M(k, j);
                            M(k, j) = M(p, j);
                            M(p, j) = tmp;
                        }
                    }
                }
                if (info.ihi + 1 < n) {
                    for (std::size_t k = info.ihi + 1; k < n; ++k) {
                        const std::size_t p = info.perm[k];
                        if (p != k) {
                            auto tmp = M(k, j);
                            M(k, j) = M(p, j);
                            M(p, j) = tmp;
                        }
                    }
                }
            }
        } // namespace detail_balance

        /**
         * @brief Inverse balance of an eigenvector matrix (per column, equivalent to dgebak)
         *
         * @param V in/out  column j = eigenvector v_b of balanced A → eigenvector v of A
         *                  inverse transform: v = P D v_b
         * @param info return value of balance_matrix
         */
        template<typename T>
        void unbalance_eigenvectors(BaseMatrix<T>& V, const BalanceInfo<T>& info) {
            if (V.rows() != info.n) {
                throw DimensionError("unbalance_eigenvectors: row count must match BalanceInfo::n");
            }
            const std::size_t cols = V.cols();
            for (std::size_t j = 0; j < cols; ++j) {
                detail_balance::apply_scaling_to_column(V, j, info);
                detail_balance::apply_permutation_to_column(V, j, info);
            }
        }

        /// Overload for complex matrices (for the eigenvectors returned by eigen())
        template<typename T>
        void unbalance_eigenvectors(BaseMatrix<Complex<T>>& V, const BalanceInfo<T>& info) {
            if (V.rows() != info.n) {
                throw DimensionError("unbalance_eigenvectors: row count must match BalanceInfo::n");
            }
            const std::size_t cols = V.cols();
            for (std::size_t j = 0; j < cols; ++j) {
                detail_balance::apply_scaling_to_column(V, j, info);
                detail_balance::apply_permutation_to_column(V, j, info);
            }
        }

        /**
         * @brief Inverse balance of a single eigenvector
         */
        template<typename T>
        void unbalance_vector(BaseVector<T>& v, const BalanceInfo<T>& info) {
            if (v.size() != info.n) {
                throw DimensionError("unbalance_vector: size must match BalanceInfo::n");
            }
            // Same processing, treating the Vector as a 1-column Matrix
            // It is clearer to implement it directly, so we do that
            for (std::size_t i = info.ilo; i <= info.ihi; ++i) {
                const T s = info.scale[i];
                if (s == T(1)) continue;
                v[i] *= s;
            }
            const std::size_t n = info.n;
            if (info.ilo > 0) {
                for (std::size_t k = info.ilo; k-- > 0; ) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            if (info.ihi + 1 < n) {
                for (std::size_t k = info.ihi + 1; k < n; ++k) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
        }

        template<typename T>
        void unbalance_vector(BaseVector<Complex<T>>& v, const BalanceInfo<T>& info) {
            if (v.size() != info.n) {
                throw DimensionError("unbalance_vector: size must match BalanceInfo::n");
            }
            for (std::size_t i = info.ilo; i <= info.ihi; ++i) {
                const T s = info.scale[i];
                if (s == T(1)) continue;
                v[i] = v[i] * Complex<T>(s, T(0));
            }
            const std::size_t n = info.n;
            if (info.ilo > 0) {
                for (std::size_t k = info.ilo; k-- > 0; ) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            if (info.ihi + 1 < n) {
                for (std::size_t k = info.ihi + 1; k < n; ++k) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
        }

        /**
         * @brief Reverse balance of an invariant subspace (LQR Laub form)
         *
         * @param Z in/out  Schur vectors of balanced A → invariant subspace of A
         *                  for each column j, Z(i, j) ← D(i, i) · Z(i, j) (scaling only)
         *                  + reverse-permute the rows by the Step A permutation
         *
         * @note The orthogonality of Z is lost, but it is correct as an invariant subspace.
         *       Valid for the U_2 U_1^{-1} computation of LQR Laub.
         */
        template<typename T>
        void rebalance_invariant_subspace(BaseMatrix<T>& Z, const BalanceInfo<T>& info) {
            // The implementation is the same as unbalance_eigenvectors (per-column scaling + inverse permutation transform)
            unbalance_eigenvectors(Z, info);
        }

        //=====================================================================
        // balance_vector — convert an original-coordinate vector to balanced coordinates (reverse direction of dgebak)
        //=====================================================================

        /**
         * @brief Convert an original-coordinate vector v to balanced coordinates v_b (= forward balance)
         *
         * @param v in/out  input is original coordinates, output is balanced coordinates
         *                  v_b = D^{-1} P^{-1} v, the exact opposite processing of unbalance_vector
         *
         * @note Used to apply a user-specified initial guess (eigenvalue iteration etc.) on the
         *       balanced matrix. After computation, return to original coordinates with unbalance_vector.
         */
        template<typename T>
        void balance_vector(BaseVector<T>& v, const BalanceInfo<T>& info) {
            if (v.size() != info.n) {
                throw DimensionError("balance_vector: size must match BalanceInfo::n");
            }
            // Step 1: apply P (permutation in original order)
            //   Right side (= the first recorded swap group): descending from perm[n-1] to perm[ihi+1]
            const std::size_t n = info.n;
            if (info.ihi + 1 < n) {
                for (std::size_t k = n; k-- > info.ihi + 1; ) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            //   Left side (= the next recorded swap group): ascending from perm[0] to perm[ilo-1]
            if (info.ilo > 0) {
                for (std::size_t k = 0; k < info.ilo; ++k) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            // Step 2: D^{-1} (= divide by the scaling factor)
            for (std::size_t i = info.ilo; i <= info.ihi; ++i) {
                const T s = info.scale[i];
                if (s == T(1)) continue;
                v[i] = v[i] / s;
            }
        }

        template<typename T>
        void balance_vector(BaseVector<Complex<T>>& v, const BalanceInfo<T>& info) {
            if (v.size() != info.n) {
                throw DimensionError("balance_vector: size must match BalanceInfo::n");
            }
            const std::size_t n = info.n;
            if (info.ihi + 1 < n) {
                for (std::size_t k = n; k-- > info.ihi + 1; ) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            if (info.ilo > 0) {
                for (std::size_t k = 0; k < info.ilo; ++k) {
                    const std::size_t p = info.perm[k];
                    if (p != k) std::swap(v[k], v[p]);
                }
            }
            for (std::size_t i = info.ilo; i <= info.ihi; ++i) {
                const T s = info.scale[i];
                if (s == T(1)) continue;
                v[i] = v[i] * Complex<T>(T(1) / s, T(0));
            }
        }

        //=====================================================================
        // EquilibrateInfo / equilibrate_matrix — independent left/right scaling for Ax = b
        //=====================================================================

        /**
         * @brief Output information of equilibrate_matrix
         *
         * D_l = diag(row_scale), D_r = diag(col_scale), A → D_l A D_r, b → D_l b.
         * The solution is recovered by x = D_r y.
         */
        template<typename T>
        struct EquilibrateInfo {
            std::vector<T> row_scale;
            std::vector<T> col_scale;
            T row_cond  = T(1);  ///< condition number of the row scaling (= max/min row_scale)
            T col_cond  = T(1);  ///< condition number of the column scaling
            T amax      = T(0);  ///< max|A(i,j)| before scaling
            bool row_equilibrated = false;
            bool col_equilibrated = false;
        };

        /**
         * @brief Independent left/right row/column scaling (equivalent to LAPACK dgeequ + dlaqge)
         *
         * @param A in/out  m×n matrix. Rewritten to D_l A D_r.
         * @return EquilibrateInfo
         *
         * @note row_scale[i] = 1 / max_j |A(i, j)|
         *       col_scale[j] = 1 / max_i |D_l(i,i) A(i, j)|
         *       The LAPACK way is to skip applying scaling if the condition-number improvement is small (cond < ~0.1), but
         *       this implementation always applies it for simplicity. If needed, the caller can decide by looking at row_cond / col_cond.
         *
         * @note When solving Ax = b, after equilibrate_matrix(A), multiply b by row_scale,
         *       and multiply the obtained y by col_scale to recover x (see unequilibrate_solution).
         */
        template<typename T>
        EquilibrateInfo<T> equilibrate_matrix(BaseMatrix<T>& A) {
            const std::size_t m = A.rows();
            const std::size_t n = A.cols();
            EquilibrateInfo<T> info;
            info.row_scale.assign(m, T(1));
            info.col_scale.assign(n, T(1));

            if (m == 0 || n == 0) return info;

            // amax = max|A(i, j)|
            T amax = T(0);
            for (std::size_t i = 0; i < m; ++i) {
                for (std::size_t j = 0; j < n; ++j) {
                    const T v = std::abs(A(i, j));
                    if (v > amax) amax = v;
                }
            }
            info.amax = amax;
            if (amax == T(0)) return info;

            // Row scale
            T r_max = T(0), r_min = std::numeric_limits<T>::max();
            for (std::size_t i = 0; i < m; ++i) {
                T row_amax = T(0);
                for (std::size_t j = 0; j < n; ++j) {
                    const T v = std::abs(A(i, j));
                    if (v > row_amax) row_amax = v;
                }
                if (row_amax > T(0)) {
                    info.row_scale[i] = T(1) / row_amax;
                    if (row_amax > r_max) r_max = row_amax;
                    if (row_amax < r_min) r_min = row_amax;
                }
            }
            info.row_cond = (r_max > T(0)) ? r_min / r_max : T(1);
            // Apply
            for (std::size_t i = 0; i < m; ++i) {
                const T s = info.row_scale[i];
                if (s == T(1)) continue;
                for (std::size_t j = 0; j < n; ++j) {
                    A(i, j) *= s;
                }
            }
            info.row_equilibrated = true;

            // Column scale
            T c_max = T(0), c_min = std::numeric_limits<T>::max();
            for (std::size_t j = 0; j < n; ++j) {
                T col_amax = T(0);
                for (std::size_t i = 0; i < m; ++i) {
                    const T v = std::abs(A(i, j));
                    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);
            for (std::size_t j = 0; j < n; ++j) {
                const T s = info.col_scale[j];
                if (s == T(1)) continue;
                for (std::size_t i = 0; i < m; ++i) {
                    A(i, j) *= s;
                }
            }
            info.col_equilibrated = true;

            return info;
        }

        /**
         * @brief Apply row_scale to the equilibrated right-hand side b (b → D_l b)
         */
        template<typename T>
        void equilibrate_rhs(BaseVector<T>& b, const EquilibrateInfo<T>& info) {
            if (!info.row_equilibrated) return;
            if (b.size() != info.row_scale.size()) {
                throw DimensionError("equilibrate_rhs: size mismatch");
            }
            for (std::size_t i = 0; i < b.size(); ++i) {
                b[i] *= info.row_scale[i];
            }
        }

        /**
         * @brief Recover the original x from the solution y obtained in the equilibrated system (x = D_r y)
         */
        template<typename T>
        void unequilibrate_solution(BaseVector<T>& y, const EquilibrateInfo<T>& info) {
            if (!info.col_equilibrated) return;
            if (y.size() != info.col_scale.size()) {
                throw DimensionError("unequilibrate_solution: size mismatch");
            }
            for (std::size_t j = 0; j < y.size(); ++j) {
                y[j] *= info.col_scale[j];
            }
        }

    } // namespace algorithms
} // namespace sangi

#endif // SANGI_BALANCE_HPP