#ifndef CK_THREADWISE_GENERIC_TENSOR_SLICE_COPY_HPP
#define CK_THREADWISE_GENERIC_TENSOR_SLICE_COPY_HPP

#include "static_kernel_common_header.hpp"
#include "static_kernel_tensor_descriptor.hpp"
#include "static_kernel_tensor_descriptor_helper.hpp"
#include "static_kernel_tensor_coordinate.hpp"

namespace ck {

// This threadwise copy allow vector access of src and dst.
// It allows the vector size to be different on src and dst.
// The dimensions of vector access should be the same on src and dst.
// The dimension access order should be the same on src and dst.
// Will do valid mapping check on src data: Read 0 if src data has a invalid mapping
// Will do valid mapping check on dst data: No write if dst data has a invalid mapping
template <typename SrcDesc,
          typename DstDesc,
          typename SliceLengths,
          typename SrcDstDimAccessOrder,
          index_t SrcDstVectorReadWriteDim,
          index_t SrcDataPerRead,
          index_t DstDataPerWrite,
          AddressSpace SrcAddressSpace     = AddressSpace::Generic,
          AddressSpace DstAddressSpace     = AddressSpace::Generic,
          InMemoryDataOperation DstInMemOp = InMemoryDataOperation::Set,
          index_t SrcDataStride            = 1,
          index_t DstDataStride            = 1>
struct ThreadwiseGenericTensorSliceCopy_v4r2
{
    static constexpr index_t nDim = SliceLengths::Size();
    using Index                   = MultiIndex<nDim>;

    using SrcCoord = typename TensorCoordinate<SrcDesc>::type;
    using DstCoord = typename TensorCoordinate<DstDesc>::type;

    __device__ constexpr ThreadwiseGenericTensorSliceCopy_v4r2(const Index& src_slice_origin,
                                                               const Index& dst_slice_origin)
        : mSrcSliceOrigin(src_slice_origin), mDstSliceOrigin(dst_slice_origin)
    {
        static_assert(nDim == SrcDesc::GetNumOfDimension() &&
                          nDim == DstDesc::GetNumOfDimension() && nDim == SliceLengths::Size() &&
                          nDim == SrcDstDimAccessOrder::Size(),
                      "wrong! # of dimensions not the same");

        static_assert(is_valid_sequence_map<SrcDstDimAccessOrder>{}, "wrong! map is not valid");

        static_assert(SliceLengths{}[SrcDstVectorReadWriteDim] %
                              math::lcm(SrcDataPerRead, DstDataPerWrite) ==
                          0,
                      "wrong! cannot evenly divide");

        // TODO:: sanity-check if vectorized memory read/write is allowed on src and dst
    }

    __device__ constexpr ThreadwiseGenericTensorSliceCopy_v4r2()
        : ThreadwiseGenericTensorSliceCopy_v4r2(make_zero_array<index_t, nDim>(),
                                                make_zero_array<index_t, nDim>())
    {
    }

    __device__ void SetSrcSliceOrigin(SrcCoord src_slice_origin)
    {
        mSrcSliceOrigin = src_slice_origin;
    }

    __device__ void SetDstSliceOrigin(DstCoord dst_slice_origin)
    {
        mDstSliceOrigin = dst_slice_origin;
    }

