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

// matrix_functions.hpp
//
// Matrix functions: expm (matrix exponential), sqrtm (matrix square root), logm (matrix logarithm)
//
// Algorithms:
//   expm  — Scaling-and-squaring + Padé [13/13] (Higham 2005)
//   sqrtm — Schur decomposition + Parlett recursion (Björck-Hammarling 1983)
//   logm  — Schur decomposition + Parlett recursion

#ifndef SANGI_MATRIX_FUNCTIONS_HPP
#define SANGI_MATRIX_FUNCTIONS_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 <algorithm>
#include <cmath>
#include <cstddef>
#include <limits>
#include <vector>

namespace sangi {
    namespace algorithms {

        namespace detail_matfun {
            // T → int conversion (for sangi::Float, demote to int via toDouble)
            // Used in expm's scaling parameter s = ceil(log2(norm/theta)) etc.
            template<typename T>
            inline int float_to_int(const T& x) {
                if constexpr (std::is_floating_point_v<T>) {
                    return static_cast<int>(x);
                } else {
                    // sangi::Float: demote to double via .toDouble(), then to int
                    return static_cast<int>(x.toDouble());
                }
            }
        }


        //=====================================================================
        // Helper functions (internal)
        //=====================================================================

        namespace detail_matfun {

            // Matrix 1-norm: max_j Σ_i |A_{ij}| (maximum absolute column sum)
            template<typename T>
            T matrix_one_norm(const BaseMatrix<T>& A) {
                const auto rows = A.rows();
                const auto cols = A.cols();
                T max_col_sum = T(0);
                for (std::size_t j = 0; j < cols; ++j) {
                    T col_sum = T(0);
                    for (std::size_t i = 0; i < rows; ++i) {
                        col_sum += std::abs(A(i, j));
                    }
                    if (col_sum > max_col_sum) {
                        max_col_sum = col_sum;
                    }
                }
                return max_col_sum;
            }

            // Detect the block structure of a quasi-upper-triangular matrix
            // block_starts[k] = starting row of the k-th block
            // block_sizes[k]  = 1 or 2
            template<typename T>
            void detect_blocks(const BaseMatrix<T>& T_mat, std::size_t n,
                               std::vector<std::size_t>& block_starts,
                               std::vector<std::size_t>& block_sizes) {
                block_starts.clear();
                block_sizes.clear();
                std::size_t i = 0;
                while (i < n) {
                    if (i + 1 < n && T_mat(i + 1, i) != T(0)) {
                        // 2×2 block
                        block_starts.push_back(i);
                        block_sizes.push_back(2);
                        i += 2;
                    } else {
                        // 1×1 block
                        block_starts.push_back(i);
                        block_sizes.push_back(1);
                        i += 1;
                    }
                }
            }

            // Compute the eigenvalues α±iβ of a 2×2 block B
            // B = [[a,b],[c,d]], α = (a+d)/2, β = sqrt(|bc + ((a-d)/2)²|)
            // Only for the complex-conjugate-pair case (disc < 0)
            template<typename T>
            void block2_eigenvalues(const BaseMatrix<T>& T_mat, std::size_t i,
                                     T& alpha, T& beta) {
                T a = T_mat(i, i), b = T_mat(i, i + 1);
                T c = T_mat(i + 1, i), d = T_mat(i + 1, i + 1);
                alpha = (a + d) / T(2);
                T half_diff = (a - d) / T(2);
                T disc = half_diff * half_diff + b * c;
                if (disc < T(0)) {
                    beta = std::sqrt(-disc);
                } else {
                    // Real-eigenvalue case (α ± sqrt(disc))
                    beta = T(0);
                }
            }

            // Compute f(B) for a 2×2 block B
            // f(B) = u·I + (v/β)·(B - α·I)
            // func_re, func_im: f(α+iβ) = func_re + i·func_im
            template<typename T>
            void apply_block2_function(
                const BaseMatrix<T>& T_mat, std::size_t i,
                Matrix<T>& F, T alpha, T beta,
                T func_re, T func_im) {
                if (beta == T(0)) {
                    // Real-eigenvalue case: a 2×2 block but β=0
                    // Set f(α) on the diagonal elements
                    F(i, i) = func_re;
                    F(i + 1, i + 1) = func_re;
                    // Approximate the upper-triangular element by f'(α) * T_ij
                    // (in practice this case is rare)
                    F(i, i + 1) = T(0);
                    F(i + 1, i) = T(0);
                    return;
                }
                T u = func_re;
                T v_over_beta = func_im / beta;
                // Set F's (i,i), (i,i+1), (i+1,i), (i+1,i+1)
                F(i, i)         = u + v_over_beta * (T_mat(i, i)     - alpha);
                F(i, i + 1)     =     v_over_beta *  T_mat(i, i + 1);
                F(i + 1, i)     =     v_over_beta *  T_mat(i + 1, i);
                F(i + 1, i + 1) = u + v_over_beta * (T_mat(i + 1, i + 1) - alpha);
            }

