#pragma once
#include "host_tensor.hpp"

template <typename TIn,
          typename TWei,
          typename TOut,
          typename ConvStrides,
          typename ConvDilations,
          typename InLeftPads,
          typename InRightPads>
void host_direct_convolution_backward_data(Tensor<TIn>& in,
                                           const Tensor<TWei>& wei,
                                           const Tensor<TOut>& out,
                                           const ConvStrides& conv_strides,
                                           const ConvDilations& conv_dilations,
                                           const InLeftPads& in_left_pads,
                                           const InRightPads& ,
                                           const ConvTensorLayout layout = ConvTensorLayout::NCHW)
{
    using namespace ck;

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

    auto loop_pattern = [&](const int stride, 
                            const int dilation, 
                            const int pad, 
                            const int Y, 
                            const int Ho, 
                            const int hi, 
                            int& yBegin, 
                            int& yEnd, 
                            int& dy,
                            int& hoBegin,
                            int& hoEnd,
                            int& dho) {

        const int gcd = std::gcd(dilation, stride);
        dy = stride / gcd; 
        dho = -dilation / gcd;
        yBegin = Y;
        yEnd = 0;
        hoBegin = 0;
        hoEnd = Ho;

        for (int y=0; y < Y; ++y)
        {
            int h_tmp = hi + pad - y * dilation;
            int ho = h_tmp / stride;
            if (h_tmp % stride == 0 && ho >= 0 && ho < Ho)
            {
                yBegin = y;
                hoBegin = ho;
            }
        }

        for (int y=Y-1; y>=0; --y)
        {
            int h_tmp = hi + pad - y * dilation;
            int ho = h_tmp / stride;
            if (h_tmp % stride == 0 && ho >= 0 && ho < Ho)
            {
                yEnd = y+1;
                hoEnd = ho;
            }
        }
    };

    auto f_nchw = [&](auto n, auto c, auto hi, auto wi) {
        std::size_t K = wei.mDesc.GetLengths()[I0];
        std::size_t Y = wei.mDesc.GetLengths()[I2];
        std::size_t X = wei.mDesc.GetLengths()[I3];

        std::size_t Ho = out.mDesc.GetLengths()[I2];
        std::size_t Wo = out.mDesc.GetLengths()[I3];

        int yBegin, yEnd, dy;
        int xBegin, xEnd, dx;
        int hoBegin, hoEnd, dho;
        int woBegin, woEnd, dwo;
        loop_pattern(conv_strides[I0], conv_dilations[I0], in_left_pads[I0], Y, Ho, hi, yBegin, yEnd, dy, hoBegin, hoEnd, dho);
        loop_pattern(conv_strides[I1], conv_dilations[I1], in_left_pads[I1], X, Wo, wi, xBegin, xEnd, dx, woBegin, woEnd, dwo);

        double v = 0;

#if 0
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int ho = (hi + in_left_pads[I0] - y * conv_dilations[I0]) / conv_strides[I0];

            for(int x = xBegin; x < xEnd; x+=dx)
            {
                int wo = (wi + in_left_pads[I1] - x * conv_dilations[I1]) / conv_strides[I1];

                for(int k = 0; k < K; ++k)
                {
                    v += out(n, k, ho, wo) * wei(k, c, y, x);
                }
            }
        }
#else
        int ho = hoBegin;
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int wo = woBegin;
            for(int x = xBegin; x < xEnd; x+=dx)
            {
                for(int k = 0; k < K; ++k)
                {
                    v += out(n, k, ho, wo) * wei(k, c, y, x);
                }
                wo += dwo;
            }
            ho += dho;
        }
#endif

        in(n, c, hi, wi) = v;
    };

    auto f_nhwc = [&](auto n, auto c, auto hi, auto wi) {
        std::size_t K = wei.mDesc.GetLengths()[I0];
        std::size_t Y = wei.mDesc.GetLengths()[I1];
        std::size_t X = wei.mDesc.GetLengths()[I2];

        std::size_t Ho = out.mDesc.GetLengths()[I1];
        std::size_t Wo = out.mDesc.GetLengths()[I2];

        int yBegin, yEnd, dy;
        int xBegin, xEnd, dx;
        int hoBegin, hoEnd, dho;
        int woBegin, woEnd, dwo;
        loop_pattern(conv_strides[I0], conv_dilations[I0], in_left_pads[I0], Y, Ho, hi, yBegin, yEnd, dy, hoBegin, hoEnd, dho);
        loop_pattern(conv_strides[I1], conv_dilations[I1], in_left_pads[I1], X, Wo, wi, xBegin, xEnd, dx, woBegin, woEnd, dwo);
        double v = 0;

#if 0
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int ho = (hi + in_left_pads[I0] - y * conv_dilations[I0]) / conv_strides[I0];

            for(int x = xBegin; x < xEnd; x+=dx)
            {
                int wo = (wi + in_left_pads[I1] - x * conv_dilations[I1]) / conv_strides[I1];

                for(int k = 0; k < K; ++k)
                {
                    v += out(n, ho, wo, k) * wei(k, y, x, c);
                }
            }
        }

