#pragma once

#include "micro-kernels/householder-updowndate.hpp"

namespace hyhound {

namespace detail {
using micro_kernels::matrix_accessor;
using micro_kernels::mut_matrix_accessor;
using micro_kernels::householder::Config;
using micro_kernels::householder::mut_W_accessor;
/// Greedily partitions the given size S into multiples of M and Ms, and calls
/// `Func<{R, M}>` for these sizes. Practically, M and Ms should
/// be the sizes of the available instantiations of @ref Func, in
/// descending order.
/// For example, `tile_tail<{4, 8}, 8, 4, 2, 1>(15)` will call
/// `Func<4, 8>` on the first block row of its arguments, then
/// `Func<4, 4>` on the second block row, `Func<4, 2>`
/// and finally `downdate_tail<4, 1>` for the bottom row.
template <class T, template <auto, class, class> class Func, Config Conf,
          class UpDown, index_t M, index_t... Ms>
inline void tile_tail(index_t rowsA, index_t colsA0, index_t colsA,
                      mut_W_accessor<T> W, T *L, index_t ldL, const T *B,
                      index_t ldB, T *A, index_t ldA, UpDown updown) noexcept {
    constexpr auto simd_M = micro_kernels::native_simd_size<T>;
    // If the block size is larger than the config allows, skip it.
    constexpr bool skip_large_M = M > Conf.block_size_s;
    // If the block size is not efficiently vectorizable, and is not yet a
    // remainder block size smaller than the vector length, skip it.
    // E.g. on AVX2, S=12, Ms={12, 8, 4, 2, 1} should use the kernel of size 12,
    // which will use 3 vector registers of size 4 internally, but on AVX512,
    // it should first use the kernel of size 8, and then 4.
    constexpr bool skip_suboptimal_M = M > simd_M && (M % simd_M) != 0;
    if constexpr (skip_large_M || skip_suboptimal_M) {
        if constexpr (sizeof...(Ms) > 0)
            tile_tail<T, Func, Conf, UpDown, Ms...>(
                rowsA, colsA0, colsA, W, L, ldL, B, ldB, A, ldA, updown);
        return;
    }
    while (rowsA >= M) {
        constexpr Config NewConf{.block_size_r = Conf.block_size_r,
                                 .block_size_s = M};
        Func<NewConf, T, UpDown>{}(colsA0, colsA, W, L, ldL, B, ldB, A, ldA,
                                   updown);
        L += M;
        A += M;
        rowsA -= M;
    }
    if constexpr (sizeof...(Ms) > 0)
        if (rowsA > 0)
            tile_tail<T, Func, Conf, UpDown, Ms...>(
                rowsA, colsA0, colsA, W, L, ldL, B, ldB, A, ldA, updown);
}

template <Config Conf, class T, class UpDown>
struct updowndate_tail_func {
    template <class... Args>
    decltype(auto) operator()(Args &&...args) const {
        return micro_kernels::householder::updowndate_tail<Conf, T, UpDown>(
            std::forward<Args>(args)...);
    }
};

} // namespace detail

/// @see @ref detail::tile_tail
/// The sizes specified here should be instantiated in the code generated by
/// CMake.
template <micro_kernels::householder::Config Conf, class T, class UpDown>
inline void updowndate_tile_tail(index_t rowsA, index_t colsA0, index_t colsA,
                                 detail::mut_W_accessor<T> W,
                                 detail::mut_matrix_accessor<T> L,
                                 detail::matrix_accessor<T> B,
                                 detail::mut_matrix_accessor<T> A,
                                 UpDown signs) {
    detail::tile_tail<T, detail::updowndate_tail_func, Conf, UpDown, //
                      32, 24, 16, 12, 8, 4, 2, 1>(
        rowsA, colsA0, colsA, W, L.data, L.outer_stride, B.data, B.outer_stride,
        A.data, A.outer_stride, signs);
}

template <micro_kernels::householder::Config Conf, class T, class UpDown>
inline void
updowndate_tail(index_t colsA0, index_t colsA, detail::mut_W_accessor<T> W,
                detail::mut_matrix_accessor<T> L, detail::matrix_accessor<T> B,
                detail::mut_matrix_accessor<T> A, UpDown signs) {
    using micro_kernels::householder::updowndate_tail;
    updowndate_tail<Conf, T, UpDown>(colsA0, colsA, W, L.data, L.outer_stride,
                                     B.data, B.outer_stride, A.data,
                                     A.outer_stride, signs);
}

template <index_t R, class T, class UpDown>
inline void updowndate_diag(index_t colsA, detail::mut_W_accessor<T> W,
                            detail::mut_matrix_accessor<T> L,
                            detail::mut_matrix_accessor<T> A, UpDown signs) {
    using micro_kernels::householder::updowndate_diag;
    updowndate_diag<R, T, UpDown>(colsA, W, L.data, L.outer_stride, A.data,
                                  A.outer_stride, signs);
}

template <index_t R, class T, class UpDown>
inline void updowndate_full(index_t colsA, detail::mut_matrix_accessor<T> L,
                            detail::mut_matrix_accessor<T> A, UpDown signs) {
    using micro_kernels::householder::updowndate_full;
    updowndate_full<R, T, UpDown>(colsA, L.data, L.outer_stride, A.data,
                                  A.outer_stride, signs);
}

} // namespace hyhound