            // Small-block Sylvester equation solver: A·X - X·B = C
            // A: p×p, B: q×q, C: p×q, X: p×q (p,q ∈ {1,2})
            // Writes the solution into the submatrix F[ri:ri+p, ci:ci+q] of F
            template<typename T>
            void solve_sylvester_block(
                const BaseMatrix<T>& F,      // F matrix (source for reading A_block and B_block)
                const BaseMatrix<T>& T_mat,  // original quasi-upper-triangular matrix
                Matrix<T>& result,       // destination for writing the solution
                std::size_t ri, std::size_t p,  // row block (start, size)
                std::size_t ci, std::size_t q,  // column block (start, size)
                std::size_t n) {
                // Compute the RHS: C = (F_ii · T_ij - T_ij · F_jj) + Σ_k (F_ik · T_kj - T_ik · F_kj)
                // where F_ii = result[ri:ri+p, ri:ri+p] (diagonal block)
                //         F_jj = result[ci:ci+q, ci:ci+q] (diagonal block)
                //         T_ij = T_mat[ri:ri+p, ci:ci+q] (upper-triangular block)

                // First compute the T_ij · (F_ii - F_jj) term
                // Then add the cross terms from the intermediate blocks

                // C = F_ii · T_ij
                T C[2][2] = {};
                for (std::size_t a = 0; a < p; ++a) {
                    for (std::size_t b = 0; b < q; ++b) {
                        T sum = T(0);
                        for (std::size_t k = 0; k < p; ++k) {
                            sum += result(ri + a, ri + k) * T_mat(ri + k, ci + b);
                        }
                        C[a][b] = sum;
                    }
                }

                // C -= T_ij · F_jj
                for (std::size_t a = 0; a < p; ++a) {
                    for (std::size_t b = 0; b < q; ++b) {
                        T sum = T(0);
                        for (std::size_t k = 0; k < q; ++k) {
                            sum += T_mat(ri + a, ci + k) * result(ci + k, ci + b);
                        }
                        C[a][b] -= sum;
                    }
                }

                // Add the cross terms from the intermediate blocks
                // Σ_{k: ri+p ≤ k < ci} (F[ri:,k:] · T[k:,ci:] - T[ri:,k:] · F[k:,ci:])
                for (std::size_t k = ri + p; k < ci; ++k) {
                    for (std::size_t a = 0; a < p; ++a) {
                        for (std::size_t b = 0; b < q; ++b) {
                            C[a][b] += result(ri + a, k) * T_mat(k, ci + b)
                                     - T_mat(ri + a, k) * result(k, ci + b);
                        }
                    }
                }

                // Solve the Sylvester equation A·X - X·B = C
                // A = T_mat[ri:ri+p, ri:ri+p] (diagonal block of T, not F_ii!)
                // Note: deriving from the relation TF=FT gives
                //   T_ii · F_ij - F_ij · T_jj = RHS
                // so A = T_ii, B = T_jj

                if (p == 1 && q == 1) {
                    // Scalar: x = c / (t_ii - t_jj)
                    T denom = T_mat(ri, ri) - T_mat(ci, ci);
                    if (std::abs(denom) < std::numeric_limits<T>::epsilon() * T(100)) {
                        result(ri, ci) = T(0); // degenerate case
                    } else {
                        result(ri, ci) = C[0][0] / denom;
                    }
                } else if (p == 1 && q == 2) {
                    // 1×2 Sylvester: a·X - X·B = C
                    // [a*x0 - x0*b00 - x1*b10, a*x1 - x0*b01 - x1*b11] = [c0, c1]
                    T a = T_mat(ri, ri);
                    T b00 = T_mat(ci, ci), b01 = T_mat(ci, ci + 1);
                    T b10 = T_mat(ci + 1, ci), b11 = T_mat(ci + 1, ci + 1);
                    // M · [x0, x1]^T = [c0, c1]^T
                    T m00 = a - b00, m01 = -b10;
                    T m10 = -b01,    m11 = a - b11;
                    T det = m00 * m11 - m01 * m10;
                    if (std::abs(det) < std::numeric_limits<T>::epsilon() * T(100)) {
                        result(ri, ci) = T(0);
                        result(ri, ci + 1) = T(0);
                    } else {
                        result(ri, ci)     = ( m11 * C[0][0] - m01 * C[0][1]) / det;
                        result(ri, ci + 1) = (-m10 * C[0][0] + m00 * C[0][1]) / det;
                    }
                } else if (p == 2 && q == 1) {
                    // 2×1 Sylvester: A·X - X·b = C
                    T b = T_mat(ci, ci);
                    T a00 = T_mat(ri, ri),     a01 = T_mat(ri, ri + 1);
                    T a10 = T_mat(ri + 1, ri), a11 = T_mat(ri + 1, ri + 1);
                    // [a00-b, a01; a10, a11-b] · [x0; x1] = [c0; c1]
                    T m00 = a00 - b, m01 = a01;
                    T m10 = a10,     m11 = a11 - b;
                    T det = m00 * m11 - m01 * m10;
                    if (std::abs(det) < std::numeric_limits<T>::epsilon() * T(100)) {
                        result(ri, ci) = T(0);
                        result(ri + 1, ci) = T(0);
                    } else {
                        result(ri, ci)     = ( m11 * C[0][0] - m01 * C[1][0]) / det;
                        result(ri + 1, ci) = (-m10 * C[0][0] + m00 * C[1][0]) / det;
                    }
                } else {
                    // 2×2 Sylvester: A·X - X·B = C → 4×4 linear system
                    T a00 = T_mat(ri, ri),     a01 = T_mat(ri, ri + 1);
                    T a10 = T_mat(ri + 1, ri), a11 = T_mat(ri + 1, ri + 1);
                    T b00 = T_mat(ci, ci),     b01 = T_mat(ci, ci + 1);
                    T b10 = T_mat(ci + 1, ci), b11 = T_mat(ci + 1, ci + 1);

                    // vec(X) = [x00, x10, x01, x11]^T (column-major)
                    // M · vec(X) = vec(C) where M = I⊗A - B^T⊗I
                    // M = [[a00-b00, a01,    -b10,    0    ],
                    //      [a10,     a11-b00, 0,     -b10  ],
                    //      [-b01,    0,       a00-b11, a01  ],
                    //      [0,      -b01,     a10,    a11-b11]]
                    T M[4][4] = {
                        {a00 - b00, a01,       -b10,      T(0)},
                        {a10,       a11 - b00,  T(0),     -b10},
                        {-b01,      T(0),       a00 - b11, a01},
                        {T(0),     -b01,        a10,       a11 - b11}
                    };
                    T rhs[4] = {C[0][0], C[1][0], C[0][1], C[1][1]};

                    // 4×4 Gaussian elimination (partial pivoting)
                    int piv[4] = {0, 1, 2, 3};
                    for (int col = 0; col < 4; ++col) {
                        // Pivot selection
                        int max_row = col;
                        T max_val = std::abs(M[piv[col]][col]);
                        for (int row = col + 1; row < 4; ++row) {
                            T v = std::abs(M[piv[row]][col]);
                            if (v > max_val) {
                                max_val = v;
                                max_row = row;
                            }
                        }
                        std::swap(piv[col], piv[max_row]);

                        if (max_val < std::numeric_limits<T>::epsilon() * T(100)) {
                            // Singular case
                            result(ri, ci) = T(0);
                            result(ri + 1, ci) = T(0);
                            result(ri, ci + 1) = T(0);
                            result(ri + 1, ci + 1) = T(0);
                            return;
                        }

                        // Elimination
                        for (int row = col + 1; row < 4; ++row) {
                            T factor = M[piv[row]][col] / M[piv[col]][col];
                            for (int k = col + 1; k < 4; ++k) {
                                M[piv[row]][k] -= factor * M[piv[col]][k];
                            }
                            rhs[piv[row]] -= factor * rhs[piv[col]];
                        }
                    }

                    // Back substitution
                    T x[4];
                    for (int i = 3; i >= 0; --i) {
                        T sum = rhs[piv[i]];
                        for (int k = i + 1; k < 4; ++k) {
                            sum -= M[piv[i]][k] * x[k];
                        }
                        x[i] = sum / M[piv[i]][i];
                    }

                    result(ri, ci)         = x[0];
                    result(ri + 1, ci)     = x[1];
                    result(ri, ci + 1)     = x[2];
                    result(ri + 1, ci + 1) = x[3];
                }
            }

        } // namespace detail_matfun