    template <typename SrcData, typename DstData>
    __device__ void Run(const SrcData* p_src,
                        DstData* p_dst,
                        SrcData src_out_of_bound_value = type_convert<SrcData>{}(0.0f)) const
    {
        constexpr auto vector_access_dim = Number<SrcDstVectorReadWriteDim>{};

        constexpr auto src_data_per_access = Number<SrcDataPerRead>{};
        constexpr auto dst_data_per_access = Number<DstDataPerWrite>{};

        constexpr auto long_vector_size = Number<math::lcm(SrcDataPerRead, DstDataPerWrite)>{};

        constexpr auto long_vector_access_lengths = SliceLengths::Modify(
            vector_access_dim, SliceLengths::Get(vector_access_dim) / long_vector_size);

        ford<decltype(long_vector_access_lengths), SrcDstDimAccessOrder>{}(
            [&](auto long_vector_access_id) {
                // data id w.r.t slicing-window
                auto long_vector_data_begin_id = long_vector_access_id;
                long_vector_data_begin_id(vector_access_dim) =
                    long_vector_size * long_vector_access_id[vector_access_dim];

                // buffer to hold a src long-vector
                SrcData p_src_long_vector[long_vector_size];

                // zero out buffer
                for(index_t i = 0; i < long_vector_size; ++i)
                {
                    p_src_long_vector[i] = src_out_of_bound_value;
                }

                // load data from src to the long-vector buffer
                for(index_t i = 0; i < long_vector_size / src_data_per_access; ++i)
                {
                    auto scalar_id               = make_zero_array<index_t, nDim>();
                    scalar_id(vector_access_dim) = i * src_data_per_access;

                    const index_t buffer_offset = i * src_data_per_access;

                    const auto src_coord =
                        mSrcSliceOrigin + (long_vector_data_begin_id + scalar_id);

                    // Check src data's valid mapping situation, only check the first data in this
                    // src
                    //   vector. It's user's responsiblity to make sure all data in the src vector
                    //   has the valid/invalid mapping situation
                    if(src_coord.IsOffsetValidAssumingUpperIndexIsValid())
                    {
                        transfer_data<SrcData,
                                      SrcDataPerRead,
                                      SrcAddressSpace,
                                      AddressSpace::Vgpr,
                                      InMemoryDataOperation::Set,
                                      SrcDataStride,
                                      1>(
                            p_src, src_coord.GetOffset(), p_src_long_vector, buffer_offset);
                    }
                }

                // SrcData to DstData conversion
                DstData p_dst_long_vector[long_vector_size];

                for(index_t i = 0; i < long_vector_size; ++i)
                {
                    p_dst_long_vector[i] = type_convert<DstData>{}(p_src_long_vector[i]);
                }

                // store data from the long-vector buffer to dst
                for(index_t i = 0; i < long_vector_size / dst_data_per_access; ++i)
                {
                    auto scalar_id               = make_zero_array<index_t, nDim>();
                    scalar_id(vector_access_dim) = i * dst_data_per_access;

                    const index_t buffer_offset = i * dst_data_per_access;

                    const auto dst_coord =
                        mDstSliceOrigin + (long_vector_data_begin_id + scalar_id);

                    // Check dst data's valid mapping situation, only check the first data in this
                    // dst
                    //   vector. It's user's responsiblity to make sure all data in the dst vector
                    //   has the valid/invalid mapping situation
                    if(dst_coord.IsOffsetValidAssumingUpperIndexIsValid())
                    {
                        transfer_data<DstData,
                                      DstDataPerWrite,
                                      AddressSpace::Vgpr,
                                      DstAddressSpace,
                                      DstInMemOp,
                                      1,
                                      DstDataStride>(
                            p_dst_long_vector, buffer_offset, p_dst, dst_coord.GetOffset());
                    }
                }
            });
    }

    // Modify Length to 1, if Mask is set to false
    // Used for isolating linear dimension from non-linear dimensions
    template <index_t... Lengths, index_t... Mask>
    __device__ static constexpr auto mask_lengths(Sequence<Lengths...>, Sequence<Mask...>)
    {
        return Sequence<(Mask ? Lengths : 1)...>{};
    }