#else
        int ho = hoBegin;
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int wo = woBegin;
            for(int x = xBegin; x < xEnd; x+=dx)
            {
                for(int k = 0; k < K; ++k)
                {
                    v += out(n, ho, wo, k) * wei(k, y, x, c);
                }
                wo += dwo;
            }
            ho += dho;
        }
#endif

        in(n, hi, wi, c) = v;
    };

    if(layout == ConvTensorLayout::NCHW)
    {
        make_ParallelTensorFunctor(f_nchw,
                                   in.mDesc.GetLengths()[0],
                                   in.mDesc.GetLengths()[1],
                                   in.mDesc.GetLengths()[2],
                                   in.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    }
    else if(layout == ConvTensorLayout::NHWC)
    {
        make_ParallelTensorFunctor(f_nhwc,
                                   in.mDesc.GetLengths()[0],
                                   in.mDesc.GetLengths()[1],
                                   in.mDesc.GetLengths()[2],
                                   in.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    }
    else
    {
        throw std::runtime_error("wrong! not supported layout");
    }
}
#pragma once
#include "host_tensor.hpp"

template <typename TIn,
          typename TWei,
          typename TOut,
          typename ConvStrides,
          typename ConvDilations,
          typename InLeftPads,
          typename InRightPads>
void host_direct_convolution_backward_data(Tensor<TIn>& in,
                                           const Tensor<TWei>& wei,
                                           const Tensor<TOut>& out,
                                           const ConvStrides& conv_strides,
                                           const ConvDilations& conv_dilations,
                                           const InLeftPads& in_left_pads,
                                           const InRightPads& /* in_right_pads */,
                                           const ConvTensorLayout layout = ConvTensorLayout::NCHW)
{
    using namespace ck;

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

    /*
     * Find the loop pattern of (y, ho) or (x, wo).
     */
    auto loop_pattern = [&](const int stride, 
                            const int dilation, 
                            const int pad, 
                            const int Y, 
                            const int Ho, 
                            const int hi, 
                            int& yBegin, 
                            int& yEnd, 
                            int& dy,
                            int& hoBegin,
                            int& hoEnd,
                            int& dho) {
        // Find those y's such that h_tmp % stride == 0, 
        // where 
        //      h_tmp = hi + pad - y * dilation. 
        // If yBegin is the minimal y that satisfies this condition, then the next one is
        //      yBegin + dy
        // where 
        //      dy = stride / greatest_common_divisor(stride, dilation) = stride / gcd.
        //
        // This can be explained as follows:
        //      hi + pads0 - (yBegin+dy) * dilation = (hi + pads0 - yBegin * dilation)- (dy * dilation). 
        // The first part is already a multiple of stride, in order for (yBegin+dy) to satisfy the condition,
        //      dy * dilation = k * stride  <===> dilation / stride = k / dy.
        // Since gcd is the greatest common divisor of stride and dilation, we have 
        //      m * gcd = dilation,
        //      n * gcd = stride
        // and dilation / stride = m / n is a irreducible fraction. Then k/dy = m/n, and the minimal dy
        // we can have is dy = n = stride / gcd.
        //
        // Another way to find out dy is to loop over y in [0, Y) to find out yBegin and dy+yBegin.
        //
        //
        // We can further find the corresponding hoBegin, dho, hoEnd, which can be used to avoid the division operation.
        // Let hoBegin = (hi + pad - yBegin * dilation) / stride, then
        // hoNext = hoBegin + dho 
        //        = (hi + pad - (yBegin + dy) * dilation) / stride
        //        = hoBegin - dy * dilation / stride
        //
        // that is, dho = dy * dilation / stride = -dilation/gcd; 

        // loop pattern of y
        const int gcd = std::gcd(dilation, stride);
        dy = stride / gcd; 
        dho = -dilation / gcd;
        yBegin = Y;
        yEnd = 0;
        hoBegin = 0;
        hoEnd = Ho;

        // I believe these 2 loops can be replaced with 
        //  yBegin = max(0, (hi + pad - (Ho-1) * stride + dilation -1) / dilation);
        //  yEnd = min(Y, (hi + pad) / dilation + 1);
        //  
        // This can be explained as follows. 
        // Since 
        //      hi + pad - y * dilation = ho * stride, 
        // and 
        //      0 <= ho <= Ho-1
        // we have 
        //      (hi + pad - (Ho-1) * stride) <= y * dilation <= hi + pad
        // This reduces to
        //      (hi + pad - (Ho-1) * stride + dilation - 1) / dilation <= y <= (hi + pad) / dilation.
        // We also have 0 <= y < Y, combining the 2 inequalities gives rise to the range [yBegin, yEnd).
        //
        for (int y=0; y < Y; ++y)
        {
            int h_tmp = hi + pad - y * dilation;
            int ho = h_tmp / stride;
            if (h_tmp % stride == 0 && ho >= 0 && ho < Ho)
            {
                yBegin = y;
                hoBegin = ho;
            }
        }


        for (int y=Y-1; y>=0; --y)
        {
            int h_tmp = hi + pad - y * dilation;
            int ho = h_tmp / stride;
            if (h_tmp % stride == 0 && ho >= 0 && ho < Ho)
            {
                yEnd = y+1;
                hoEnd = ho;
            }
        }
    };


    auto f_nchw = [&](auto n, auto c, auto hi, auto wi) {
        std::size_t K = wei.mDesc.GetLengths()[I0];
        std::size_t Y = wei.mDesc.GetLengths()[I2];
        std::size_t X = wei.mDesc.GetLengths()[I3];

        std::size_t Ho = out.mDesc.GetLengths()[I2];
        std::size_t Wo = out.mDesc.GetLengths()[I3];

        int yBegin, yEnd, dy;
        int xBegin, xEnd, dx;
        int hoBegin, hoEnd, dho;
        int woBegin, woEnd, dwo;
        loop_pattern(conv_strides[I0], conv_dilations[I0], in_left_pads[I0], Y, Ho, hi, yBegin, yEnd, dy, hoBegin, hoEnd, dho);
        loop_pattern(conv_strides[I1], conv_dilations[I1], in_left_pads[I1], X, Wo, wi, xBegin, xEnd, dx, woBegin, woEnd, dwo);

        double v = 0;
        
#if 0
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int ho = (hi + in_left_pads[I0] - y * conv_dilations[I0]) / conv_strides[I0];

            for(int x = xBegin; x < xEnd; x+=dx)
            {
                int wo = (wi + in_left_pads[I1] - x * conv_dilations[I1]) / conv_strides[I1];

                for(int k = 0; k < K; ++k)
                {
                    v += out(n, k, ho, wo) * wei(k, c, y, x);
                }
            }
        }
#else
        int ho = hoBegin;
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int wo = woBegin;
            for(int x = xBegin; x < xEnd; x+=dx)
            {
                for(int k = 0; k < K; ++k)
                {
                    v += out(n, k, ho, wo) * wei(k, c, y, x);
                }
                wo += dwo;
            }
            ho += dho;
        }
#endif

        in(n, c, hi, wi) = v;
    };

    auto f_nhwc = [&](auto n, auto c, auto hi, auto wi) {
        std::size_t K = wei.mDesc.GetLengths()[I0];
        std::size_t Y = wei.mDesc.GetLengths()[I1];
        std::size_t X = wei.mDesc.GetLengths()[I2];

        std::size_t Ho = out.mDesc.GetLengths()[I1];
        std::size_t Wo = out.mDesc.GetLengths()[I2];

        int yBegin, yEnd, dy;
        int xBegin, xEnd, dx;
        int hoBegin, hoEnd, dho;
        int woBegin, woEnd, dwo;
        loop_pattern(conv_strides[I0], conv_dilations[I0], in_left_pads[I0], Y, Ho, hi, yBegin, yEnd, dy, hoBegin, hoEnd, dho);
        loop_pattern(conv_strides[I1], conv_dilations[I1], in_left_pads[I1], X, Wo, wi, xBegin, xEnd, dx, woBegin, woEnd, dwo);
        double v = 0;

#if 0
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int ho = (hi + in_left_pads[I0] - y * conv_dilations[I0]) / conv_strides[I0];

            for(int x = xBegin; x < xEnd; x+=dx)
            {
                int wo = (wi + in_left_pads[I1] - x * conv_dilations[I1]) / conv_strides[I1];

                for(int k = 0; k < K; ++k)
                {
                    v += out(n, ho, wo, k) * wei(k, y, x, c);
                }
            }
        }

#else
        int ho = hoBegin;
        for (int y=yBegin; y<yEnd; y+=dy)
        {
            int wo = woBegin;
            for(int x = xBegin; x < xEnd; x+=dx)
            {
                for(int k = 0; k < K; ++k)
                {
                    v += out(n, ho, wo, k) * wei(k, y, x, c);
                }
                wo += dwo;
            }
            ho += dho;
        }
#endif

        in(n, hi, wi, c) = v;
    };

    if(layout == ConvTensorLayout::NCHW)
    {
        make_ParallelTensorFunctor(f_nchw,
                                   in.mDesc.GetLengths()[0],
                                   in.mDesc.GetLengths()[1],
                                   in.mDesc.GetLengths()[2],
                                   in.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    }
    else if(layout == ConvTensorLayout::NHWC)
    {
        make_ParallelTensorFunctor(f_nhwc,
                                   in.mDesc.GetLengths()[0],
                                   in.mDesc.GetLengths()[1],
                                   in.mDesc.GetLengths()[2],
                                   in.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    }
    else
    {
        throw std::runtime_error("wrong! not supported layout");
    }
}