        //=====================================================================
        // expm — matrix exponential
        //=====================================================================

        // expm via the Schur-Parlett method for arbitrary-precision T (= sangi::Float etc.).
        // Based on Higham (2008) "Functions of Matrices" §10.4.
        //   1. Schur decomposition A = Q T Q^T (T is quasi-upper-triangular)
        //   2. Compute exp(T_ii) directly on the diagonal blocks (1×1 / 2×2 complex conjugate)
        //   3. Fill in the upper-triangular off-diagonal via Parlett iteration
        //   4. exp(A) = Q exp(T) Q^T
        // Advantage: avoids cancellation inside Padé, no precision loss at arbitrary precision.
        // Drawback: O(n³) cost of the Schur decomposition (large constant); Parlett iteration breaks down on repeated eigenvalues
        //       (precision degrades when eigenvalues are closely spaced; a separate block reorder would be needed but
        //        this is accepted as the same scope of limitation as sqrtm/logm).
        template<typename T>
        Matrix<T> expm_schur_parlett(const BaseMatrix<T>& A_basemat) {
            const auto n = A_basemat.rows();
            if (n == 0) return Matrix<T>(A_basemat);
            if (n == 1) {
                Matrix<T> result(1, 1);
                result(0, 0) = std::exp(A_basemat(0, 0));
                return result;
            }

            // Schur decomposition A = Q T Q^T
            auto [T_mat, Q] = schur_decomposition(A_basemat);

            // Detect the block structure
            std::vector<std::size_t> block_starts, block_sizes;
            detail_matfun::detect_blocks(T_mat, n, block_starts, block_sizes);

            // Initialize the result matrix F = exp(T)
            Matrix<T> F(n, n, T(0));

            // Step 1: apply exp directly to the diagonal blocks
            for (std::size_t bi = 0; bi < block_starts.size(); ++bi) {
                std::size_t i = block_starts[bi];
                std::size_t sz = block_sizes[bi];

                if (sz == 1) {
                    // 1×1 block: exp(T_ii) — high precision via sangi::exp
                    F(i, i) = std::exp(T_mat(i, i));
                } else {
                    // 2×2 block: complex-conjugate eigenvalues α±iβ
                    // exp(α+iβ) = exp(α)·cos(β) + i·exp(α)·sin(β)
                    T alpha, beta;
                    detail_matfun::block2_eigenvalues(T_mat, i, alpha, beta);

                    T ea = std::exp(alpha);
                    T u = ea * std::cos(beta);  // real part of exp(α+iβ)
                    T v = ea * std::sin(beta);  // imaginary part

                    detail_matfun::apply_block2_function(T_mat, i, F, alpha, beta, u, v);
                }
            }

            // Step 2: compute the super-diagonal blocks via Parlett iteration
            for (std::size_t d = 1; d < block_starts.size(); ++d) {
                for (std::size_t bi = 0; bi + d < block_starts.size(); ++bi) {
                    std::size_t bj = bi + d;
                    std::size_t ri = block_starts[bi];
                    std::size_t p  = block_sizes[bi];
                    std::size_t ci = block_starts[bj];
                    std::size_t q  = block_sizes[bj];

                    detail_matfun::solve_sylvester_block(F, T_mat, F, ri, p, ci, q, n);
                }
            }

            // Step 3: recover exp(A) = Q·F·Q^T
            Matrix<T> Qt = Q.transpose();
            return Q * F * Qt;
        }