    // Will do valid mapping check on src data: Read 0 if src data has a invalid mapping
    // Will do valid mapping check on dst data: No write if dst data has a invalid mapping
    // This version is optimized for address calculation of src tensor
    // TODO: this function is not compiled to expected ISA
    template <typename SrcData, typename DstData>
    __device__ void Run_optimized_src_address_calculation(
        const SrcData* p_src,
        DstData* p_dst,
        SrcData src_out_of_bound_value = type_convert<SrcData>{}(0.0f)) const
    {
        constexpr auto vector_access_dim = Number<SrcDstVectorReadWriteDim>{};

        constexpr auto src_data_per_access = Number<SrcDataPerRead>{};
        constexpr auto dst_data_per_access = Number<DstDataPerWrite>{};

        constexpr auto long_vector_size = Number<math::lcm(SrcDataPerRead, DstDataPerWrite)>{};

        constexpr auto long_vector_access_lengths = SliceLengths::Modify(
            vector_access_dim, SliceLengths::Get(vector_access_dim) / long_vector_size);

        // separate linear dimensions from non-linear dimensions
        constexpr auto src_linear_dim_mask    = SrcDesc::GetLinearDimensionMask();
        constexpr auto src_nonlinear_dim_mask = SrcDesc::GetNonLinearDimensionMask();

        static_assert(
            src_linear_dim_mask.At(SrcDstVectorReadWriteDim) || long_vector_size == SrcDataPerRead,
            "Warning! SrcDstVectorReadWriteDim is not SrcDesc's linear dimension, performance "
            "would drop");

        // separate steps into linear and non-linear components, accoording to src tensor
        constexpr auto linear_long_vector_access_lengths =
            mask_lengths(long_vector_access_lengths, src_linear_dim_mask);

        constexpr auto nonlinear_long_vector_access_lengths =
            mask_lengths(long_vector_access_lengths, src_nonlinear_dim_mask);

        // loop over src's non-linear dimensions
        ford<decltype(nonlinear_long_vector_access_lengths)>{}(
            [&](auto nonlinear_dim_long_vector_access_id) {
                // calculate step-sizes along src's nonlinear dimensions
                auto nonlinear_dim_data_steps = nonlinear_dim_long_vector_access_id;
                nonlinear_dim_data_steps(vector_access_dim) =
                    long_vector_size * nonlinear_dim_long_vector_access_id[vector_access_dim];

                // move src cooridnate along nonlinear dimensions
                // this coordinate contains run-time per-thread offset
                const auto src_nonlinear_coord = mSrcSliceOrigin + nonlinear_dim_data_steps;

                // loop over src's linear dimensions
                ford<decltype(linear_long_vector_access_lengths)>{}(
                    [&](auto linear_dim_long_vector_access_id) {
                        // step-sizes along src's linear dimensions
                        auto linear_dim_data_steps = linear_dim_long_vector_access_id;
                        linear_dim_data_steps(vector_access_dim) =
                            long_vector_size * linear_dim_long_vector_access_id[vector_access_dim];

                        // buffer to hold a long-vector
                        SrcData p_src_long_vector[long_vector_size];

                        // zero out buffer
                        for(index_t i = 0; i < long_vector_size; ++i)
                        {
                            p_src_long_vector[i] = src_out_of_bound_value;
                        }

                        // Loop over SrcDstVectorReadWriteDim, and load data from src to the
                        //   long-vector buffer.
                        // If SrcDstVectorReadWriteDim is src's linear dimension, then src's
                        //   offset-diff due to this looping is known at compile-time. If
                        //   SrcDstVectorReadWriteDim is src's nonlinear dimension, then src's
                        //   offset-diff due to this looping is only known at run-time. For best
                        //   performance, SrcDstVectorReadWriteDim, should be src's linear dimension
                        for(index_t i = 0; i < long_vector_size / src_data_per_access; ++i)
                        {
                            auto scalar_id               = make_zero_array<index_t, nDim>();
                            scalar_id(vector_access_dim) = i * src_data_per_access;

                            const index_t buffer_offset = i * src_data_per_access;

                            // move src cooridnate along linear dimensions
                            const auto src_coord =
                                src_nonlinear_coord + (linear_dim_data_steps + scalar_id);

#if CK_EXPERIMENTAL_TENSOR_COORDINATE_USE_CALCULATE_OFFSET_DIFF // tweaking
                            // this is src compile-time offset
                            const index_t src_linear_offset =
                                src_nonlinear_coord.CalculateOffsetDiff(linear_dim_data_steps +
                                                                        scalar_id);
#else
                            // this is src compile-time offset
                            const index_t src_linear_offset =
                                src_coord.GetOffset() - src_nonlinear_coord.GetOffset();
#endif

                            // Check src data's valid mapping situation, only check the first data
                            // in this src
                            //   vector. It's user's responsiblity to make sure all data in the src
                            //   vector has the valid/invalid mapping situation
                            if(src_coord.IsOffsetValidAssumingUpperIndexIsValid())
                            {
                                transfer_data<SrcData,
                                              SrcDataPerRead,
                                              SrcAddressSpace,
                                              AddressSpace::Vgpr,
                                              InMemoryDataOperation::Set>(
                                    p_src,
                                    src_nonlinear_coord.GetOffset() + src_linear_offset,
                                    p_src_long_vector,
                                    buffer_offset);
                            }
                        }

                        // SrcData to DstData conversion
                        DstData p_dst_long_vector[long_vector_size];

                        for(index_t i = 0; i < long_vector_size; ++i)
                        {
                            p_dst_long_vector[i] = type_convert<DstData>{}(p_src_long_vector[i]);
                        }

                        // store data from the long-vector buffer to dst
                        for(index_t i = 0; i < long_vector_size / dst_data_per_access; ++i)
                        {
                            auto scalar_id               = make_zero_array<index_t, nDim>();
                            scalar_id(vector_access_dim) = i * dst_data_per_access;

                            const index_t buffer_offset = i * dst_data_per_access;

                            // dst offset is calculated here, without explicitly separating into
                            //   compile-time and per-thread component
                            const auto dst_coord =
                                mDstSliceOrigin +
                                (nonlinear_dim_data_steps + linear_dim_data_steps + scalar_id);

                            // Check dst data's valid mapping situation, only check the first data
                            // in this dst
                            //   vector. It's user's responsiblity to make sure all data in the dst
                            //   vector has the valid/invalid mapping situation
                            if(dst_coord.IsOffsetValidAssumingUpperIndexIsValid())
                            {
                                transfer_data<DstData,
                                              DstDataPerWrite,
                                              AddressSpace::Vgpr,
                                              DstAddressSpace,
                                              DstInMemOp>(
                                    p_dst_long_vector, buffer_offset, p_dst, dst_coord.GetOffset());
                            }
                        }
                    });
            });
    }

