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

// blas.hpp
// BLAS level 1/2/3 interface
//
// Provides a BLAS (Basic Linear Algebra Subprograms) style API.
// Foundation for switching to an MKL/OpenBLAS backend in the future.
// Currently a native implementation (operates directly on sangi's Vector/Matrix).

#ifndef SANGI_BLAS_HPP
#define SANGI_BLAS_HPP

#include <math/core/matrix.hpp>
#include <math/core/computation/simd_backend.hpp>
#include <math/linalg/static_dims.hpp>
#include <concepts>
#include <type_traits>
#include <cmath>
#include <stdexcept>
#include <cassert>

namespace sangi {
namespace blas {

    // ====================================================================
    // Level 1: vector-vector operations
    // ====================================================================

    /**
     * @brief Dot product: result = x^T * y
     */
    template<typename VX, typename VY>
        requires sangi::BaseVectorLike<VX>
              && sangi::BaseVectorLike<VY>
              && std::same_as<sangi::element_t<VX>, sangi::element_t<VY>>
    sangi::element_t<VX> dot(const VX& x, const VY& y) {
        using T = sangi::element_t<VX>;
        // Compile-time size check when both vectors carry static sizes.
        if constexpr (sangi::dims_both_static(
                          std::remove_cvref_t<VX>::static_size,
                          std::remove_cvref_t<VY>::static_size)) {
            static_assert(std::remove_cvref_t<VX>::static_size
                          == std::remove_cvref_t<VY>::static_size,
                          "sangi::blas::dot: vector sizes must match (static check)");
        }

        if (x.size() != y.size()) {
            assert(false && "DimensionError: blas::dot: size mismatch");
            throw DimensionError("blas::dot: size mismatch");
        }
        if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
            if (x.is_contiguous() && y.is_contiguous())
                return computation::simd::dot_product_simd<T>(x.data(), y.data(), x.size());
        }
        T result = T{0};
        for (std::size_t i = 0; i < x.size(); ++i)
            result += x[i] * y[i];
        return result;
    }

