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

// sylvester.hpp
//
// Sylvester equation solver: A*X + X*B = C
//
// Algorithm: Bartels-Stewart method (1972)
//   1. Real Schur decomposition: A = U R U^T, B = V S V^T
//   2. Transform: C' = U^T C V
//   3. Solve the quasi-upper-triangular system R Y + Y S = C' column block by column block (Y = U^T X V)
//   4. Recover: X = U Y V^T
//
// Complexity: O(n^3) (n = max(rows(A), rows(B)))
// Solvability condition: spec(A) ∩ spec(-B) = ∅ (A and -B share no eigenvalue)
//
// Uses: coordinate transforms in control systems, foundation of the Lyapunov equation, matrix function computation

#ifndef SANGI_SYLVESTER_HPP
#define SANGI_SYLVESTER_HPP

#include <math/core/common.hpp>
#include <math/core/matrix.hpp>
#include <math/core/traits.hpp>
#include <math/linalg/decomposition.hpp>
#include <math/linalg/solvers.hpp>
#include <cassert>
#include <cstddef>
#include <limits>
#include <vector>

namespace sangi {
    namespace algorithms {

        namespace detail_sylvester {

            // Detect the block structure of the quasi-upper-triangular matrix R (n×n)
            template<typename T>
            void detect_blocks(const BaseMatrix<T>& R,
                               std::vector<std::size_t>& starts,
                               std::vector<std::size_t>& sizes) {
                starts.clear();
                sizes.clear();
                const auto n = R.rows();
                std::size_t i = 0;
                while (i < n) {
                    if (i + 1 < n && R(i + 1, i) != T(0)) {
                        starts.push_back(i);
                        sizes.push_back(2);
                        i += 2;
                    } else {
                        starts.push_back(i);
                        sizes.push_back(1);
                        i += 1;
                    }
                }
            }

            // Small-block Sylvester: R_b * X + X * S_b = D
            //   R_b: p×p (diagonal block of R), S_b: q×q (diagonal block of S)
            //   X, D: p×q, p,q ∈ {1,2}
            // Directly solve the Kronecker form (I_q ⊗ R_b + S_b^T ⊗ I_p) vec(X) = vec(D)
            template<typename T>
            void solve_small_block(
                const BaseMatrix<T>& R, std::size_t ri, std::size_t p,
                const BaseMatrix<T>& S, std::size_t ci, std::size_t q,
                const T D[2][2],
                T X[2][2])
            {
                const std::size_t m = p * q;
                T M[4][4] = {};
                T rhs[4];

                // Column-major vec: vec(X)[a + p*b] = X[a][b]
                // (I_q ⊗ R_b): M[a + p*b][a' + p*b] += R_b[a][a']  (same b)
                // (S_b^T ⊗ I_p): M[a + p*b][a + p*b'] += S_b[b'][b] (same a, (b',b) -> S_b[b'][b])
                for (std::size_t b = 0; b < q; ++b) {
                    for (std::size_t a = 0; a < p; ++a) {
                        const std::size_t row = a + p * b;
                        // R_b term
                        for (std::size_t ap = 0; ap < p; ++ap) {
                            const std::size_t col = ap + p * b;
                            M[row][col] += R(ri + a, ri + ap);
                        }
                        // S_b^T term (S_b[b'][b] is the (b', b) entry)
                        for (std::size_t bp = 0; bp < q; ++bp) {
                            const std::size_t col = a + p * bp;
                            M[row][col] += S(ci + bp, ci + b);
                        }
                        rhs[row] = D[a][b];
                    }
                }

                // Gaussian elimination (partial pivoting)
                int piv[4] = {0, 1, 2, 3};
                for (std::size_t col = 0; col < m; ++col) {
                    std::size_t max_row = col;
                    T max_val = std::abs(M[piv[col]][col]);
                    for (std::size_t r = col + 1; r < m; ++r) {
                        T v = std::abs(M[piv[r]][col]);
                        if (v > max_val) { max_val = v; max_row = r; }
                    }
                    std::swap(piv[col], piv[max_row]);
                    if (max_val < std::numeric_limits<T>::epsilon() * T(100)) {
                        // Near-singular: close to spec(A) ∩ spec(-B) ≠ ∅
                        for (std::size_t a = 0; a < p; ++a)
                            for (std::size_t b = 0; b < q; ++b)
                                X[a][b] = T(0);
                        return;
                    }
                    for (std::size_t r = col + 1; r < m; ++r) {
                        T f = M[piv[r]][col] / M[piv[col]][col];
                        for (std::size_t k = col + 1; k < m; ++k)
                            M[piv[r]][k] -= f * M[piv[col]][k];
                        rhs[piv[r]] -= f * rhs[piv[col]];
                    }
                }
                T x[4];
                for (int r = static_cast<int>(m) - 1; r >= 0; --r) {
                    T s = rhs[piv[r]];
                    for (std::size_t k = r + 1; k < m; ++k)
                        s -= M[piv[r]][k] * x[k];
                    x[r] = s / M[piv[r]][r];
                }
                for (std::size_t b = 0; b < q; ++b)
                    for (std::size_t a = 0; a < p; ++a)
                        X[a][b] = x[a + p * b];
            }