        /**
         * @brief Matrix exponential exp(A)
         *
         * Scaling-and-squaring method + Padé [13/13] approximation (float/double).
         * Based on Higham (2005) "The Scaling and Squaring Method for
         * the Matrix Exponential Revisited".
         *
         * For arbitrary-precision types (= sangi::Float etc.), dispatches to the Schur-Parlett method:
         *   - avoids the V±U cancellation of Padé
         *   - can deliver precision according to target_precision
         *   - cost: Schur decomposition O(n³) (Hessenberg + QR iteration)
         *
         * @tparam T element type (float, double, sangi::Float)
         * @param A input matrix (n×n square matrix)
         * @return exp(A)
         * @throws DimensionError if A is not square
         */
        template<typename T>
        Matrix<T> expm(const BaseMatrix<T>& A_basemat) {
            // Arbitrary-precision types dispatch to Schur-Parlett (suited for high precision)
            if constexpr (!std::is_floating_point_v<T>) {
                return expm_schur_parlett(A_basemat);
            }

            Matrix<T> A(A_basemat);  // 1 copy via BaseMatrix ctor
            const auto n = A.rows();
            if (n != A.cols()) {
                assert(false && "DimensionError: expm: matrix must be square");
                throw DimensionError("expm: matrix must be square");
            }
            if (n == 0) {
                return A;
            }
            if (n == 1) {
                Matrix<T> result(1, 1);
                result(0, 0) = std::exp(A(0, 0));
                return result;
            }

            auto I = Matrix<T>::identity(n);

            // θ values for the Padé approximation (Higham Table 10.2)
            const T theta3  = T(1.495585217958292e-02);
            const T theta5  = T(2.539398330063230e-01);
            const T theta7  = T(9.504178996162932e-01);
            const T theta9  = T(2.097847961257068e+00);
            const T theta13 = T(5.371920351148152e+00);

            T norm1 = detail_matfun::matrix_one_norm(A);

            // Precompute A² (used at all orders)
            Matrix<T> A2 = A * A;

            // When a low-order Padé can be used (||A||₁ is sufficiently small)
            auto pade_3 = [&]() -> Matrix<T> {
                // b₀=120, b₁=60, b₂=12, b₃=1
                const T b0 = T(120), b1 = T(60), b2 = T(12), b3 = T(1);
                Matrix<T> U = A * (b3 * A2 + b1 * I);
                Matrix<T> V = b2 * A2 + b0 * I;
                return solve_multi(Matrix<T>(V - U), Matrix<T>(V + U));
            };

            auto pade_5 = [&]() -> Matrix<T> {
                const T b0 = T(30240), b1 = T(15120), b2 = T(3360);
                const T b3 = T(420), b4 = T(30), b5 = T(1);
                Matrix<T> A4 = A2 * A2;
                Matrix<T> U = A * (b5 * A4 + b3 * A2 + b1 * I);
                Matrix<T> V = b4 * A4 + b2 * A2 + b0 * I;
                return solve_multi(Matrix<T>(V - U), Matrix<T>(V + U));
            };

            auto pade_7 = [&]() -> Matrix<T> {
                const T b0 = T(17297280), b1 = T(8648640), b2 = T(1995840);
                const T b3 = T(277200), b4 = T(25200), b5 = T(1512);
                const T b6 = T(56), b7 = T(1);
                Matrix<T> A4 = A2 * A2;
                Matrix<T> A6 = A4 * A2;
                Matrix<T> U = A * (b7 * A6 + b5 * A4 + b3 * A2 + b1 * I);
                Matrix<T> V = b6 * A6 + b4 * A4 + b2 * A2 + b0 * I;
                return solve_multi(Matrix<T>(V - U), Matrix<T>(V + U));
            };

            auto pade_9 = [&]() -> Matrix<T> {
                const T b0 = T(17643225600.0), b1 = T(8821612800.0);
                const T b2 = T(2075673600.0), b3 = T(302702400.0);
                const T b4 = T(30270240.0), b5 = T(2162160.0);
                const T b6 = T(110880.0), b7 = T(3960.0);
                const T b8 = T(90.0), b9 = T(1);
                Matrix<T> A4 = A2 * A2;
                Matrix<T> A6 = A4 * A2;
                Matrix<T> A8 = A4 * A4;
                Matrix<T> U = A * (b9 * A8 + b7 * A6 + b5 * A4 + b3 * A2 + b1 * I);
                Matrix<T> V = b8 * A8 + b6 * A6 + b4 * A4 + b2 * A2 + b0 * I;
                return solve_multi(Matrix<T>(V - U), Matrix<T>(V + U));
            };

            // Arbitrary-precision types already dispatched to Schur-Parlett at the top of the function. The following is float/double only.
            if (norm1 <= theta3) return pade_3();
            if (norm1 <= theta5) return pade_5();
            if (norm1 <= theta7) return pade_7();
            if (norm1 <= theta9) return pade_9();

            // Padé [13/13] with scaling (theta13 = 5.37 as the target norm)
            int s = std::max(0, detail_matfun::float_to_int(std::ceil(std::log2(norm1 / theta13))));
            T scale = std::ldexp(T(1), -s);
            Matrix<T> As = A * scale;
            Matrix<T> As2 = As * As;
            Matrix<T> As4 = As2 * As2;
            Matrix<T> As6 = As4 * As2;

            // Padé [13/13] coefficients
            const T b0  = T(64764752532480000.0);
            const T b1  = T(32382376266240000.0);
            const T b2  = T(7771770303897600.0);
            const T b3  = T(1187353796428800.0);
            const T b4  = T(129060195264000.0);
            const T b5  = T(10559470521600.0);
            const T b6  = T(670442572800.0);
            const T b7  = T(33522128640.0);
            const T b8  = T(1323241920.0);
            const T b9  = T(40840800.0);
            const T b10 = T(960960.0);
            const T b11 = T(16380.0);
            const T b12 = T(182.0);
            const T b13 = T(1.0);

            // U = As · (As6·(b13·As6 + b11·As4 + b9·As2) + b7·As6 + b5·As4 + b3·As2 + b1·I)
            Matrix<T> inner_u = b13 * As6 + b11 * As4 + b9 * As2;
            Matrix<T> U = As * (As6 * inner_u + b7 * As6 + b5 * As4 + b3 * As2 + b1 * I);

            // V = As6·(b12·As6 + b10·As4 + b8·As2) + b6·As6 + b4·As4 + b2·As2 + b0·I
            Matrix<T> inner_v = b12 * As6 + b10 * As4 + b8 * As2;
            Matrix<T> V = As6 * inner_v + b6 * As6 + b4 * As4 + b2 * As2 + b0 * I;

            // Solve (V - U) · F = V + U
            Matrix<T> F = solve_multi(Matrix<T>(V - U), Matrix<T>(V + U));

            // Square s times
            for (int i = 0; i < s; ++i) {
                F = F * F;
            }

            return F;
        }

