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

// computation_policy.hpp
#pragma once

#include <cstddef>
#include <cstring>
#include <type_traits>
#include <algorithm>
#include <future>
#include <math/core/numeric_state.hpp>
#include <math/core/computation/simd_backend.hpp>
#include <math/core/simd_traits.hpp>
#include <math/core/mp/ThreadPool.hpp>
#include <math/core/computation/mkl_backend.hpp>

namespace sangi {

    // Base class for computation policies
    template <typename T>
    class ComputePolicyBase {
    public:
        using value_type = T;
        using size_type = std::size_t;
    };

    // Forward declarations
    template <typename T> class DefaultComputePolicy;
    template <typename T> class SIMDComputePolicy;
    template <typename T> class MKLComputePolicy;

    // Default computation policy
    template <typename T>
    class DefaultComputePolicy : public ComputePolicyBase<T> {
    public:
        using value_type = T;
        using size_type = std::size_t;

        // Vector addition
        template <typename VecA, typename VecB, typename VecResult>
        static void vector_add(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for addition");
            }

            for (size_type i = 0; i < a.size(); ++i) {
                result[i] = a[i] + b[i];

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(a[i]) ||
                        numeric_state_traits<T>::isSpecial(b[i])) {
                        // NumericState propagation
                    }
                }
            }
        }

        // Vector subtraction
        template <typename VecA, typename VecB, typename VecResult>
        static void vector_subtract(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for subtraction");
            }

            for (size_type i = 0; i < a.size(); ++i) {
                result[i] = a[i] - b[i];

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(a[i]) ||
                        numeric_state_traits<T>::isSpecial(b[i])) {
                        // NumericState propagation
                    }
                }
            }
        }

        // Vector dot product
        template <typename VecA, typename VecB>
        static T dot(const VecA& a, const VecB& b) {
            if (a.size() != b.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for dot product");
            }

            T result = static_cast<T>(0);
            for (size_type i = 0; i < a.size(); ++i) {
                result += a[i] * b[i];

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(a[i]) ||
                        numeric_state_traits<T>::isSpecial(b[i])) {
                        // NumericState propagation
                    }
                }
            }

            return result;
        }

        // Vector norm (L2 norm)
        template <typename Vec>
        static T norm(const Vec& vec) {
            T sum_of_squares = static_cast<T>(0);

            for (size_type i = 0; i < vec.size(); ++i) {
                sum_of_squares += vec[i] * vec[i];

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(vec[i])) {
                        // NumericState propagation
                    }
                }
            }

            return std::sqrt(sum_of_squares);
        }

        // Vector-scalar multiplication
        template <typename Vec, typename VecResult>
        static void vector_scalar_multiply(const Vec& vec, const T& scalar, VecResult& result) {
            if (vec.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for scalar multiplication");
            }

            for (size_type i = 0; i < vec.size(); ++i) {
                result[i] = vec[i] * scalar;

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(vec[i]) ||
                        numeric_state_traits<T>::isSpecial(scalar)) {
                        // NumericState propagation
                    }
                }
            }
        }

        // Element-wise vector multiplication
        template <typename VecA, typename VecB, typename VecResult>
        static void elementWiseMultiply(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for element-wise multiplication");
            }

            for (size_type i = 0; i < a.size(); ++i) {
                result[i] = a[i] * b[i];

                // Apply the NumericState system
                if (numeric_state_traits<T>::is_supported) {
                    if (numeric_state_traits<T>::isSpecial(a[i]) ||
                        numeric_state_traits<T>::isSpecial(b[i])) {
                        // NumericState propagation
                    }
                }
            }
        }

        // Element-wise vector division
        template <typename VecA, typename VecB, typename VecResult>
        static void elementWiseDivide(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for element-wise division");
            }

            for (size_type i = 0; i < a.size(); ++i) {
                if (b[i] == static_cast<T>(0)) {
                    // Handle division by zero (set a special value or throw an exception)
                    if (numeric_state_traits<T>::is_supported && numeric_state_traits<T>::isSpecial(a[i])) {
                        // Set a special state
                    }
                    else {
                        throw std::invalid_argument("Division by zero");
                    }
                }
                else {
                    result[i] = a[i] / b[i];

                    // Apply the NumericState system
                    if (numeric_state_traits<T>::is_supported) {
                        if (numeric_state_traits<T>::isSpecial(a[i]) ||
                            numeric_state_traits<T>::isSpecial(b[i])) {
                            // NumericState propagation
                        }
                    }
                }
            }
        }

        // Functions for matrix operations

        // Matrix addition
        template <typename MatA, typename MatB, typename MatResult>
        static void add(const MatA& a, const MatB& b, MatResult& result) {
            if (a.rows() != b.rows() || a.cols() != b.cols() ||
                a.rows() != result.rows() || a.cols() != result.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for addition");
            }

            for (size_type i = 0; i < a.rows(); ++i) {
                for (size_type j = 0; j < a.cols(); ++j) {
                    result(i, j) = a(i, j) + b(i, j);
                }
            }
        }

        // Matrix subtraction
        template <typename MatA, typename MatB, typename MatResult>
        static void subtract(const MatA& a, const MatB& b, MatResult& result) {
            if (a.rows() != b.rows() || a.cols() != b.cols() ||
                a.rows() != result.rows() || a.cols() != result.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for subtraction");
            }

            for (size_type i = 0; i < a.rows(); ++i) {
                for (size_type j = 0; j < a.cols(); ++j) {
                    result(i, j) = a(i, j) - b(i, j);
                }
            }
        }

        // Matrix multiplication: C = A * B
        template <typename MatA, typename MatB, typename MatResult>
        static void multiply(const MatA& a, const MatB& b, MatResult& result) {
            if (a.cols() != b.rows() || result.rows() != a.rows() || result.cols() != b.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for multiplication");
            }

            const size_type M = a.rows();
            const size_type N = b.cols();
            const size_type K = a.cols();

            // Raw pointer access (assumes row-major layout)
            const T* a_data = a.data();
            const T* b_data = b.data();
            T* c_data = result.data();
            const size_type lda = a.cols();
            const size_type ldb = b.cols();
            const size_type ldc = result.cols();

#if SANGI_HAS_CBLAS
            // Dispatch to cblas_?gemm when MKL is available and T is float/double
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                // MKL's gemm is column-major, so we convert from row-major:
                //   C = A*B (row-major) <=> C^T = B^T * A^T (column-major)
                // i.e. viewed as column-major: C(N,M) = B(N,K)^T_row * A(K,M)^T_row
                // We pass row-major data as-is and use the transpose flag:
                //   cblas_?gemm(RowMajor, NoTrans, NoTrans, M, N, K, 1, A, lda, B, ldb, 0, C, ldc)
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0f, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0f, c_data, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0, c_data, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            // Zero-initialize
            const T zero_val = numeric_traits<T>::zero();
            for (size_type i = 0; i < M * N; ++i)
                c_data[i] = zero_val;

            gemm_blocked(a_data, b_data, c_data, M, N, K, lda, ldb, ldc);
        }

        // Matrix multiplication (B-transposed): C = A * B^T
        // A is (M, K), B is (N, K) -- B^T is not materialized; computed directly using cblas TransB=true
        // Used for output = input * weight^T in linear layers
        template <typename MatA, typename MatB, typename MatResult>
        static void multiply_trans_b(const MatA& a, const MatB& b, MatResult& result) {
            // A: M×K, B: N×K (pre-transpose), C: M×N
            if (a.cols() != b.cols() || result.rows() != a.rows() || result.cols() != b.rows()) {
                throw std::invalid_argument("Matrix dimensions mismatch for multiply_trans_b");
            }

            const size_type M = a.rows();
            const size_type N = b.rows();
            const size_type K = a.cols();

            const T* a_data = a.data();
            const T* b_data = b.data();
            T* c_data = result.data();
            const size_type lda = a.cols();  // A row stride = K
            const size_type ldb = b.cols();  // B row stride = K
            const size_type ldc = result.cols();  // C row stride = N

#if SANGI_HAS_CBLAS
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                // cblas_?gemm(RowMajor, NoTrans, Trans, M, N, K, 1, A, K, B, K, 0, C, N)
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0f, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0f, c_data, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0, c_data, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            // Fallback: naively form B^T and delegate to multiply
            // (Poor performance but ensures correctness)
            const T zero_val = numeric_traits<T>::zero();
            for (size_type i = 0; i < M; ++i) {
                for (size_type j = 0; j < N; ++j) {
                    T sum = zero_val;
                    for (size_type k = 0; k < K; ++k)
                        sum += a_data[i * lda + k] * b_data[j * ldb + k];
                    c_data[i * ldc + j] = sum;
                }
            }
        }

        // Matrix multiplication (B-transposed, accumulating): C = alpha * A * B^T + beta * C
        // Allows fusing the linear bias broadcast + GEMM into a single cblas call
        template <typename MatA, typename MatB, typename MatResult>
        static void multiply_trans_b_accumulate(const MatA& a, const MatB& b, MatResult& result,
                                                 T alpha, T beta) {
            if (a.cols() != b.cols() || result.rows() != a.rows() || result.cols() != b.rows()) {
                throw std::invalid_argument("Matrix dimensions mismatch for multiply_trans_b_accumulate");
            }

            const size_type M = a.rows();
            const size_type N = b.rows();
            const size_type K = a.cols();

            const T* a_data = a.data();
            const T* b_data = b.data();
            T* c_data = result.data();
            const size_type lda = a.cols();
            const size_type ldb = b.cols();
            const size_type ldc = result.cols();

#if SANGI_HAS_CBLAS
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                beta, c_data, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                beta, c_data, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            // Fallback: naively compute A * B^T, then form beta*C + alpha*result
            for (size_type i = 0; i < M; ++i) {
                for (size_type j = 0; j < N; ++j) {
                    T sum = T{0};
                    for (size_type k = 0; k < K; ++k)
                        sum += a_data[i * lda + k] * b_data[j * ldb + k];
                    c_data[i * ldc + j] = alpha * sum + beta * c_data[i * ldc + j];
                }
            }
        }

        // Matrix multiplication (A-transposed): C = A^T * B
        // A is (K, M), B is (K, N) -- A^T is not materialized; computed directly using cblas TransA=true
        // Used for gradWeight = gradOutput^T * input in linear-layer backward
        template <typename MatA, typename MatB, typename MatResult>
        static void multiply_trans_a(const MatA& a, const MatB& b, MatResult& result) {
            // A: K×M, B: K×N, C: M×N
            if (a.rows() != b.rows() || result.rows() != a.cols() || result.cols() != b.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for multiply_trans_a");
            }

            const size_type M = a.cols();
            const size_type N = b.cols();
            const size_type K = a.rows();

            const T* a_data = a.data();
            const T* b_data = b.data();
            T* c_data = result.data();
            const size_type lda = a.cols();  // A row stride = M
            const size_type ldb = b.cols();  // B row stride = N
            const size_type ldc = result.cols();  // C row stride = N

#if SANGI_HAS_CBLAS
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0f, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0f, c_data, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                1.0, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                0.0, c_data, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            // Fallback
            const T zero_val = numeric_traits<T>::zero();
            for (size_type i = 0; i < M; ++i) {
                for (size_type j = 0; j < N; ++j) {
                    T sum = zero_val;
                    for (size_type k = 0; k < K; ++k)
                        sum += a_data[k * lda + i] * b_data[k * ldb + j];
                    c_data[i * ldc + j] = sum;
                }
            }
        }

        // fused GEMM: C = alpha * A * B + beta * C
        template <typename MatA, typename MatB, typename MatResult>
        static void gemm(const MatA& a, const MatB& b, MatResult& result,
                         T alpha = T{1}, T beta = T{0}) {
            if (a.cols() != b.rows() || result.rows() != a.rows() || result.cols() != b.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for GEMM");
            }

            const size_type M = a.rows();
            const size_type N = b.cols();
            const size_type K = a.cols();

            const T* a_data = a.data();
            const T* b_data = b.data();
            T* c_data = result.data();
            const size_type lda = a.cols();
            const size_type ldb = b.cols();
            const size_type ldc = result.cols();

#if SANGI_HAS_CBLAS
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                beta, c_data, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a_data, static_cast<sangi::blas_int_t>(lda),
                                b_data, static_cast<sangi::blas_int_t>(ldb),
                                beta, c_data, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            // Software GEMM: C = alpha * A * B + beta * C
            gemm_fused(a_data, b_data, c_data, M, N, K, lda, ldb, ldc, alpha, beta);
        }

    private:
        // ================================================================
        // Blocked gemm: tiling + register blocking + B packing
        // ================================================================
        //
        // MR x NR micro-kernel: hold C in YMM registers and accumulate FMA across the k loop
        // -> C load/store drops from one per k to one overall (a BK-fold bandwidth reduction)
        //
        // B packing: rearrange the B panel into NR-wide column panels
        // -> converts strided access along k (jumping by ldb) into contiguous access

        static constexpr size_type MR = 4;   // Register block row count

        // NR: register block column count (two SIMD packets)
        static constexpr size_type gemm_NR() {
            if constexpr (has_simd_packet_v<T>)
                return PacketTraits<T>::size * 2;  // double: 8, float: 16
            else
                return 4;
        }

        static constexpr size_type BM = 64;    // Row block
        static constexpr size_type BN = 64;    // Column block (multiple of NR)
        static constexpr size_type BK = 128;   // Inner-product block (can be enlarged via register blocking)

        // B packing: rearrange B[k0:k0+bk, j0:j0+bn] into NR-wide column panels
        // Output layout: the k-th entry of panel p -> packB[p * bk * NR + k * NR + 0..NR-1]
        static void pack_b_panel(const T* b, T* packB,
                                  size_type k0, size_type bk,
                                  size_type j0, size_type bn,
                                  size_type ldb) {
            const size_type NR = gemm_NR();
            const size_type n_full = bn / NR;
            // Full panels
            for (size_type p = 0; p < n_full; ++p) {
                T* dst = packB + p * bk * NR;
                for (size_type k = 0; k < bk; ++k) {
                    const T* src = &b[(k0 + k) * ldb + j0 + p * NR];
                    std::memcpy(dst + k * NR, src, NR * sizeof(T));
                }
            }
            // Remainder panel (less than NR) -- zero-padded
            const size_type rem = bn % NR;
            if (rem > 0) {
                T* dst = packB + n_full * bk * NR;
                for (size_type k = 0; k < bk; ++k) {
                    const T* src = &b[(k0 + k) * ldb + j0 + n_full * NR];
                    size_type r = 0;
                    for (; r < rem; ++r)
                        dst[k * NR + r] = src[r];
                    for (; r < NR; ++r)
                        dst[k * NR + r] = T{0};
                }
            }
        }

        // Register-block micro-kernel (packed-B variant): C[i:i+4, 0:NR] += A * packB
        // C is held in 8 YMM registers and accumulated across the entire k loop
        static void gemm_micro_kernel(const T* a, const T* packB_panel, T* c,
                                       size_type i, size_type bk,
                                       size_type lda, size_type ldc) {
            if constexpr (has_simd_packet_v<T>) {
                using PT = PacketTraits<T>;
                constexpr size_type PS = PT::size;
                constexpr size_type NR = PS * 2;

                auto c00 = PT::load(&c[(i+0)*ldc]);
                auto c01 = PT::load(&c[(i+0)*ldc + PS]);
                auto c10 = PT::load(&c[(i+1)*ldc]);
                auto c11 = PT::load(&c[(i+1)*ldc + PS]);
                auto c20 = PT::load(&c[(i+2)*ldc]);
                auto c21 = PT::load(&c[(i+2)*ldc + PS]);
                auto c30 = PT::load(&c[(i+3)*ldc]);
                auto c31 = PT::load(&c[(i+3)*ldc + PS]);

                for (size_type k = 0; k < bk; ++k) {
                    const auto b0 = PT::load(&packB_panel[k * NR]);
                    const auto b1 = PT::load(&packB_panel[k * NR + PS]);

                    const auto a0 = PT::set1(a[(i+0)*lda + k]);
                    c00 = PT::fmadd(a0, b0, c00);
                    c01 = PT::fmadd(a0, b1, c01);

                    const auto a1 = PT::set1(a[(i+1)*lda + k]);
                    c10 = PT::fmadd(a1, b0, c10);
                    c11 = PT::fmadd(a1, b1, c11);

                    const auto a2 = PT::set1(a[(i+2)*lda + k]);
                    c20 = PT::fmadd(a2, b0, c20);
                    c21 = PT::fmadd(a2, b1, c21);

                    const auto a3 = PT::set1(a[(i+3)*lda + k]);
                    c30 = PT::fmadd(a3, b0, c30);
                    c31 = PT::fmadd(a3, b1, c31);
                }

                PT::store(&c[(i+0)*ldc], c00);
                PT::store(&c[(i+0)*ldc + PS], c01);
                PT::store(&c[(i+1)*ldc], c10);
                PT::store(&c[(i+1)*ldc + PS], c11);
                PT::store(&c[(i+2)*ldc], c20);
                PT::store(&c[(i+2)*ldc + PS], c21);
                PT::store(&c[(i+3)*ldc], c30);
                PT::store(&c[(i+3)*ldc + PS], c31);
            }
        }

        // Register-block micro-kernel (direct-B variant): no packing
        // Avoids packing overhead on small matrices
        static void gemm_micro_kernel_direct(const T* a, const T* b, T* c,
                                              size_type i, size_type j,
                                              size_type k0, size_type bk,
                                              size_type lda, size_type ldb, size_type ldc) {
            if constexpr (has_simd_packet_v<T>) {
                using PT = PacketTraits<T>;
                constexpr size_type PS = PT::size;

                auto c00 = PT::load(&c[(i+0)*ldc + j]);
                auto c01 = PT::load(&c[(i+0)*ldc + j + PS]);
                auto c10 = PT::load(&c[(i+1)*ldc + j]);
                auto c11 = PT::load(&c[(i+1)*ldc + j + PS]);
                auto c20 = PT::load(&c[(i+2)*ldc + j]);
                auto c21 = PT::load(&c[(i+2)*ldc + j + PS]);
                auto c30 = PT::load(&c[(i+3)*ldc + j]);
                auto c31 = PT::load(&c[(i+3)*ldc + j + PS]);

                for (size_type k = k0; k < k0 + bk; ++k) {
                    const auto b0 = PT::load(&b[k * ldb + j]);
                    const auto b1 = PT::load(&b[k * ldb + j + PS]);

                    const auto a0 = PT::set1(a[(i+0)*lda + k]);
                    c00 = PT::fmadd(a0, b0, c00);
                    c01 = PT::fmadd(a0, b1, c01);

                    const auto a1 = PT::set1(a[(i+1)*lda + k]);
                    c10 = PT::fmadd(a1, b0, c10);
                    c11 = PT::fmadd(a1, b1, c11);

                    const auto a2 = PT::set1(a[(i+2)*lda + k]);
                    c20 = PT::fmadd(a2, b0, c20);
                    c21 = PT::fmadd(a2, b1, c21);

                    const auto a3 = PT::set1(a[(i+3)*lda + k]);
                    c30 = PT::fmadd(a3, b0, c30);
                    c31 = PT::fmadd(a3, b1, c31);
                }

                PT::store(&c[(i+0)*ldc + j], c00);
                PT::store(&c[(i+0)*ldc + j + PS], c01);
                PT::store(&c[(i+1)*ldc + j], c10);
                PT::store(&c[(i+1)*ldc + j + PS], c11);
                PT::store(&c[(i+2)*ldc + j], c20);
                PT::store(&c[(i+2)*ldc + j + PS], c21);
                PT::store(&c[(i+3)*ldc + j], c30);
                PT::store(&c[(i+3)*ldc + j + PS], c31);
            }
        }

        // Small-matrix path: no packing, register blocking only
        // Used when the packing overhead is large relative to the compute volume
        static void gemm_direct(const T* a, const T* b, T* c,
                                 size_type M, size_type N, size_type K,
                                 size_type lda, size_type ldb, size_type ldc) {
            if constexpr (has_simd_packet_v<T>) {
                using PT = PacketTraits<T>;
                constexpr size_type PS = PT::size;
                constexpr size_type NR = PS * 2;

                // MR x NR register block (reads B directly)
                size_type i = 0;
                for (; i + MR <= M; i += MR) {
                    size_type j = 0;
                    for (; j + NR <= N; j += NR) {
                        gemm_micro_kernel_direct(a, b, c, i, j, 0, K, lda, ldb, ldc);
                    }
                    // Remainder columns: one packet wide
                    for (; j + PS <= N; j += PS) {
                        auto c00 = PT::load(&c[(i+0)*ldc + j]);
                        auto c10 = PT::load(&c[(i+1)*ldc + j]);
                        auto c20 = PT::load(&c[(i+2)*ldc + j]);
                        auto c30 = PT::load(&c[(i+3)*ldc + j]);
                        for (size_type k = 0; k < K; ++k) {
                            const auto bv = PT::load(&b[k * ldb + j]);
                            c00 = PT::fmadd(PT::set1(a[(i+0)*lda + k]), bv, c00);
                            c10 = PT::fmadd(PT::set1(a[(i+1)*lda + k]), bv, c10);
                            c20 = PT::fmadd(PT::set1(a[(i+2)*lda + k]), bv, c20);
                            c30 = PT::fmadd(PT::set1(a[(i+3)*lda + k]), bv, c30);
                        }
                        PT::store(&c[(i+0)*ldc + j], c00);
                        PT::store(&c[(i+1)*ldc + j], c10);
                        PT::store(&c[(i+2)*ldc + j], c20);
                        PT::store(&c[(i+3)*ldc + j], c30);
                    }
                    // Scalar remainder
                    for (; j < N; ++j) {
                        T s0 = c[(i+0)*ldc+j], s1 = c[(i+1)*ldc+j];
                        T s2 = c[(i+2)*ldc+j], s3 = c[(i+3)*ldc+j];
                        for (size_type k = 0; k < K; ++k) {
                            T bkj = b[k*ldb+j];
                            s0 += a[(i+0)*lda+k] * bkj; s1 += a[(i+1)*lda+k] * bkj;
                            s2 += a[(i+2)*lda+k] * bkj; s3 += a[(i+3)*lda+k] * bkj;
                        }
                        c[(i+0)*ldc+j] = s0; c[(i+1)*ldc+j] = s1;
                        c[(i+2)*ldc+j] = s2; c[(i+3)*ldc+j] = s3;
                    }
                }
                // Remainder rows
                for (; i < M; ++i) {
                    for (size_type k = 0; k < K; ++k) {
                        const auto a_ik = PT::set1(a[i*lda + k]);
                        T a_val = a[i*lda + k];
                        size_type j = 0;
                        const size_type j_vec_end = (N / PS) * PS;
                        for (; j < j_vec_end; j += PS) {
                            auto cv = PT::load(&c[i*ldc + j]);
                            PT::store(&c[i*ldc + j], PT::fmadd(a_ik, PT::load(&b[k*ldb + j]), cv));
                        }
                        for (; j < N; ++j)
                            c[i*ldc + j] += a_val * b[k*ldb + j];
                    }
                }
            } else {
                for (size_type i = 0; i < M; ++i)
                    for (size_type k = 0; k < K; ++k) {
                        T a_val = a[i*lda + k];
                        for (size_type j = 0; j < N; ++j)
                            c[i*ldc + j] += a_val * b[k*ldb + j];
                    }
            }
        }

        // Large-matrix path: B packing + register blocking
        static void gemm_packed(const T* a, const T* b, T* c,
                                 size_type M, size_type N, size_type K,
                                 size_type lda, size_type ldb, size_type ldc) {
            if constexpr (has_simd_packet_v<T>) {
                using PT = PacketTraits<T>;
                constexpr size_type PS = PT::size;
                constexpr size_type NR = PS * 2;

                constexpr size_type PACK_B_SIZE = (BN / NR + 1) * BK * NR;
                alignas(32) T packB[PACK_B_SIZE];

                for (size_type jj = 0; jj < N; jj += BN) {
                    const size_type bn = ((jj + BN) < N) ? BN : (N - jj);
                    for (size_type kk = 0; kk < K; kk += BK) {
                        const size_type bk = ((kk + BK) < K) ? BK : (K - kk);

                        pack_b_panel(b, packB, kk, bk, jj, bn, ldb);

                        for (size_type ii = 0; ii < M; ii += BM) {
                            const size_type iend = ((ii + BM) < M) ? (ii + BM) : M;

                            size_type i = ii;
                            for (; i + MR <= iend; i += MR) {
                                const size_type n_full = bn / NR;
                                for (size_type p = 0; p < n_full; ++p) {
                                    gemm_micro_kernel(
                                        a + kk, packB + p * bk * NR,
                                        c + jj + p * NR,
                                        i, bk, lda, ldc);
                                }
                                const size_type j_rem_start = jj + n_full * NR;
                                for (size_type j = j_rem_start; j < jj + bn; ++j) {
                                    T sum = c[i*ldc+j];      T sum1 = c[(i+1)*ldc+j];
                                    T sum2 = c[(i+2)*ldc+j]; T sum3 = c[(i+3)*ldc+j];
                                    for (size_type k = 0; k < bk; ++k) {
                                        T bkj = b[(kk+k)*ldb+j];
                                        sum  += a[ i     *lda+kk+k] * bkj;
                                        sum1 += a[(i+1)*lda+kk+k] * bkj;
                                        sum2 += a[(i+2)*lda+kk+k] * bkj;
                                        sum3 += a[(i+3)*lda+kk+k] * bkj;
                                    }
                                    c[i*ldc+j] = sum;      c[(i+1)*ldc+j] = sum1;
                                    c[(i+2)*ldc+j] = sum2; c[(i+3)*ldc+j] = sum3;
                                }
                            }
                            for (; i < iend; ++i) {
                                for (size_type k = 0; k < bk; ++k) {
                                    const auto a_ik = PT::set1(a[i*lda+kk+k]);
                                    T a_val = a[i*lda+kk+k];
                                    const T* b_row = &b[(kk+k)*ldb+jj];
                                    T* c_row = &c[i*ldc+jj];
                                    size_type j = 0;
                                    const size_type j_vec_end = (bn / PS) * PS;
                                    for (; j < j_vec_end; j += PS) {
                                        auto cv = PT::load(&c_row[j]);
                                        PT::store(&c_row[j], PT::fmadd(a_ik, PT::load(&b_row[j]), cv));
                                    }
                                    for (; j < bn; ++j)
                                        c_row[j] += a_val * b_row[j];
                                }
                            }
                        }
                    }
                }
            } else {
                for (size_type jj = 0; jj < N; jj += BN) {
                    const size_type jend = ((jj + BN) < N) ? (jj + BN) : N;
                    for (size_type kk = 0; kk < K; kk += BK) {
                        const size_type kend = ((kk + BK) < K) ? (kk + BK) : K;
                        for (size_type ii = 0; ii < M; ii += BM) {
                            const size_type iend = ((ii + BM) < M) ? (ii + BM) : M;
                            for (size_type i = ii; i < iend; ++i)
                                for (size_type k = kk; k < kend; ++k) {
                                    T a_val = a[i*lda+k];
                                    for (size_type j = jj; j < jend; ++j)
                                        c[i*ldc+j] += a_val * b[k*ldb+j];
                                }
                        }
                    }
                }
            }
        }

        // ================================================================
        // Packing decision: M*N*K alone is not sufficient
        //
        // Packing cost    = O(K*N) (copy of the B panel)
        // Compute volume  = O(M*N*K)
        // Amortization    = M / MR (number of micro-kernel calls along rows)
        //
        // -> If M is small, the packing cost is large relative to the compute volume
        // -> If N < 2*NR there is at most one full panel, so packing is less meaningful
        // -> However, when B exceeds L2 (K*N > 64K) cache misses dominate, so
        //    packing is beneficial even for small M
        // ================================================================

        static bool should_pack(size_type M, size_type N, size_type K) {
            const size_type NR = gemm_NR();
            // N less than 2*NR: at most one panel, packing not needed
            if (N < NR * 2) return false;
            // If M is large enough and there is enough compute, pack
            // Below 2M (100^3 = 1M) B fits in L2 and packing is unnecessary
            if (M >= 32 && M * N * K >= 2'000'000) return true;
            // If B is large and does not fit in L2, pack even when M is small
            if (K * N >= 65536) return true;
            return false;
        }

        // ================================================================
        // Parallelism decision: heuristic for splitting along M
        //
        // Thread-pool launch and synchronization overhead ~= 300-400 us
        // -> No benefit from parallelization unless single-thread time exceeds this
        // -> When M is small, rows per thread are insufficient and efficiency drops
        //
        // Benchmark results (128 HW threads):
        //   192^3 (7M):       ratio 1.33 -> PARALLEL wins (crossover)
        //   150^3 (3.4M):     ratio 0.84 -> SINGLE wins
        //   32x500x500 (8M):  ratio 0.63 -> SINGLE (M too small)
        //   64x500x500 (16M): ratio 1.30 -> PARALLEL (M=64 is the lower bound)
        // ================================================================
        static inline bool force_single_thread_ = false;
    public:
        /// Disable GEMM parallelism (for benchmarks)
        static void setSingleThread(bool st) { force_single_thread_ = st; }
    private:
        static bool should_parallel(size_type M, size_type N, size_type K) {
            if (force_single_thread_) return false;
            // M less than 64: rows per thread are insufficient
            if (M < 64) return false;
            // Parallelize only when total compute is at least 6M FLOPs
            if (M * N * K < size_type(6'000'000)) return false;
            // Verify that rows-per-thread x K is sufficient
            // Work per thread = (M/n_threads) x N x K FLOPs
            // To amortize ~50us thread launch cost, ~500K FLOPs/thread is required
            auto hw = std::thread::hardware_concurrency();
            if (hw <= 1) return false;
            size_type rows_per_thread = M / hw;
            if (rows_per_thread < MR) return false;  // Fewer than MR rows yields poor micro-kernel efficiency
            return rows_per_thread * N * K >= size_type(500'000);
        }

        // Parallel gemm: split along M and distribute among threads
        // Each thread owns a distinct row range of C -> no synchronization required
        static void gemm_parallel(const T* a, const T* b, T* c,
                                   size_type M, size_type N, size_type K,
                                   size_type lda, size_type ldb, size_type ldc) {
            auto& pool = threadPool();
            const unsigned hw = std::thread::hardware_concurrency();
            const unsigned n_threads = (hw > 1) ? hw : 1;

            // Split rows in units of MR (register-blocking alignment)
            const size_type rows_per_thread = ((M / n_threads + MR - 1) / MR) * MR;

            std::vector<std::future<void>> futures;
            futures.reserve(n_threads);

            for (unsigned t = 0; t < n_threads; ++t) {
                const size_type i_start = t * rows_per_thread;
                if (i_start >= M) break;
                const size_type i_end = (i_start + rows_per_thread < M)
                                         ? i_start + rows_per_thread : M;
                const size_type sub_M = i_end - i_start;

                futures.push_back(pool.submit([=]() {
                    gemm_packed(a + i_start * lda, b,
                                c + i_start * ldc,
                                sub_M, N, K, lda, ldb, ldc);
                }));
            }

            for (auto& f : futures)
                f.get();
        }

        static void gemm_blocked(const T* a, const T* b, T* c,
                                  size_type M, size_type N, size_type K,
                                  size_type lda, size_type ldb, size_type ldc) {
            if (!should_pack(M, N, K))
                gemm_direct(a, b, c, M, N, K, lda, ldb, ldc);
            else if (!should_parallel(M, N, K))
                gemm_packed(a, b, c, M, N, K, lda, ldb, ldc);
            else
                gemm_parallel(a, b, c, M, N, K, lda, ldb, ldc);
        }

        // Fused GEMM: C = alpha * A * B + beta * C (software implementation)
        //
        // beta=0: zero-initialize C, accumulate A*B, then apply alpha
        // beta!=0: scale C by beta, accumulate A*B, then apply alpha
        // alpha=1, beta=0: equivalent to ordinary matrix multiplication
        static void gemm_fused(const T* a, const T* b, T* c,
                                size_type M, size_type N, size_type K,
                                size_type lda, size_type ldb, size_type ldc,
                                T alpha, T beta) {
            const size_type total = M * N;

            // Pre-process C based on beta
            if (beta == T{0}) {
                const T zero_val = T{0};
                for (size_type i = 0; i < total; ++i)
                    c[i] = zero_val;
            } else if (beta != T{1}) {
                for (size_type i = 0; i < total; ++i)
                    c[i] *= beta;
            }
            // For beta == 1, leave C unchanged (accumulate)

            // Accumulate A * B into C
            gemm_blocked(a, b, c, M, N, K, lda, ldb, ldc);

            // Post-scale by alpha
            if (alpha != T{1}) {
                for (size_type i = 0; i < total; ++i)
                    c[i] *= alpha;
            }
        }

    public:

        // Raw-pointer gemm (C += A * B; C must be zero-initialized beforehand)
        static void gemm(const T* a, const T* b, T* c,
                         size_type M, size_type N, size_type K,
                         size_type lda, size_type ldb, size_type ldc) {
            gemm_blocked(a, b, c, M, N, K, lda, ldb, ldc);
        }

        // Raw-pointer fused GEMM: C = alpha * A * B + beta * C
        static void gemm(const T* a, const T* b, T* c,
                         size_type M, size_type N, size_type K,
                         size_type lda, size_type ldb, size_type ldc,
                         T alpha, T beta) {
#if SANGI_HAS_CBLAS
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if constexpr (std::is_same_v<T, float>) {
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a, static_cast<sangi::blas_int_t>(lda),
                                b, static_cast<sangi::blas_int_t>(ldb),
                                beta, c, static_cast<sangi::blas_int_t>(ldc));
                } else {
                    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                                static_cast<sangi::blas_int_t>(M), static_cast<sangi::blas_int_t>(N), static_cast<sangi::blas_int_t>(K),
                                alpha, a, static_cast<sangi::blas_int_t>(lda),
                                b, static_cast<sangi::blas_int_t>(ldb),
                                beta, c, static_cast<sangi::blas_int_t>(ldc));
                }
                return;
            }