    /**
     * @brief 2-norm: result = ‖x‖₂
     */
    template<typename VX>
        requires sangi::BaseVectorLike<VX>
    sangi::element_t<VX> nrm2(const VX& x) {
        using T = sangi::element_t<VX>;
        if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
            if (x.is_contiguous()) {
                T sum = computation::simd::dot_product_simd<T>(x.data(), x.data(), x.size());
                return std::sqrt(sum);
            }
            T sum = T{0};
            for (std::size_t i = 0; i < x.size(); ++i)
                sum += x[i] * x[i];
            return std::sqrt(sum);
        } else {
            T sum = T{0};
            for (std::size_t i = 0; i < x.size(); ++i)
                sum += x[i] * x[i];
            return static_cast<T>(std::sqrt(static_cast<double>(sum)));
        }
    }

    /**
     * @brief 1-norm (sum of absolute values): result = ‖x‖₁ = Σ|x_i|
     *
     * Equivalent to BLAS dasum.
     */
    template<typename VX>
        requires sangi::BaseVectorLike<VX>
    sangi::element_t<VX> asum(const VX& x) {
        using T = sangi::element_t<VX>;
        T sum = T{0};
        for (std::size_t i = 0; i < x.size(); ++i)
            sum += static_cast<T>(std::abs(static_cast<double>(x[i])));
        return sum;
    }

    /**
     * @brief Index of the element with largest absolute value: result = argmax_i |x_i|
     *
     * Equivalent to BLAS idamax. Returns 0 for an empty vector.
     */
    template<typename VX>
        requires sangi::BaseVectorLike<VX>
    std::size_t iamax(const VX& x) {
        using T = sangi::element_t<VX>;
        if (x.size() == 0) return 0;
        std::size_t idx = 0;
        double max_val = std::abs(static_cast<double>(x[0]));
        for (std::size_t i = 1; i < x.size(); ++i) {
            double ai = std::abs(static_cast<double>(x[i]));
            if (ai > max_val) { max_val = ai; idx = i; }
        }
        return idx;
    }

    /**
     * @brief Scaling: x ← alpha * x
     */
    template<typename T>
    void scal(T alpha, Vector<T>& x) {
        if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
            computation::simd::scale_simd<T>(x.data(), alpha, x.size());
        } else {
            for (std::size_t i = 0; i < x.size(); ++i)
                x[i] *= alpha;
        }
    }

    /**
     * @brief AXPY: y ← alpha * x + y
     */
    template<typename T>
    void axpy(T alpha, const BaseVector<T>& x, Vector<T>& y) {
        if (x.size() != y.size()) {
            assert(false && "DimensionError: blas::axpy: size mismatch");
            throw DimensionError("blas::axpy: size mismatch");
        }
        if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
            if (x.is_contiguous()) {
                computation::simd::axpy_simd<T>(y.data(), alpha, x.data(), x.size());
                return;
            }
        }
        for (std::size_t i = 0; i < x.size(); ++i)
            y[i] += alpha * x[i];
    }

    /**
     * @brief Copy: y ← x
     */
    template<typename T>
    void copy(const BaseVector<T>& x, Vector<T>& y) {
        if (x.size() != y.size()) {
            assert(false && "DimensionError: blas::copy: size mismatch");
            throw DimensionError("blas::copy: size mismatch");
        }
        for (std::size_t i = 0; i < x.size(); ++i)
            y[i] = x[i];
    }

    /**
     * @brief Swap: x ↔ y
     */
    template<typename T>
    void swap(Vector<T>& x, Vector<T>& y) {
        if (x.size() != y.size()) {
            assert(false && "DimensionError: blas::swap: size mismatch");
            throw DimensionError("blas::swap: size mismatch");
        }
        for (std::size_t i = 0; i < x.size(); ++i)
            std::swap(x[i], y[i]);
    }

    // ====================================================================
    // Level 2: matrix-vector operations
    // ====================================================================

    /**
     * @brief General matrix-vector product: y ← alpha * op(A) * x + beta * y
     *
     * @param trans false: op(A) = A, true: op(A) = A^T
     * @param alpha scalar coefficient
     * @param A m×n matrix
     * @param x input vector
     * @param beta scalar coefficient
     * @param y output vector (updated in-place)
     */
    template<typename T>
    void gemv(bool trans, T alpha, const BaseMatrix<T>& A,
              const BaseVector<T>& x, T beta, BaseVector<T>& y)
    {
        const auto m = A.rows();
        const auto n = A.cols();
        // The SIMD path is only valid when each row of A is contiguous (= col_stride == 1)
        // and y is contiguous. Safe even when a transposed view, etc. arrives via BaseMatrix.
        const bool A_row_contig = (A.col_stride() == 1);

        if (!trans) {
            // y ← alpha * A * x + beta * y
            if (x.size() != n || y.size() != m) {
                assert(false && "DimensionError: blas::gemv: dimension mismatch");
                throw DimensionError("blas::gemv: dimension mismatch");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if (A_row_contig && x.is_contiguous() && y.is_contiguous()) {
                    for (std::size_t i = 0; i < m; ++i) {
                        T sum = computation::simd::dot_product_simd<T>(&A(i, 0), x.data(), n);
                        y[i] = alpha * sum + beta * y[i];
                    }
                    return;
                }
            }
            for (std::size_t i = 0; i < m; ++i) {
                T sum = T{0};
                for (std::size_t j = 0; j < n; ++j)
                    sum += A(i, j) * x[j];
                y[i] = alpha * sum + beta * y[i];
            }
        } else {
            // y ← alpha * A^T * x + beta * y
            if (x.size() != m || y.size() != n) {
                assert(false && "DimensionError: blas::gemv: dimension mismatch (transpose)");
                throw DimensionError("blas::gemv: dimension mismatch (transpose)");
            }
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if (A_row_contig && y.is_contiguous()) {
                    computation::simd::scale_simd<T>(y.data(), beta, n);
                    for (std::size_t i = 0; i < m; ++i) {
                        T ax = alpha * x[i];
                        computation::simd::axpy_simd<T>(y.data(), ax, &A(i, 0), n);
                    }
                    return;
                }
            }
            for (std::size_t j = 0; j < n; ++j)
                y[j] *= beta;
            for (std::size_t i = 0; i < m; ++i) {
                T ax = alpha * x[i];
                for (std::size_t j = 0; j < n; ++j)
                    y[j] += ax * A(i, j);
            }
        }
    }

    /**
     * @brief Solves a triangular system of equations: x ← op(A)⁻¹ * x
     *
     * @param upper true: upper triangular, false: lower triangular
     * @param trans true: A^T, false: A
     * @param unit_diag true: treat diagonal entries as 1
     * @param A triangular matrix
     * @param x right-hand-side vector (returns the solution via in-place update)
     */
    template<typename T>
    void trsv(bool upper, bool trans, bool unit_diag,
              const BaseMatrix<T>& A, BaseVector<T>& x)
    {
        const auto n = A.rows();
        if (A.cols() != n || x.size() != n) {
            assert(false && "DimensionError: blas::trsv: dimension mismatch");
            throw DimensionError("blas::trsv: dimension mismatch");
        }

        if (!trans) {
            if (upper) {
                // Upper triangular, back substitution
                for (std::size_t ii = 0; ii < n; ++ii) {
                    std::size_t i = n - 1 - ii;
                    for (std::size_t j = i + 1; j < n; ++j)
                        x[i] -= A(i, j) * x[j];
                    if (!unit_diag) x[i] /= A(i, i);
                }
            } else {
                // Lower triangular, forward substitution
                for (std::size_t i = 0; i < n; ++i) {
                    for (std::size_t j = 0; j < i; ++j)
                        x[i] -= A(i, j) * x[j];
                    if (!unit_diag) x[i] /= A(i, i);
                }
            }
        } else {
            if (upper) {
                // Upper triangular transposed → forward substitution as lower triangular
                for (std::size_t i = 0; i < n; ++i) {
                    for (std::size_t j = 0; j < i; ++j)
                        x[i] -= A(j, i) * x[j];
                    if (!unit_diag) x[i] /= A(i, i);
                }
            } else {
                // Lower triangular transposed → back substitution as upper triangular
                for (std::size_t ii = 0; ii < n; ++ii) {
                    std::size_t i = n - 1 - ii;
                    for (std::size_t j = i + 1; j < n; ++j)
                        x[i] -= A(j, i) * x[j];
                    if (!unit_diag) x[i] /= A(i, i);
                }
            }
        }
    }

    /**
     * @brief Rank-1 update: A ← alpha * x * y^T + A
     */
    template<typename T>
    void ger(T alpha, const BaseVector<T>& x, const BaseVector<T>& y, BaseMatrix<T>& A) {
        if (x.size() != A.rows() || y.size() != A.cols()) {
            assert(false && "DimensionError: blas::ger: dimension mismatch");
            throw DimensionError("blas::ger: dimension mismatch");
        }
        for (std::size_t i = 0; i < A.rows(); ++i) {
            T ax = alpha * x[i];
            for (std::size_t j = 0; j < A.cols(); ++j)
                A(i, j) += ax * y[j];
        }
    }

    /**
     * @brief Symmetric rank-1 update: A ← alpha * x * x^T + A
     *
     * A is a symmetric matrix (update only the upper triangle, then copy to the lower).
     */
    template<typename T>
    void syr(bool upper, T alpha, const BaseVector<T>& x, BaseMatrix<T>& A) {
        const auto n = x.size();
        if (A.rows() != n || A.cols() != n) {
            assert(false && "DimensionError: blas::syr: dimension mismatch");
            throw DimensionError("blas::syr: dimension mismatch");
        }
        if (upper) {
            for (std::size_t i = 0; i < n; ++i) {
                T ax = alpha * x[i];
                for (std::size_t j = i; j < n; ++j) {
                    A(i, j) += ax * x[j];
                    if (i != j) A(j, i) = A(i, j);
                }
            }
        } else {
            for (std::size_t j = 0; j < n; ++j) {
                T ax = alpha * x[j];
                for (std::size_t i = j; i < n; ++i) {
                    A(i, j) += ax * x[i];
                    if (i != j) A(j, i) = A(i, j);
                }
            }
        }
    }

    // ====================================================================
    // Level 3: matrix-matrix operations
    // ====================================================================

    /**
     * @brief General matrix product: C ← alpha * op(A) * op(B) + beta * C
     *
     * @param transA false: op(A) = A, true: op(A) = A^T
     * @param transB false: op(B) = B, true: op(B) = B^T
     */
    template<typename T>
    void gemm(bool transA, bool transB, T alpha,
              const BaseMatrix<T>& A, const BaseMatrix<T>& B,
              T beta, BaseMatrix<T>& C)
    {
        const auto opA_rows = transA ? A.cols() : A.rows();
        const auto opA_cols = transA ? A.rows() : A.cols();
        const auto opB_rows = transB ? B.cols() : B.rows();
        const auto opB_cols = transB ? B.rows() : B.cols();

        if (opA_cols != opB_rows || C.rows() != opA_rows || C.cols() != opB_cols) {
            assert(false && "DimensionError: blas::gemm: dimension mismatch");
            throw DimensionError("blas::gemm: dimension mismatch");
        }

        const auto m = opA_rows;
        const auto n = opB_cols;
        const auto k = opA_cols;

        // The SIMD blocking path requires each row of B / C to be row-contiguous (col_stride == 1)
        const bool BC_row_contig = (B.col_stride() == 1) && (C.col_stride() == 1);

        // C ← beta * C
        if (beta == T{0}) {
            for (std::size_t i = 0; i < m; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    C(i, j) = T{0};
        } else if (beta != T{1}) {
            for (std::size_t i = 0; i < m; ++i)
                for (std::size_t j = 0; j < n; ++j)
                    C(i, j) *= beta;
        }

        // C += alpha * op(A) * op(B)
        // 4 transpose combinations
        if (!transA && !transB) {
            if constexpr (std::is_same_v<T, float> || std::is_same_v<T, double>) {
                if (BC_row_contig) {
                    // Blocking + SIMD (ipj loop order)
                    constexpr std::size_t BLK = 64;
                    for (std::size_t i0 = 0; i0 < m; i0 += BLK) {
                        const std::size_t i1 = (std::min)(i0 + BLK, m);
                        for (std::size_t p0 = 0; p0 < k; p0 += BLK) {
                            const std::size_t p1 = (std::min)(p0 + BLK, k);
                            for (std::size_t j0 = 0; j0 < n; j0 += BLK) {
                                const std::size_t j1 = (std::min)(j0 + BLK, n);
                                const std::size_t jlen = j1 - j0;
                                // Micro-kernel
                                for (std::size_t i = i0; i < i1; ++i) {
                                    for (std::size_t p = p0; p < p1; ++p) {
                                        T aip = alpha * A(i, p);
                                        computation::simd::axpy_simd<T>(
                                            &C(i, j0), aip, &B(p, j0), jlen);
                                    }
                                }
                            }
                        }
                    }
                } else {
                    // Non-contiguous fallback
                    for (std::size_t i = 0; i < m; ++i)
                        for (std::size_t p = 0; p < k; ++p) {
                            T aip = alpha * A(i, p);
                            for (std::size_t j = 0; j < n; ++j)
                                C(i, j) += aip * B(p, j);
                        }
                }
            } else {
                for (std::size_t i = 0; i < m; ++i)
                    for (std::size_t p = 0; p < k; ++p) {
                        T aip = alpha * A(i, p);
                        for (std::size_t j = 0; j < n; ++j)
                            C(i, j) += aip * B(p, j);
                    }
            }
        } else if (transA && !transB) {
            for (std::size_t i = 0; i < m; ++i)
                for (std::size_t p = 0; p < k; ++p) {
                    T api = alpha * A(p, i);
                    for (std::size_t j = 0; j < n; ++j)
                        C(i, j) += api * B(p, j);
                }
        } else if (!transA && transB) {
            for (std::size_t i = 0; i < m; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    T sum = T{0};
                    for (std::size_t p = 0; p < k; ++p)
                        sum += A(i, p) * B(j, p);
                    C(i, j) += alpha * sum;
                }
        } else {
            // transA && transB
            for (std::size_t i = 0; i < m; ++i)
                for (std::size_t j = 0; j < n; ++j) {
                    T sum = T{0};
                    for (std::size_t p = 0; p < k; ++p)
                        sum += A(p, i) * B(j, p);
                    C(i, j) += alpha * sum;
                }
        }
    }

    /**
     * @brief Solves a triangular system of equations (matrix version): B ← op(A)⁻¹ * B  or  B ← B * op(A)⁻¹
     *
     * @param side true: left (A⁻¹B), false: right (BA⁻¹)
     * @param upper true: upper triangular, false: lower triangular
     * @param trans true: A^T, false: A
     * @param unit_diag true: treat diagonal entries as 1
     * @param alpha scalar coefficient (B ← alpha * ...)
     * @param A triangular matrix (n×n)
     * @param B input/output matrix
     */
    template<typename T>
    void trsm(bool side_left, bool upper, bool trans, bool unit_diag,
              T alpha, const BaseMatrix<T>& A, BaseMatrix<T>& B)
    {
        const auto n = A.rows();
        if (A.cols() != n) {
            assert(false && "DimensionError: blas::trsm: A must be square");
            throw DimensionError("blas::trsm: A must be square");
        }

        // alpha scaling
        if (alpha != T{1}) {
            for (std::size_t i = 0; i < B.rows(); ++i)
                for (std::size_t j = 0; j < B.cols(); ++j)
                    B(i, j) *= alpha;
        }

        if (side_left) {
            // B ← A⁻¹ * B: trsv each column independently
            if (B.rows() != n) {
                assert(false && "DimensionError: blas::trsm: dimension mismatch (left)");
                throw DimensionError("blas::trsm: dimension mismatch (left)");
            }
            for (std::size_t col = 0; col < B.cols(); ++col) {
                Vector<T> b_col(n);
                for (std::size_t i = 0; i < n; ++i) b_col[i] = B(i, col);
                trsv(upper, trans, unit_diag, A, b_col);
                for (std::size_t i = 0; i < n; ++i) B(i, col) = b_col[i];
            }
        } else {
            // B ← B * A⁻¹: trsv each row independently (equivalent to trsv of A^T)
            if (B.cols() != n) {
                assert(false && "DimensionError: blas::trsm: dimension mismatch (right)");
                throw DimensionError("blas::trsm: dimension mismatch (right)");
            }
            for (std::size_t row = 0; row < B.rows(); ++row) {
                Vector<T> b_row(n);
                for (std::size_t j = 0; j < n; ++j) b_row[j] = B(row, j);
                trsv(upper, !trans, unit_diag, A, b_row);
                for (std::size_t j = 0; j < n; ++j) B(row, j) = b_row[j];
            }
        }
    }

    /**
     * @brief Symmetric rank-k update: C ← alpha * A * A^T + beta * C  (trans=false)
     *                        C ← alpha * A^T * A + beta * C  (trans=true)
     *
     * C is a symmetric matrix (n×n).
     * trans=false: A is n×k, C ← alpha*A*A^T + beta*C
     * trans=true:  A is k×n, C ← alpha*A^T*A + beta*C
     */
    template<typename T>
    void syrk(bool upper, bool trans, T alpha,
              const BaseMatrix<T>& A, T beta, BaseMatrix<T>& C)
    {
        const auto n = C.rows();
        if (C.cols() != n) {
            assert(false && "DimensionError: blas::syrk: C must be square");
            throw DimensionError("blas::syrk: C must be square");
        }

        std::size_t k;
        if (!trans) {
            if (A.rows() != n) {
                assert(false && "DimensionError: blas::syrk: dimension mismatch");
                throw DimensionError("blas::syrk: dimension mismatch");
            }
            k = A.cols();
        } else {
            if (A.cols() != n) {
                assert(false && "DimensionError: blas::syrk: dimension mismatch");
                throw DimensionError("blas::syrk: dimension mismatch");
            }
            k = A.rows();
        }

        // C ← beta * C (upper or lower triangle only)
        for (std::size_t i = 0; i < n; ++i) {
            std::size_t j_start = upper ? i : 0;
            std::size_t j_end   = upper ? n : i + 1;
            for (std::size_t j = j_start; j < j_end; ++j)
                C(i, j) *= beta;
        }

        // C += alpha * op(A) * op(A)^T
        if (!trans) {
            // C += alpha * A * A^T
            for (std::size_t i = 0; i < n; ++i) {
                std::size_t j_start = upper ? i : 0;
                std::size_t j_end   = upper ? n : i + 1;
                for (std::size_t j = j_start; j < j_end; ++j) {
                    T sum = T{0};
                    for (std::size_t p = 0; p < k; ++p)
                        sum += A(i, p) * A(j, p);
                    C(i, j) += alpha * sum;
                }
            }
        } else {
            // C += alpha * A^T * A
            for (std::size_t i = 0; i < n; ++i) {
                std::size_t j_start = upper ? i : 0;
                std::size_t j_end   = upper ? n : i + 1;
                for (std::size_t j = j_start; j < j_end; ++j) {
                    T sum = T{0};
                    for (std::size_t p = 0; p < k; ++p)
                        sum += A(p, i) * A(p, j);
                    C(i, j) += alpha * sum;
                }
            }
        }

        // Copy the symmetric part
        for (std::size_t i = 0; i < n; ++i)
            for (std::size_t j = i + 1; j < n; ++j) {
                if (upper) C(j, i) = C(i, j);
                else       C(i, j) = C(j, i);
            }
    }

} // namespace blas
} // namespace sangi

#endif // SANGI_BLAS_HPP