    // This version is optimized for address calculation of dst tensor
    // TODO: this function is not compiled to expected ISA
    template <typename SrcData, typename DstData>
    __device__ void Run_optimized_dst_address_calculation(
        const SrcData* p_src,
        DstData* p_dst,
        SrcData src_out_of_bound_value = type_convert<SrcData>{}(0.0f)) const
    {
        constexpr auto vector_access_dim = Number<SrcDstVectorReadWriteDim>{};

        constexpr auto src_data_per_access = Number<SrcDataPerRead>{};
        constexpr auto dst_data_per_access = Number<DstDataPerWrite>{};

        constexpr auto long_vector_size = Number<math::lcm(SrcDataPerRead, DstDataPerWrite)>{};

        constexpr auto long_vector_access_lengths = SliceLengths::Modify(
            vector_access_dim, SliceLengths::Get(vector_access_dim) / long_vector_size);

        // separate linear dimensions from non-linear dimensions
        constexpr auto dst_linear_dim_mask    = DstDesc::GetLinearDimensionMask();
        constexpr auto dst_nonlinear_dim_mask = DstDesc::GetNonLinearDimensionMask();

        static_assert(
            dst_linear_dim_mask.At(SrcDstVectorReadWriteDim) || long_vector_size == DstDataPerWrite,
            "Warning! SrcDstVectorReadWriteDim is not DstDesc's linear dimension, performance "
            "would drop");

        // separate steps into linear and non-linear components, accoording to dst tensor
        constexpr auto linear_long_vector_access_lengths =
            mask_lengths(long_vector_access_lengths, dst_linear_dim_mask);

        constexpr auto nonlinear_long_vector_access_lengths =
            mask_lengths(long_vector_access_lengths, dst_nonlinear_dim_mask);

        // loop over dst's non-linear dimensions
        ford<decltype(nonlinear_long_vector_access_lengths)>{}(
            [&](auto nonlinear_dim_long_vector_access_id) {
                // calculate step-sizes along dst's nonlinear dimensions
                auto nonlinear_dim_data_steps = nonlinear_dim_long_vector_access_id;
                nonlinear_dim_data_steps(vector_access_dim) =
                    long_vector_size * nonlinear_dim_long_vector_access_id[vector_access_dim];

                // move dst cooridnate along nonlinear dimensions
                // this coordinate contains run-time per-thread offset
                const auto dst_nonlinear_coord = mDstSliceOrigin + nonlinear_dim_data_steps;

                // loop over dst's linear dimensions
                ford<decltype(linear_long_vector_access_lengths)>{}(
                    [&](auto linear_dim_long_vector_access_id) {
                        // step-sizes along dst's linear dimensions
                        auto linear_dim_data_steps = linear_dim_long_vector_access_id;
                        linear_dim_data_steps(vector_access_dim) =
                            long_vector_size * linear_dim_long_vector_access_id[vector_access_dim];

                        // buffer to hold a long-vector
                        SrcData p_src_long_vector[long_vector_size];

                        // zero out buffer
                        for(index_t i = 0; i < long_vector_size; ++i)
                        {
                            p_src_long_vector[i] = src_out_of_bound_value;
                        }

                        // Loop over SrcDstVectorReadWriteDim, and load data from src to the
                        //   long-vector buffer.
                        // If SrcDstVectorReadWriteDim is dst's linear dimension, then dst's
                        //   offset-diff due to this looping is known at compile-time. If
                        //   SrcDstVectorReadWriteDim is dst's nonlinear dimension, then dst's
                        //   offset-diff due to this looping is only known at run-time. For best
                        //   performance, SrcDstVectorReadWriteDim, should be dst's linear dimension
                        for(index_t i = 0; i < long_vector_size / src_data_per_access; ++i)
                        {
                            auto scalar_id               = make_zero_array<index_t, nDim>();
                            scalar_id(vector_access_dim) = i * src_data_per_access;

                            const index_t buffer_offset = i * src_data_per_access;

                            // src offset is calculated here, without explicitly separating into
                            //   compile-time and per-thread component
                            const auto src_coord =
                                mSrcSliceOrigin +
                                (nonlinear_dim_data_steps + linear_dim_data_steps + scalar_id);

                            // Check src data's valid mapping situation, only check the first data
                            // in this src
                            //   vector. It's user's responsiblity to make sure all data in the src
                            //   vector has the valid/invalid mapping situation
                            if(src_coord.IsOffsetValidAssumingUpperIndexIsValid())
                            {
                                transfer_data<SrcData,
                                              SrcDataPerRead,
                                              SrcAddressSpace,
                                              AddressSpace::Vgpr,
                                              InMemoryDataOperation::Set>(
                                    p_src, src_coord.GetOffset(), p_src_long_vector, buffer_offset);
                            }
                        }

                        // SrcData to DstData conversion
                        DstData p_dst_long_vector[long_vector_size];

                        for(index_t i = 0; i < long_vector_size; ++i)
                        {
                            p_dst_long_vector[i] = type_convert<DstData>{}(p_src_long_vector[i]);
                        }

                        // store data from the long-vector buffer to dst
                        for(index_t i = 0; i < long_vector_size / dst_data_per_access; ++i)
                        {
                            auto scalar_id               = make_zero_array<index_t, nDim>();
                            scalar_id(vector_access_dim) = i * dst_data_per_access;

                            const index_t buffer_offset = i * dst_data_per_access;

                            // move dst cooridnate along linear dimensions
                            const auto dst_coord =
                                dst_nonlinear_coord + (linear_dim_data_steps + scalar_id);

#if CK_EXPERIMENTAL_TENSOR_COORDINATE_USE_CALCULATE_OFFSET_DIFF // tweaking
                            // this is dst compile-time offset
                            const index_t dst_linear_offset =
                                dst_nonlinear_coord.CalculateOffsetDiff(linear_dim_data_steps +
                                                                        scalar_id);
#else
                            // this is dst compile-time offset
                            const index_t dst_linear_offset =
                                dst_coord.GetOffset() - dst_nonlinear_coord.GetOffset();
#endif

                            // Check dst data's valid mapping situation, only check the first data
                            // in this dst
                            //   vector. It's user's responsiblity to make sure all data in the dst
                            //   vector has the valid/invalid mapping situation
                            if(dst_coord.IsOffsetValidAssumingUpperIndexIsValid())
                            {
                                transfer_data<DstData,
                                              DstDataPerWrite,
                                              AddressSpace::Vgpr,
                                              DstAddressSpace,
                                              DstInMemOp>(p_dst_long_vector,
                                                          buffer_offset,
                                                          p_dst,
                                                          dst_nonlinear_coord.GetOffset() +
                                                              dst_linear_offset);
                            }
                        }
                    });
            });
    }