#endif
            gemm_fused(a, b, c, M, N, K, lda, ldb, ldc, alpha, beta);
        }

        // Scalar multiplication
        template <typename Mat, typename MatResult>
        static void scalarMultiply(const Mat& mat, const T& scalar, MatResult& result) {
            if (mat.rows() != result.rows() || mat.cols() != result.cols()) {
                throw std::invalid_argument("Matrix dimensions mismatch for scalar multiplication");
            }

            for (size_type i = 0; i < mat.rows(); ++i) {
                for (size_type j = 0; j < mat.cols(); ++j) {
                    result(i, j) = mat(i, j) * scalar;
                }
            }
        }

        // Transpose (AVX2 micro-kernel + block transpose + large-scale parallelization)
        //
        // double: 4x4 AVX2 kernel (256-bit = 4 doubles)
        // float:  8x8 AVX2 kernel (256-bit = 8 floats)
        // others: scalar block transpose
        // Large matrices (>= 512x512): thread-pool parallelization
        template <typename Mat, typename MatResult>
        static void transpose(const Mat& mat, MatResult& result) {
            if (mat.rows() != result.cols() || mat.cols() != result.rows()) {
                throw std::invalid_argument("Matrix dimensions mismatch for transpose");
            }

            const size_type R = mat.rows();
            const size_type C = mat.cols();

            if constexpr ((std::is_same_v<T, double> || std::is_same_v<T, float>)
                          && has_simd_packet_v<T>) {
                transpose_block_avx2_dispatch(mat.data(), result.data(), R, C);
            } else if constexpr (std::is_arithmetic_v<T>) {
                transpose_block_scalar(mat.data(), result.data(), R, C);
            } else {
                for (size_type i = 0; i < R; ++i)
                    for (size_type j = 0; j < C; ++j)
                        result(j, i) = mat(i, j);
            }
        }

    private:
        // Parallelism threshold: parallelize when element count is at least this large
        static constexpr size_type TRANSPOSE_PARALLEL_THRESHOLD = 512 * 512;

        // Scalar block transpose
        static void transpose_block_scalar(const T* src, T* dst,
                                            size_type R, size_type C) {
            constexpr size_type BS = 32;
            for (size_type ii = 0; ii < R; ii += BS) {
                const size_type iend = (ii + BS < R) ? ii + BS : R;
                for (size_type jj = 0; jj < C; jj += BS) {
                    const size_type jend = (jj + BS < C) ? jj + BS : C;
                    for (size_type i = ii; i < iend; ++i)
                        for (size_type j = jj; j < jend; ++j)
                            dst[j * R + i] = src[i * C + j];
                }
            }
        }

