// coding=utf-8
//
// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation
// and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "disco.h"
#include "disco_cuda.cuh"

template <int BDIM_X, int ELXTH, typename REAL_T>
__device__ void disco_bwd_d(const int Hi, const int Wi, const int K, const int Ho, const int Wo, const int pscale,
                            const int64_t *__restrict__ roff, const int64_t *__restrict__ kers,
                            const int64_t *__restrict__ rows, const int64_t *__restrict__ cols,
                            const REAL_T *__restrict__ vals, const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
{

    const int tid = threadIdx.x;

    const int64_t bidx = blockIdx.x; // gloabl row
    const int64_t bidy = blockIdx.y; // bc

    int64_t soff = roff[bidx];
    int64_t eoff = roff[bidx + 1];

    const int64_t ker = kers[soff];
    const int64_t row = rows[soff];

    inp += bidy * K * Hi * Wi + ker * Hi * Wi + row * Wi;
    out += bidy * Ho * Wo;

    // align to larger supported fp type
    extern __shared__ __align__(sizeof(double)) unsigned char __sh_ptr[]; // REAL_T __sh[2*(BDIM_X*ELXTH)*pscale]

    REAL_T(*__sh)[BDIM_X * ELXTH * 2] = reinterpret_cast<REAL_T(*)[BDIM_X * ELXTH * 2]>(__sh_ptr);

    // copy current inp row in regs
    REAL_T __reg[ELXTH];

#pragma unroll
    for (int i = 0; i < ELXTH; i++) { __reg[i] = (i * BDIM_X + tid < Wi) ? inp[i * BDIM_X + tid] : REAL_T(0); }

    // reset shared row up to Wo+2, remaining
    // ppscale*(BDIM_X*ELXTH - Wo) locations
    // will be written to but never copied to
    // global mem
    for (int i = 0; i < pscale; i++) {
#pragma unroll
        for (int j = 0; j < 2 * BDIM_X * ELXTH; j += BDIM_X) { __sh[i][j + tid] = 0; }
    }
    __syncthreads();

    int col_prev = cols[soff];

    int h_prev = col_prev / Wo;
    int w_prev = col_prev % Wo;

    // loops along the colums of CTA's row
    for (int64_t nz = soff; nz < eoff; nz++) {

        const int col = cols[nz];
        const REAL_T val = vals[nz];

        // if we are processing a nz with a col value
        // leading to a new row of inp then copy it
        // to shmem;
        // we read a col that points to a new output
        // row if (col / Wo) > (col_prev / Wo)
        if (col >= col_prev - w_prev + Wo) {
            __syncthreads();
            for (int i = 0; i < pscale; i++) {
                for (int j = tid; j < Wi; j += BDIM_X) {

                    const REAL_T v = __sh[i][j] + __sh[i][Wi + j];

                    atomicAdd(&out[h_prev * Wo + j * pscale + i], v);

                    __sh[i][j] = 0;
                    __sh[i][Wi + j] = 0;
                }
            }
            __syncthreads();

            col_prev = col;
            h_prev = col / Wo;
            w_prev = col % Wo;
        }

        const int w = w_prev + (col - col_prev);
        const int w_mod_ps = w % pscale;
        const int w_div_ps = w / pscale;

#pragma unroll
        for (int i = 0; i < ELXTH; i++) {

            const int pp = i * BDIM_X + tid;
            __sh[w_mod_ps][w_div_ps + pp] += val * __reg[i];
        }

        // to avoid race conditions on __sh[]
        // among consecutive iterations along nz
        __syncthreads();
    }
    __syncthreads();

    // write last row
    for (int i = 0; i < pscale; i++) {

        for (int j = tid; j < Wi; j += BDIM_X) {

            const REAL_T v = __sh[i][j] + __sh[i][Wi + j];
            atomicAdd(&out[h_prev * Wo + j * pscale + i], v);
        }
    }
    return;
}

template <int BDIM_X, int ELXTH, int PSCALE, typename REAL_T>
__global__
    __launch_bounds__(BDIM_X) void disco_bwd_blk_k(const int Hi, const int Wi, const int K, const int Ho, const int Wo,
                                                   const int pscale, const int64_t *__restrict__ roff,
                                                   const int64_t *__restrict__ kers, const int64_t *__restrict__ rows,
                                                   const int64_t *__restrict__ cols, const REAL_T *__restrict__ vals,
                                                   const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
{

    if constexpr (PSCALE != 0) {
        disco_bwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, PSCALE, roff, kers, rows, cols, vals, inp, out);
    } else {
        disco_bwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, pscale, roff, kers, rows, cols, vals, inp, out);
    }

    return;
}