            // Quasi-upper-triangular system R Y + Y S = D (in-place D -> Y)
            // R is n_r×n_r quasi-upper-triangular, S is n_s×n_s quasi-upper-triangular, D is n_r×n_s
            template<typename T>
            void solve_quasi_triangular(
                const BaseMatrix<T>& R, const BaseMatrix<T>& S, Matrix<T>& D)
            {
                const std::size_t nr = R.rows();
                const std::size_t ns = S.rows();

                std::vector<std::size_t> r_starts, r_sizes;
                std::vector<std::size_t> s_starts, s_sizes;
                detect_blocks(R, r_starts, r_sizes);
                detect_blocks(S, s_starts, s_sizes);

                // Process the column blocks from left to right
                for (std::size_t jb = 0; jb < s_starts.size(); ++jb) {
                    const std::size_t cj = s_starts[jb];
                    const std::size_t q = s_sizes[jb];

                    // Subtract the contribution of the already-solved left column blocks from the RHS
                    // D[:, cj:cj+q] -= Y[:, 0:cj] * S[0:cj, cj:cj+q]
                    for (std::size_t i = 0; i < nr; ++i) {
                        for (std::size_t b = 0; b < q; ++b) {
                            T s = T(0);
                            for (std::size_t k = 0; k < cj; ++k) {
                                s += D(i, k) * S(k, cj + b);
                            }
                            D(i, cj + b) -= s;
                        }
                    }

                    // Process the row blocks from the bottom up (back substitution since R is quasi-upper-triangular)
                    for (std::size_t ib_idx = r_starts.size(); ib_idx > 0; --ib_idx) {
                        const std::size_t ib = ib_idx - 1;
                        const std::size_t ri = r_starts[ib];
                        const std::size_t p = r_sizes[ib];

                        // Subtract the contribution of the already-solved lower row blocks from the RHS
                        // D[ri:ri+p, cj:cj+q] -= R[ri:ri+p, ri+p:] * Y[ri+p:, cj:cj+q]
                        for (std::size_t a = 0; a < p; ++a) {
                            for (std::size_t b = 0; b < q; ++b) {
                                T s = T(0);
                                for (std::size_t k = ri + p; k < nr; ++k) {
                                    s += R(ri + a, k) * D(k, cj + b);
                                }
                                D(ri + a, cj + b) -= s;
                            }
                        }

                        // Solve the small-block Sylvester
                        T Db[2][2];
                        for (std::size_t a = 0; a < p; ++a)
                            for (std::size_t b = 0; b < q; ++b)
                                Db[a][b] = D(ri + a, cj + b);
                        T Xb[2][2];
                        solve_small_block(R, ri, p, S, cj, q, Db, Xb);
                        for (std::size_t a = 0; a < p; ++a)
                            for (std::size_t b = 0; b < q; ++b)
                                D(ri + a, cj + b) = Xb[a][b];
                    }
                }
            }

        } // namespace detail_sylvester