#ifdef __AVX2__
        // =====================================================================
        // AVX2 4×4 double in-register transpose
        // =====================================================================
        static void transpose_4x4d_kernel(const double* src, size_type src_stride,
                                            double* dst, size_type dst_stride) {
            auto r0 = _mm256_loadu_pd(src);
            auto r1 = _mm256_loadu_pd(src + src_stride);
            auto r2 = _mm256_loadu_pd(src + 2 * src_stride);
            auto r3 = _mm256_loadu_pd(src + 3 * src_stride);

            auto t0 = _mm256_unpacklo_pd(r0, r1);
            auto t1 = _mm256_unpackhi_pd(r0, r1);
            auto t2 = _mm256_unpacklo_pd(r2, r3);
            auto t3 = _mm256_unpackhi_pd(r2, r3);

            _mm256_storeu_pd(dst,                    _mm256_permute2f128_pd(t0, t2, 0x20));
            _mm256_storeu_pd(dst + dst_stride,       _mm256_permute2f128_pd(t1, t3, 0x20));
            _mm256_storeu_pd(dst + 2 * dst_stride,   _mm256_permute2f128_pd(t0, t2, 0x31));
            _mm256_storeu_pd(dst + 3 * dst_stride,   _mm256_permute2f128_pd(t1, t3, 0x31));
        }

        // double block transpose (one block strip: rows [ii, iend) x all columns)
        static void transpose_block_avx2_double_strip(
                const double* src, double* dst,
                size_type R, size_type C,
                size_type ii, size_type iend) {
            constexpr size_type BS = 32;
            for (size_type jj = 0; jj < C; jj += BS) {
                const size_type jend = (jj + BS < C) ? jj + BS : C;
                size_type i = ii;
                for (; i + 4 <= iend; i += 4) {
                    size_type j = jj;
                    for (; j + 4 <= jend; j += 4) {
                        transpose_4x4d_kernel(
                            src + i * C + j, C,
                            dst + j * R + i, R);
                    }
                    for (; j < jend; ++j)
                        for (size_type k = i; k < i + 4; ++k)
                            dst[j * R + k] = src[k * C + j];
                }
                for (; i < iend; ++i)
                    for (size_type j = jj; j < jend; ++j)
                        dst[j * R + i] = src[i * C + j];
            }
        }

        // =====================================================================
        // AVX2 8×8 float in-register transpose
        // =====================================================================
        static void transpose_8x8f_kernel(const float* src, size_type src_stride,
                                            float* dst, size_type dst_stride) {
            // Load 8 rows
            __m256 r0 = _mm256_loadu_ps(src);
            __m256 r1 = _mm256_loadu_ps(src + src_stride);
            __m256 r2 = _mm256_loadu_ps(src + 2 * src_stride);
            __m256 r3 = _mm256_loadu_ps(src + 3 * src_stride);
            __m256 r4 = _mm256_loadu_ps(src + 4 * src_stride);
            __m256 r5 = _mm256_loadu_ps(src + 5 * src_stride);
            __m256 r6 = _mm256_loadu_ps(src + 6 * src_stride);
            __m256 r7 = _mm256_loadu_ps(src + 7 * src_stride);

            // Phase 1: in-lane 4x4 transpose within each 128-bit lane (unpacklo/hi + shuffle)
            __m256 t0 = _mm256_unpacklo_ps(r0, r1);  // a0b0 a1b1 a4b4 a5b5
            __m256 t1 = _mm256_unpackhi_ps(r0, r1);  // a2b2 a3b3 a6b6 a7b7
            __m256 t2 = _mm256_unpacklo_ps(r2, r3);
            __m256 t3 = _mm256_unpackhi_ps(r2, r3);
            __m256 t4 = _mm256_unpacklo_ps(r4, r5);
            __m256 t5 = _mm256_unpackhi_ps(r4, r5);
            __m256 t6 = _mm256_unpacklo_ps(r6, r7);
            __m256 t7 = _mm256_unpackhi_ps(r6, r7);

            // Phase 2: form 2x2 blocks via shuffle_ps
            __m256 s0 = _mm256_shuffle_ps(t0, t2, 0x44); // row0-lo, row0-hi (128-bit lanes)
            __m256 s1 = _mm256_shuffle_ps(t0, t2, 0xEE);
            __m256 s2 = _mm256_shuffle_ps(t1, t3, 0x44);
            __m256 s3 = _mm256_shuffle_ps(t1, t3, 0xEE);
            __m256 s4 = _mm256_shuffle_ps(t4, t6, 0x44);
            __m256 s5 = _mm256_shuffle_ps(t4, t6, 0xEE);
            __m256 s6 = _mm256_shuffle_ps(t5, t7, 0x44);
            __m256 s7 = _mm256_shuffle_ps(t5, t7, 0xEE);

            // Phase 3: swap the 128-bit lanes via permute2f128
            _mm256_storeu_ps(dst,                    _mm256_permute2f128_ps(s0, s4, 0x20));
            _mm256_storeu_ps(dst + dst_stride,       _mm256_permute2f128_ps(s1, s5, 0x20));
            _mm256_storeu_ps(dst + 2 * dst_stride,   _mm256_permute2f128_ps(s2, s6, 0x20));
            _mm256_storeu_ps(dst + 3 * dst_stride,   _mm256_permute2f128_ps(s3, s7, 0x20));
            _mm256_storeu_ps(dst + 4 * dst_stride,   _mm256_permute2f128_ps(s0, s4, 0x31));
            _mm256_storeu_ps(dst + 5 * dst_stride,   _mm256_permute2f128_ps(s1, s5, 0x31));
            _mm256_storeu_ps(dst + 6 * dst_stride,   _mm256_permute2f128_ps(s2, s6, 0x31));
            _mm256_storeu_ps(dst + 7 * dst_stride,   _mm256_permute2f128_ps(s3, s7, 0x31));
        }

        // float block transpose (one block strip: rows [ii, iend) x all columns)
        static void transpose_block_avx2_float_strip(
                const float* src, float* dst,
                size_type R, size_type C,
                size_type ii, size_type iend) {
            constexpr size_type BS = 32;
            for (size_type jj = 0; jj < C; jj += BS) {
                const size_type jend = (jj + BS < C) ? jj + BS : C;
                size_type i = ii;
                for (; i + 8 <= iend; i += 8) {
                    size_type j = jj;
                    for (; j + 8 <= jend; j += 8) {
                        transpose_8x8f_kernel(
                            src + i * C + j, C,
                            dst + j * R + i, R);
                    }
                    for (; j < jend; ++j)
                        for (size_type k = i; k < i + 8; ++k)
                            dst[j * R + k] = src[k * C + j];
                }
                for (; i < iend; ++i)
                    for (size_type j = jj; j < jend; ++j)
                        dst[j * R + i] = src[i * C + j];
            }
        }

        // AVX2 dispatch (shared entry point for double/float)
        static void transpose_block_avx2_dispatch(const T* src, T* dst,
                                                    size_type R, size_type C) {
            constexpr size_type BS = 32;
            const size_type total = R * C;

            // Large matrices: thread-pool parallelization (per row block strip)
            if (total >= TRANSPOSE_PARALLEL_THRESHOLD) {
                auto& pool = threadPool();
                const unsigned hw = std::thread::hardware_concurrency();
                const size_type nthreads = (hw > 1) ? hw : 1;
                if (nthreads > 1) {
                    const size_type n_strips = (R + BS - 1) / BS;
                    const size_type strips_per_thread = (n_strips + nthreads - 1) / nthreads;

                    std::vector<std::future<void>> futures;
                    futures.reserve(nthreads);

                    for (size_type t = 0; t < nthreads; ++t) {
                        const size_type strip_begin = t * strips_per_thread;
                        if (strip_begin >= n_strips) break;
                        const size_type strip_end = std::min((t + 1) * strips_per_thread, n_strips);
                        const size_type ii = strip_begin * BS;
                        const size_type iend = std::min(strip_end * BS, R);

                        futures.push_back(pool.submit([=]() {
                            if constexpr (std::is_same_v<T, double>) {
                                transpose_block_avx2_double_strip(src, dst, R, C, ii, iend);
                            } else {
                                transpose_block_avx2_float_strip(src, dst, R, C, ii, iend);
                            }
                        }));
                    }
                    for (auto& f : futures) f.get();
                    return;
                }
            }

            // Single thread
            if constexpr (std::is_same_v<T, double>) {
                transpose_block_avx2_double_strip(src, dst, R, C, 0, R);
            } else {
                transpose_block_avx2_float_strip(src, dst, R, C, 0, R);
            }
        }