template <int NTH, int ELXTH, typename REAL_T>
static void launch_kernel(int BC, int Hi, int Wi, int K, int Ho, int Wo, int64_t nrows, int64_t *roff_d, int64_t *ker_d,
                          int64_t *row_d, int64_t *col_d, REAL_T *val_d, REAL_T *inp_d, REAL_T *out_d,
                          cudaStream_t stream)
{

    static_assert(sizeof(REAL_T) == 2 || sizeof(REAL_T) == 4 || sizeof(REAL_T) == 8);

    if constexpr (ELXTH <= ELXTH_MAX) {
        if (NTH * ELXTH >= Wi) {
            dim3 grid(nrows, BC);

            const int pscale = Wo / Wi;
            size_t shmem = sizeof(*out_d) * (2 * (NTH * ELXTH) * pscale);

            switch (pscale) {
            case 1:
                disco_bwd_blk_k<NTH, ELXTH, 1><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
                                                                             row_d, col_d, val_d, inp_d, out_d);
                break;
            case 2:
                disco_bwd_blk_k<NTH, ELXTH, 2><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
                                                                             row_d, col_d, val_d, inp_d, out_d);
                break;
            case 3:
                disco_bwd_blk_k<NTH, ELXTH, 3><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
                                                                             row_d, col_d, val_d, inp_d, out_d);
                break;
            default:
                disco_bwd_blk_k<NTH, ELXTH, 0><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
                                                                             row_d, col_d, val_d, inp_d, out_d);
            }
        } else {
            launch_kernel<NTH, ELXTH + 1>(BC, Hi, Wi, K, Ho, Wo, nrows, roff_d, ker_d, row_d, col_d, val_d, inp_d,
                                          out_d, stream);
        }
    }
    return;
}

torch::Tensor disco_cuda_bwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo)
{

    // some sanity checks
    CHECK_CUDA_INPUT_TENSOR(inp);
    CHECK_CUDA_INPUT_TENSOR(roff_idx);
    CHECK_CUDA_INPUT_TENSOR(ker_idx);
    CHECK_CUDA_INPUT_TENSOR(row_idx);
    CHECK_CUDA_INPUT_TENSOR(col_idx);
    CHECK_CUDA_INPUT_TENSOR(val);

    // extract some shapes
    int64_t B = inp.size(0);
    int64_t C = inp.size(1);
    int64_t BC = B * C;
    int64_t Hi = inp.size(3);
    int64_t Wi = inp.size(4);
    int64_t nrows = roff_idx.size(0) - 1;

    // allocate output
    int64_t out_dims[] = {B, C, Ho, Wo};
    auto options = torch::TensorOptions().device(inp.device()).dtype(inp.dtype());
    torch::Tensor out = torch::zeros(out_dims, options);

    // get stream
    auto stream = at::cuda::getCurrentCUDAStream().stream();

    // assert
    static_assert(0 == (ELXTH_MAX % 2));

    if (Wo <= 64 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
                                       launch_kernel<64, 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 128 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
                                       launch_kernel<128, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 256 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
                                       launch_kernel<256, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 512 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
                                       launch_kernel<512, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else if (Wo <= 1024 * ELXTH_MAX) {
        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
                                       launch_kernel<1024, (ELXTH_MAX / 2) + 1, scalar_t>(
                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
                                   }));
    } else {
        fprintf(stderr, "%s:%d: error, unsupported Wo value (%ld), max supported is %d\n", __FILE__, __LINE__, Wo,
                1024 * ELXTH_MAX);
        exit(EXIT_FAILURE);
    }

    return out;
}

// PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
//  m.def("backward", &disco_cuda_bwd, "DISCO backward (CUDA)");
//}