        /**
         * @brief Solves the Sylvester equation A*X + X*B = C
         *
         * Bartels-Stewart method (real Schur decomposition + quasi-upper-triangular back substitution).
         *
         * @tparam T element type (float, double)
         * @param A m×m square matrix
         * @param B n×n square matrix
         * @param C m×n right-hand-side matrix
         * @return X m×n solution matrix
         * @throws DimensionError dimension mismatch
         *
         * @note Solvability condition: spec(A) ∩ spec(-B) = ∅
         *       If there are shared eigenvalues, near-singular diagonal blocks are replaced with a zero solution.
         */
        // Kronecker-expansion based Sylvester solver (for exact types = closure within Q).
        // Expand A X + X B = C into vec form:
        //   (I_n ⊗ A + B^T ⊗ I_m) vec(X) = vec(C)
        //   ↑ solve the (mn)×(mn) linear system via LU.
        // Cost: O((mn)^3), slower than Bartels-Stewart (O(n^3)), but requires no sqrt and
        // produces exact solutions even for exact types such as sangi::Rational. For small-to-medium matrices.
        //
        // Solvability condition: spec(A) ∩ spec(-B) = ∅ (same as Bartels-Stewart).
        template<typename T>
        Matrix<T> solve_sylvester_kronecker(const BaseMatrix<T>& A,
                                             const BaseMatrix<T>& B,
                                             const BaseMatrix<T>& C) {
            const auto m = A.rows();
            const auto n = B.rows();
            if (A.cols() != m) throw DimensionError("solve_sylvester_kronecker: A must be square");
            if (B.cols() != n) throw DimensionError("solve_sylvester_kronecker: B must be square");
            if (C.rows() != m || C.cols() != n)
                throw DimensionError("solve_sylvester_kronecker: C must be m×n");
            if (m == 0 || n == 0) return Matrix<T>(m, n);

            // Expand into Kronecker form
            // vec(X)[a + m*b] = X(a, b) (column-major vec)
            // (I_n ⊗ A) is block-diagonal A, entries: M[a+mb][a'+mb] = A(a, a')
            // (B^T ⊗ I_m) expands B(b',b) onto I_m, entries: M[a+mb][a+mb'] = B(b', b)
            const std::size_t mn = m * n;
            Matrix<T> M(mn, mn, numeric_traits<T>::zero());
            Vector<T> rhs(mn);

            for (std::size_t b = 0; b < n; ++b) {
                for (std::size_t a = 0; a < m; ++a) {
                    const std::size_t row = a + m * b;
                    // I_n ⊗ A term
                    for (std::size_t ap = 0; ap < m; ++ap) {
                        const std::size_t col = ap + m * b;
                        M(row, col) = M(row, col) + A(a, ap);
                    }
                    // B^T ⊗ I_m term
                    for (std::size_t bp = 0; bp < n; ++bp) {
                        const std::size_t col = a + m * bp;
                        M(row, col) = M(row, col) + B(bp, b);
                    }
                    rhs[row] = C(a, b);
                }
            }

            // Solve via LU (exact for exact types)
            Vector<T> x = lu_solve(M, rhs);

            // unvec
            Matrix<T> X(m, n);
            for (std::size_t b = 0; b < n; ++b)
                for (std::size_t a = 0; a < m; ++a)
                    X(a, b) = x[a + m * b];
            return X;
        }

        template<typename T>
        Matrix<T> solve_sylvester(const BaseMatrix<T>& A,
                                  const BaseMatrix<T>& B,
                                  const BaseMatrix<T>& C_basemat)
        {
            // Exact types (Rational, etc.) dispatch to the Kronecker expansion
            if constexpr (!numeric_traits<T>::is_floating_point) {
                return solve_sylvester_kronecker(A, B, C_basemat);
            } else {
            Matrix<T> C(C_basemat);  // 1 copy: body uses Matrix * C * Matrix
            const auto m = A.rows();
            const auto n = B.rows();
            if (A.cols() != m) {
                assert(false && "DimensionError: solve_sylvester: A must be square");
                throw DimensionError("solve_sylvester: A must be square");
            }
            if (B.cols() != n) {
                assert(false && "DimensionError: solve_sylvester: B must be square");
                throw DimensionError("solve_sylvester: B must be square");
            }
            if (C.rows() != m || C.cols() != n) {
                assert(false && "DimensionError: solve_sylvester: C must be m×n");
                throw DimensionError("solve_sylvester: C must be m×n");
            }
            if (m == 0 || n == 0) {
                return Matrix<T>(m, n);
            }

            // Step 1: Schur decomposition A = U R U^T, B = V S V^T
            auto [R, U] = schur_decomposition(A);
            auto [S, V] = schur_decomposition(B);

            // Step 2: C' = U^T C V
            Matrix<T> Ut = U.transpose();
            Matrix<T> Cp = Ut * C * V;

            // Step 3: solve the quasi-upper-triangular system R Y + Y S = C' (in-place, Cp -> Y)
            detail_sylvester::solve_quasi_triangular(R, S, Cp);

            // Step 4: X = U Y V^T
            Matrix<T> Vt = V.transpose();
            return U * Cp * Vt;
            }
        }

    } // namespace algorithms
} // namespace sangi

#endif // SANGI_SYLVESTER_HPP