#else
        // When AVX2 is unsupported, fall back to scalar block transpose
        static void transpose_block_avx2_dispatch(const T* src, T* dst,
                                                    size_type R, size_type C) {
            transpose_block_scalar(src, dst, R, C);
        }
#endif

    public:
        // Determinant
        template <typename Mat>
        static T determinant(const Mat& mat) {
            if (mat.rows() != mat.cols()) {
                throw std::invalid_argument("Matrix must be square to calculate determinant");
            }

            const size_type n = mat.rows();

            // Fast path for 2x2 matrices
            if (n == 2) {
                return mat(0, 0) * mat(1, 1) - mat(0, 1) * mat(1, 0);
            }

            // Fast path for 3x3 matrices
            if (n == 3) {
                return mat(0, 0) * (mat(1, 1) * mat(2, 2) - mat(1, 2) * mat(2, 1)) -
                    mat(0, 1) * (mat(1, 0) * mat(2, 2) - mat(1, 2) * mat(2, 0)) +
                    mat(0, 2) * (mat(1, 0) * mat(2, 1) - mat(1, 1) * mat(2, 0));
            }

            // For larger matrices, use cofactor expansion
            T det = static_cast<T>(0);
            int sign = 1;

            for (size_type j = 0; j < n; ++j) {
                // Build the cofactor matrix
                Mat minor_matrix(n - 1, n - 1);
                for (size_type i = 1; i < n; ++i) {
                    size_type col_idx = 0;
                    for (size_type k = 0; k < n; ++k) {
                        if (k != j) {
                            minor_matrix(i - 1, col_idx++) = mat(i, k);
                        }
                    }
                }

                det += sign * mat(0, j) * determinant(minor_matrix);
                sign = -sign;
            }

            return det;
        }

        // Matrix inverse
        template <typename Mat, typename MatResult>
        static bool inverse(const Mat& mat, MatResult& result) {
            if (mat.rows() != mat.cols() || result.rows() != mat.rows() || result.cols() != mat.cols()) {
                throw std::invalid_argument("Matrix must be square for inverse calculation");
            }

            const size_type n = mat.rows();

            // Build the augmented matrix [A|I]
            Mat augmented(n, 2 * n);

            // Initialize the augmented matrix
            for (size_type i = 0; i < n; ++i) {
                for (size_type j = 0; j < n; ++j) {
                    augmented(i, j) = mat(i, j);
                    augmented(i, j + n) = (i == j) ? static_cast<T>(1) : static_cast<T>(0);
                }
            }

            // Gauss-Jordan elimination
            for (size_type i = 0; i < n; ++i) {
                // Find the pivot element
                size_type pivot_row = i;
                T max_val = std::abs(augmented(i, i));

                for (size_type k = i + 1; k < n; ++k) {
                    if (std::abs(augmented(k, i)) > max_val) {
                        max_val = std::abs(augmented(k, i));
                        pivot_row = k;
                    }
                }

                // The matrix is singular
                if (max_val < std::numeric_limits<T>::epsilon()) {
                    return false;  // The inverse does not exist
                }

                // Swap rows
                if (pivot_row != i) {
                    for (size_type j = 0; j < 2 * n; ++j) {
                        std::swap(augmented(i, j), augmented(pivot_row, j));
                    }
                }

                // Scale the pivot element to 1
                T pivot = augmented(i, i);
                for (size_type j = 0; j < 2 * n; ++j) {
                    augmented(i, j) /= pivot;
                }

                // Zero the pivot column in the other rows
                for (size_type k = 0; k < n; ++k) {
                    if (k != i) {
                        T factor = augmented(k, i);
                        for (size_type j = 0; j < 2 * n; ++j) {
                            augmented(k, j) -= factor * augmented(i, j);
                        }
                    }
                }
            }

            // Copy the inverse portion into the result
            for (size_type i = 0; i < n; ++i) {
                for (size_type j = 0; j < n; ++j) {
                    result(i, j) = augmented(i, j + n);
                }
            }

            return true;
        }
    };

    // SIMD computation policy
    template <typename T>
    class SIMDComputePolicy : public DefaultComputePolicy<T> {
    public:
        using value_type = T;
        using size_type = std::size_t;

        // Vector addition (SIMD)
        template <typename VecA, typename VecB, typename VecResult>
        static void vector_add(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for addition");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                computation::simd::add_simd<T>(&a[0], &b[0], &result[0], a.size());
            } else {
                DefaultComputePolicy<T>::vector_add(a, b, result);
            }
        }

        // Vector subtraction (SIMD)
        template <typename VecA, typename VecB, typename VecResult>
        static void vector_subtract(const VecA& a, const VecB& b, VecResult& result) {
            if (a.size() != b.size() || a.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for subtraction");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                computation::simd::subtract_simd<T>(&a[0], &b[0], &result[0], a.size());
            } else {
                DefaultComputePolicy<T>::vector_subtract(a, b, result);
            }
        }

        // Vector dot product (SIMD)
        template <typename VecA, typename VecB>
        static T dot(const VecA& a, const VecB& b) {
            if (a.size() != b.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for dot product");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                return computation::simd::dot_product_simd<T>(&a[0], &b[0], a.size());
            } else {
                return DefaultComputePolicy<T>::dot(a, b);
            }
        }

        // Vector norm (SIMD)
        template <typename Vec>
        static T norm(const Vec& vec) {
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                T sum = computation::simd::dot_product_simd<T>(&vec[0], &vec[0], vec.size());
                return std::sqrt(sum);
            } else {
                return DefaultComputePolicy<T>::norm(vec);
            }
        }

        // Vector-scalar multiplication (SIMD)
        template <typename Vec, typename VecResult>
        static void vector_scalar_multiply(const Vec& vec, const T& scalar, VecResult& result) {
            if (vec.size() != result.size()) {
                throw std::invalid_argument("Vector dimensions mismatch for scalar multiplication");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                // Copy then scale
                std::memcpy(&result[0], &vec[0], vec.size() * sizeof(T));
                computation::simd::scale_simd<T>(&result[0], scalar, result.size());
            } else {
                DefaultComputePolicy<T>::vector_scalar_multiply(vec, scalar, result);
            }
        }
    };

    // MKL computation policy (to be implemented in the future)
    template <typename T>
    class MKLComputePolicy : public DefaultComputePolicy<T> {
        // Add Intel MKL-optimized implementations here
    };

    // Automatic policy selection
    template <typename T>
    struct OptimalPolicy {
#ifdef USE_MKL
        using type = MKLComputePolicy<T>;
#elif defined(SANGI_HAS_AVX512) || defined(SANGI_HAS_AVX2) || defined(SANGI_HAS_AVX) || \
      defined(SANGI_HAS_SSE2) || defined(SANGI_HAS_NEON) || defined(USE_SIMD)
        using type = SIMDComputePolicy<T>;
#else
        using type = DefaultComputePolicy<T>;
#endif
    };

    template <typename T>
    using OptimalPolicyType = typename OptimalPolicy<T>::type;

} // namespace sangi