        //=====================================================================
        // sqrtm — matrix square root
        //=====================================================================

        /**
         * @brief Matrix square root A^{1/2} (principal value)
         *
         * Schur decomposition + Parlett recursion (Björck-Hammarling 1983).
         * The principal square root exists when no eigenvalue of A lies on the negative real axis.
         *
         * @tparam T element type (float, double)
         * @param A input matrix (n×n square matrix)
         * @return A^{1/2}
         * @throws DimensionError if A is not square
         * @throws MathError if a negative real eigenvalue exists
         */
        template<typename T>
        Matrix<T> sqrtm(const BaseMatrix<T>& A) {
            const auto n = A.rows();
            if (n != A.cols()) {
                assert(false && "DimensionError: sqrtm: matrix must be square");
                throw DimensionError("sqrtm: matrix must be square");
            }
            if (n == 0) {
                return Matrix<T>(A);
            }
            if (n == 1) {
                if (A(0, 0) < T(0)) {
                    throw MathError("sqrtm: matrix has negative real eigenvalue");
                }
                Matrix<T> result(1, 1);
                result(0, 0) = std::sqrt(A(0, 0));
                return result;
            }

            // Schur decomposition: A = Q·T·Qᵀ
            auto [T_mat, Q] = schur_decomposition(A);

            // Detect the block structure
            std::vector<std::size_t> block_starts, block_sizes;
            detail_matfun::detect_blocks(T_mat, n, block_starts, block_sizes);

            // Initialize the result matrix F = sqrtm(T)
            Matrix<T> F(n, n, T(0));

            // Step 1: compute the diagonal blocks
            for (std::size_t bi = 0; bi < block_starts.size(); ++bi) {
                std::size_t i = block_starts[bi];
                std::size_t sz = block_sizes[bi];

                if (sz == 1) {
                    // 1×1 block: sqrt(T_ii)
                    if (T_mat(i, i) < T(0)) {
                        throw MathError("sqrtm: matrix has negative real eigenvalue");
                    }
                    F(i, i) = std::sqrt(T_mat(i, i));
                } else {
                    // 2×2 block: complex-conjugate eigenvalues α±iβ
                    // √(α+iβ) = √r · (cos(θ/2) + i·sin(θ/2))
                    // r = |λ| = √(α²+β²), θ = arg(λ) = atan2(β, α)
                    T alpha, beta;
                    detail_matfun::block2_eigenvalues(T_mat, i, alpha, beta);

                    T r = std::sqrt(alpha * alpha + beta * beta);
                    T theta = std::atan2(beta, alpha);
                    T sqrt_r = std::sqrt(r);
                    T u = sqrt_r * std::cos(theta / T(2));
                    T v = sqrt_r * std::sin(theta / T(2));

                    detail_matfun::apply_block2_function(T_mat, i, F, alpha, beta, u, v);
                }
            }

            // Step 2: compute the super-diagonal blocks (Parlett recursion)
            // Process block columns moving away from the diagonal
            for (std::size_t d = 1; d < block_starts.size(); ++d) {
                for (std::size_t bi = 0; bi + d < block_starts.size(); ++bi) {
                    std::size_t bj = bi + d;
                    std::size_t ri = block_starts[bi];
                    std::size_t p  = block_sizes[bi];
                    std::size_t ci = block_starts[bj];
                    std::size_t q  = block_sizes[bj];

                    detail_matfun::solve_sylvester_block(F, T_mat, F, ri, p, ci, q, n);
                }
            }

            // Step 3: recover sqrtm(A) = Q·F·Qᵀ
            Matrix<T> Qt = Q.transpose();
            return Q * F * Qt;
        }