    __device__ static constexpr bool HasWorkingOptimizedAddressCalculation()
    {
#if CK_EXPERIMENTAL_THREADWISE_COPY_V4R2_USE_OPTIMIZED_ADDRESS_CACLULATION // tweaking
        return true;
#else
        return false;
#endif
    }

    template <typename T, bool PositiveDirection>
    __device__ void MoveSrcSliceWindow(const T& step_sizes_,
                                       integral_constant<bool, PositiveDirection>)
    {
        const auto step_sizes = to_array(step_sizes_);

        static_if<PositiveDirection>{}([&](auto) {
            mSrcSliceOrigin += to_array(step_sizes);
        }).Else([&](auto) { mSrcSliceOrigin -= step_sizes; });
    }

    template <typename T, bool PositiveDirection>
    __device__ void MoveDstSliceWindow(const T& step_sizes_,
                                       integral_constant<bool, PositiveDirection>)
    {
        const auto step_sizes = to_array(step_sizes_);

        static_if<PositiveDirection>{}([&](auto) {
            mDstSliceOrigin += step_sizes;
        }).Else([&](auto) { mDstSliceOrigin -= step_sizes; });
    }

private:
    SrcCoord mSrcSliceOrigin;
    DstCoord mDstSliceOrigin;
};

} // namespace ck
#endif
