/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#include <cugraph-ops/operators/common.hpp>

#include <algorithm>
#include <cstddef>
#include <cstring>
#include <initializer_list>
#include <limits>
#include <memory>
#include <vector>

namespace cugraph::ops::update_efeat_csc {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_kernel(DataT* out,
                          const DataT* edge_feat,
                          const DataT* src_feat,
                          const DataT* dst_feat,
                          size_t dim_edge,
                          size_t dim_src,
                          size_t dim_dst,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          size_t n_dst_nodes,
                          size_t n_indices,
                          bool add_features)
{
  // for add, we implicitly assume dim_edge + dim_src + dim_dst
  auto out_stride =
    add_features ? std::max({dim_edge, dim_src, dim_dst}) : (dim_edge + dim_src + dim_dst);
  auto acc = std::make_unique<DataT[]>(out_stride);
  std::memset(out, 0, n_indices * out_stride * sizeof(DataT));

  for (auto dst_node_id = 0; dst_node_id < n_dst_nodes; ++dst_node_id) {
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];

    // loop over neighbors
    for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
      auto neigh = indices[neigh_off];
      auto eid   = map_csc_to_coo == nullptr ? neigh_off : map_csc_to_coo[neigh_off];

      // reset accumulators
      if (add_features) {
        for (size_t d = 0; d < out_stride; ++d) {
          acc[d] = DataT{0};

          if (edge_feat != nullptr) { acc[d] += edge_feat[eid * dim_edge + d]; }
          if (src_feat != nullptr) { acc[d] += src_feat[neigh * dim_src + d]; }
          if (dst_feat != nullptr) { acc[d] += dst_feat[dst_node_id * dim_dst + d]; }
        }
      } else {
        if (edge_feat != nullptr) {
          for (size_t d = 0; d < dim_edge; ++d) {
            acc[d] = edge_feat[eid * dim_edge + d];
          }
        }

        if (src_feat != nullptr) {
          for (size_t d = 0; d < dim_src; ++d) {
            acc[dim_edge + d] = src_feat[neigh * dim_src + d];
          }
        }

        if (dst_feat != nullptr) {
          for (size_t d = 0; d < dim_dst; ++d) {
            acc[dim_edge + dim_src + d] = dst_feat[dst_node_id * dim_dst + d];
          }
        }
      }

      // output aggregated or concatenated values
      for (size_t d = 0; d < out_stride; ++d) {
        out[eid * out_stride + d] += acc[d];
      }
    }
  }
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_kernel(DataT* grad_edge_feat,
                          DataT* grad_src_feat,
                          DataT* grad_dst_feat,
                          const DataT* grad_out,
                          size_t dim_edge,
                          size_t dim_src,
                          size_t dim_dst,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          size_t n_src_nodes,
                          size_t n_dst_nodes,
                          size_t n_edges,
                          bool add_features)
{
  if (grad_edge_feat != nullptr) {
    std::memset(grad_edge_feat, 0, n_edges * dim_edge * sizeof(DataT));
  }
  if (grad_src_feat != nullptr) {
    std::memset(grad_src_feat, 0, n_src_nodes * dim_src * sizeof(DataT));
  }
  if (grad_dst_feat != nullptr) {
    std::memset(grad_dst_feat, 0, n_dst_nodes * dim_dst * sizeof(DataT));
  }

  auto out_stride =
    add_features ? std::max({dim_edge, dim_src, dim_dst}) : (dim_edge + dim_src + dim_dst);

  for (auto dst_node_id = 0; dst_node_id < n_dst_nodes; ++dst_node_id) {
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];

    // loop over neighbors
    for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
      auto neigh = indices[neigh_off];
      auto eid   = map_csc_to_coo == nullptr ? neigh_off : map_csc_to_coo[neigh_off];

      DataT grad;

      if (add_features) {
        for (size_t d = 0; d < out_stride; ++d) {
          grad = grad_out[eid * out_stride + d];
          if (grad_edge_feat != nullptr) { grad_edge_feat[eid * dim_edge + d] += grad; }
          if (grad_src_feat != nullptr) { grad_src_feat[neigh * dim_src + d] += grad; }
          if (grad_dst_feat != nullptr) { grad_dst_feat[dst_node_id * dim_dst + d] += grad; }
        }
      } else {
        // get gradient for edge_feat
        if (grad_edge_feat != nullptr) {
          for (size_t d = 0; d < dim_edge; ++d) {
            grad = grad_out[eid * out_stride + d];
            grad_edge_feat[eid * dim_edge + d] += grad;
          }
        }

        // get gradient for src_feat
        if (grad_src_feat != nullptr) {
          for (size_t d = 0; d < dim_src; ++d) {
            grad = grad_out[eid * out_stride + dim_edge + d];
            grad_src_feat[neigh * dim_src + d] += grad;
          }
        }

        // get gradient for dst_feat
        if (grad_dst_feat != nullptr) {
          for (size_t d = 0; d < dim_dst; ++d) {
            grad = grad_out[eid * out_stride + dim_edge + dim_src + d];
            grad_dst_feat[dst_node_id * dim_dst + d] += grad;
          }
        }
      }
    }
  }
}

}  // namespace cugraph::ops::update_efeat_csc