        //=====================================================================
        // logm — matrix logarithm
        //=====================================================================

        /**
         * @brief Matrix logarithm log(A) (principal value)
         *
         * Schur decomposition + Parlett recursion.
         * The principal logarithm exists when no eigenvalue of A lies on the negative real axis.
         *
         * @tparam T element type (float, double)
         * @param A input matrix (n×n square matrix)
         * @return log(A)
         * @throws DimensionError if A is not square
         * @throws MathError if a non-positive real eigenvalue exists
         */
        template<typename T>
        Matrix<T> logm(const BaseMatrix<T>& A) {
            const auto n = A.rows();
            if (n != A.cols()) {
                assert(false && "DimensionError: logm: matrix must be square");
                throw DimensionError("logm: matrix must be square");
            }
            if (n == 0) {
                return Matrix<T>(A);
            }
            if (n == 1) {
                if (A(0, 0) <= T(0)) {
                    throw MathError("logm: matrix has non-positive real eigenvalue");
                }
                Matrix<T> result(1, 1);
                result(0, 0) = std::log(A(0, 0));
                return result;
            }

            // Schur decomposition: A = Q·T·Qᵀ
            auto [T_mat, Q] = schur_decomposition(A);

            // Detect the block structure
            std::vector<std::size_t> block_starts, block_sizes;
            detail_matfun::detect_blocks(T_mat, n, block_starts, block_sizes);

            // Initialize the result matrix F = logm(T)
            Matrix<T> F(n, n, T(0));

            // Step 1: compute the diagonal blocks
            for (std::size_t bi = 0; bi < block_starts.size(); ++bi) {
                std::size_t i = block_starts[bi];
                std::size_t sz = block_sizes[bi];

                if (sz == 1) {
                    // 1×1 block: log(T_ii)
                    if (T_mat(i, i) <= T(0)) {
                        throw MathError("logm: matrix has non-positive real eigenvalue");
                    }
                    F(i, i) = std::log(T_mat(i, i));
                } else {
                    // 2×2 block: complex-conjugate eigenvalues α±iβ
                    T alpha, beta;
                    detail_matfun::block2_eigenvalues(T_mat, i, alpha, beta);

                    T r2 = alpha * alpha + beta * beta;
                    T u = std::log(r2) / T(2); // ½·log(|λ|²)
                    T v = std::atan2(beta, alpha);

                    detail_matfun::apply_block2_function(T_mat, i, F, alpha, beta, u, v);
                }
            }

            // Step 2: compute the super-diagonal blocks (Parlett recursion)
            for (std::size_t d = 1; d < block_starts.size(); ++d) {
                for (std::size_t bi = 0; bi + d < block_starts.size(); ++bi) {
                    std::size_t bj = bi + d;
                    std::size_t ri = block_starts[bi];
                    std::size_t p  = block_sizes[bi];
                    std::size_t ci = block_starts[bj];
                    std::size_t q  = block_sizes[bj];

                    detail_matfun::solve_sylvester_block(F, T_mat, F, ri, p, ci, q, n);
                }
            }

            // Step 3: recover logm(A) = Q·F·Qᵀ
            Matrix<T> Qt = Q.transpose();
            return Q * F * Qt;
        }

        // ================================================================
        // Frechet derivative (block-triangular method)
        // ================================================================
        //
        // The Frechet derivative L_f(A)[E] of a matrix function f satisfies:
        //   f(A + εE) = f(A) + ε·L_f(A)[E] + O(ε²)
        //
        // Block-triangular method:
        //   M = [A, E; 0, A]  (2n × 2n)
        //   f(M) = [f(A), L_f(A)[E]; 0, f(A)]
        //   → L_f(A)[E] = top-right n×n block of f(M)
        //
        // Reference: Higham (2008) "Functions of Matrices", Chapter 3

        /// Frechet derivative of expm: L_exp(A)[E]
        /// @param A  n×n matrix
        /// @param E  n×n perturbation direction
        /// @return   L_exp(A)[E]  (n×n matrix)
        template<typename T>
        [[nodiscard]]
        Matrix<T> frechet_expm(const BaseMatrix<T>& A, const BaseMatrix<T>& E) {
            const auto n = A.rows();
            if (n != A.cols() || n != E.rows() || n != E.cols())
                throw DimensionError("frechet_expm: A and E must be square and same size");

            // M = [A, E; 0, A]
            Matrix<T> M(2*n, 2*n, T(0));
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    M(i, j) = A(i, j);
                    M(i, j + n) = E(i, j);
                    M(i + n, j + n) = A(i, j);
                }

            auto fM = expm(M);

            // Extract top-right block
            Matrix<T> L(n, n);
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    L(i, j) = fM(i, j + n);
            return L;
        }

        /// Frechet derivative of logm: L_log(A)[E]
        template<typename T>
        [[nodiscard]]
        Matrix<T> frechet_logm(const BaseMatrix<T>& A, const BaseMatrix<T>& E) {
            const auto n = A.rows();
            if (n != A.cols() || n != E.rows() || n != E.cols())
                throw DimensionError("frechet_logm: A and E must be square and same size");

            Matrix<T> M(2*n, 2*n, T(0));
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    M(i, j) = A(i, j);
                    M(i, j + n) = E(i, j);
                    M(i + n, j + n) = A(i, j);
                }

            auto fM = logm(M);

            Matrix<T> L(n, n);
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    L(i, j) = fM(i, j + n);
            return L;
        }

        /// Frechet derivative of sqrtm: L_sqrt(A)[E]
        template<typename T>
        [[nodiscard]]
        Matrix<T> frechet_sqrtm(const BaseMatrix<T>& A, const BaseMatrix<T>& E) {
            const auto n = A.rows();
            if (n != A.cols() || n != E.rows() || n != E.cols())
                throw DimensionError("frechet_sqrtm: A and E must be square and same size");

            Matrix<T> M(2*n, 2*n, T(0));
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    M(i, j) = A(i, j);
                    M(i, j + n) = E(i, j);
                    M(i + n, j + n) = A(i, j);
                }

            auto fM = sqrtm(M);

            Matrix<T> L(n, n);
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    L(i, j) = fM(i, j + n);
            return L;
        }

        /// Generic Frechet derivative: compute L_f(A)[E] for any matrix function f
        /// @param f  matrix function f: Matrix<T> → Matrix<T>
        /// @param A  n×n matrix
        /// @param E  n×n perturbation direction
        /// @return   L_f(A)[E]  (n×n matrix)
        template<typename T, typename MatFunc>
        [[nodiscard]]
        Matrix<T> frechet(MatFunc f, const BaseMatrix<T>& A, const BaseMatrix<T>& E) {
            const auto n = A.rows();
            if (n != A.cols() || n != E.rows() || n != E.cols())
                throw DimensionError("frechet: A and E must be square and same size");

            Matrix<T> M(2*n, 2*n, T(0));
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    M(i, j) = A(i, j);
                    M(i, j + n) = E(i, j);
                    M(i + n, j + n) = A(i, j);
                }

            auto fM = f(M);

            Matrix<T> L(n, n);
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    L(i, j) = fM(i, j + n);
            return L;
        }

        //=====================================================================
        // Differentials of decompositions (Decomposition Differentials)
        //=====================================================================

        // ================================================================
        // diff_svd_sigma — singular-value differential of the SVD
        // ================================================================
        //
        // For A = U · diag(σ) · Vt, compute dσ from dA.
        //   dσ_i = u_i^T · dA · v_i
        //
        // Note: dU, dV are omitted as they are unstable with repeated singular values.
        //       The singular-values-only differential is stable and fast.

        /// Singular-value differential of the SVD: dσ_i = u_i^T · dA · v_i
        /// @param U   left singular vector matrix (m×m or m×k)
        /// @param sigma  singular value vector (k entries)
        /// @param Vt  transpose of the right singular vector matrix (n×n or k×n)
        /// @param dA  perturbation matrix (m×n)
        /// @return dσ vector (k entries)
        template<typename T>
        [[nodiscard]]
        Vector<T> diff_svd_sigma(const BaseMatrix<T>& U, const BaseVector<T>& sigma,
                                 const BaseMatrix<T>& Vt, const BaseMatrix<T>& dA) {
            const auto m = dA.rows();
            const auto n = dA.cols();
            const auto k = sigma.size();

            if (U.rows() != m || Vt.cols() != n)
                throw DimensionError("diff_svd_sigma: dimension mismatch between U/Vt and dA");

            // dA_rotated = U^T · dA · V = U^T · dA · Vt^T
            // dσ_i = (U^T · dA · V)_{ii}
            // Since V = Vt^T, dA · V = dA · Vt^T
            auto V = Vt.transposed();  // 0-copy view

            Vector<T> dsigma(k);
            for (std::size_t i = 0; i < k; ++i) {
                // Compute u_i^T · dA · v_i as a dot product
                T sum = T(0);
                for (std::size_t r = 0; r < m; ++r) {
                    T dA_vi_r = T(0);
                    for (std::size_t c = 0; c < n; ++c) {
                        dA_vi_r += dA(r, c) * V(c, i);
                    }
                    sum += U(r, i) * dA_vi_r;
                }
                dsigma[i] = sum;
            }
            return dsigma;
        }

        // ================================================================
        // diff_symmetric_eigen — differential of the symmetric eigendecomposition
        // ================================================================
        //
        // For A = V · diag(λ) · V^T, compute dλ, dV from dA.
        //   dλ_i = v_i^T · dA · v_i
        //   (dV)_{:,j} += v_i · (v_i^T · dA · v_j) / (λ_j - λ_i)  for i ≠ j
        //
        // Note: dV is unstable when λ_i ≈ λ_j (repeated eigenvalues).
        //       If λ_j - λ_i is at or below the machine epsilon, set that component to 0.

        /// Differential of the symmetric eigendecomposition
        /// @param eigenvalues  eigenvalue vector (n entries)
        /// @param eigenvectors eigenvector matrix (n×n, columns are eigenvectors)
        /// @param dA  perturbation matrix (n×n symmetric matrix)
        /// @return the pair {dλ, dV}
        template<typename T>
        [[nodiscard]]
        std::pair<Vector<T>, Matrix<T>> diff_symmetric_eigen(
            const BaseVector<T>& eigenvalues_basevec, const Matrix<T>& eigenvectors,
            const Matrix<T>& dA) {
            Vector<T> eigenvalues(eigenvalues_basevec);  // 1 copy via BaseVector ctor
            const auto n = eigenvalues.size();

            if (eigenvectors.rows() != n || eigenvectors.cols() != n)
                throw DimensionError("diff_symmetric_eigen: eigenvectors must be n x n");
            if (dA.rows() != n || dA.cols() != n)
                throw DimensionError("diff_symmetric_eigen: dA must be n x n");

            // Compute V^T · dA · V
            Matrix<T> Vt = eigenvectors.transpose();
            Matrix<T> VtdAV = Vt * dA * eigenvectors;

            // dλ_i = (V^T · dA · V)_{ii}
            Vector<T> dlambda(n);
            for (std::size_t i = 0; i < n; ++i) {
                dlambda[i] = VtdAV(i, i);
            }

            // dV = V · Ω  where Ω_{ij} = VtdAV_{ij} / (λ_j - λ_i)  (i ≠ j), Ω_{ii} = 0
            const T eps = std::numeric_limits<T>::epsilon() * T(100);
            Matrix<T> Omega(n, n, T(0));
            for (std::size_t i = 0; i < n; ++i) {
                for (std::size_t j = 0; j < n; ++j) {
                    if (i != j) {
                        T diff = eigenvalues[j] - eigenvalues[i];
                        if (std::abs(diff) > eps) {
                            Omega(i, j) = VtdAV(i, j) / diff;
                        }
                        // Leave as 0 for repeated eigenvalues
                    }
                }
            }

            Matrix<T> dV = eigenvectors * Omega;
            return { dlambda, dV };
        }

        // ================================================================
        // diff_cholesky — differential of the Cholesky decomposition
        // ================================================================
        //
        // For A = L · L^T, compute dA → dL.
        //   S = L^{-1} · dA · L^{-T}
        //   dL = L · tril(S) where the diagonal elements are S_{ii}/2
        //
        // Reference: Giles (2008) "An extended collection of matrix derivative results
        //        for forward and reverse mode automatic differentiation"

        /// Differential of the Cholesky decomposition
        /// @param L   Cholesky factor (lower-triangular matrix, n×n)
        /// @param dA  perturbation matrix (n×n symmetric matrix)
        /// @return dL (lower-triangular matrix, n×n)
        template<typename T>
        [[nodiscard]]
        Matrix<T> diff_cholesky(const BaseMatrix<T>& L, const BaseMatrix<T>& dA) {
            const auto n = L.rows();

            if (L.cols() != n)
                throw DimensionError("diff_cholesky: L must be square");
            if (dA.rows() != n || dA.cols() != n)
                throw DimensionError("diff_cholesky: dA must be n x n");

            // Since cholesky_decomposition may leave original values in the upper-triangular part,
            // copy and use only the lower-triangular part
            Matrix<T> Lc(n, n, T(0));
            for (std::size_t i = 0; i < n; ++i)
                for (std::size_t j = 0; j <= i; ++j)
                    Lc(i, j) = L(i, j);

            // Direct method: solve the Lyapunov equation L · dL^T + dL · L^T = dA
            // dL is lower-triangular. Compute column by column in the order (j, j), (j+1, j), ..., (n-1, j).
            //
            // Diagonal (i = j):
            //   dA_{jj} = 2 L_{jj} dL_{jj} + 2 Σ_{k<j} L_{jk} dL_{jk}
            //   → dL_{jj} = (dA_{jj} - 2 Σ_{k<j} L_{jk} dL_{jk}) / (2 L_{jj})
            //
            // Lower-triangular (i > j):
            //   dA_{ij} = dL_{ij} L_{jj} + Σ_{k<j} dL_{ik} L_{jk} + Σ_{k<=j} L_{ik} dL_{jk}
            //   → dL_{ij} = (dA_{ij} - Σ_{k<j} dL_{ik} L_{jk} - Σ_{k<=j} L_{ik} dL_{jk}) / L_{jj}

            Matrix<T> dL(n, n, T(0));

            for (std::size_t j = 0; j < n; ++j) {
                // Diagonal element dL_{jj}
                T sum_diag = T(0);
                for (std::size_t k = 0; k < j; ++k) {
                    sum_diag += Lc(j, k) * dL(j, k);
                }
                dL(j, j) = (dA(j, j) - T(2) * sum_diag) / (T(2) * Lc(j, j));

                // Off-diagonal element dL_{ij} (i > j)
                for (std::size_t i = j + 1; i < n; ++i) {
                    T sum = T(0);
                    // Σ_{k<j} dL_{ik} L_{jk}
                    for (std::size_t k = 0; k < j; ++k) {
                        sum += dL(i, k) * Lc(j, k);
                    }
                    // Σ_{k<=j} L_{ik} dL_{jk}
                    for (std::size_t k = 0; k <= j; ++k) {
                        sum += Lc(i, k) * dL(j, k);
                    }
                    dL(i, j) = (dA(i, j) - sum) / Lc(j, j);
                }
            }

            return dL;
        }

    } // namespace algorithms
} // namespace sangi

#endif // SANGI_MATRIX_FUNCTIONS_HPP