/*
 * 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 <limits>
#include <memory>
#include <vector>

namespace cugraph::ops::dmpnn {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_kernel(DataT* out,
                          const DataT* edge_feat,
                          size_t dim_edge,
                          const IdxT* offsets,
                          const IdxT* /*indices*/,
                          const IdxT* map_csc_to_coo,
                          const IdxT* rev_map_csc_to_coo,
                          size_t n_nodes,
                          bool use_mean,
                          bool concat_own)
{
  auto out_stride = concat_own ? dim_edge * 2 : dim_edge;
  // this is the aggregation
  auto acc = std::make_unique<DataT[]>(dim_edge);

  for (auto dst_node_id = 0; dst_node_id < n_nodes; ++dst_node_id) {
    // reset accumulators
    for (size_t d = 0; d < dim_edge; ++d) {
      acc[d] = DataT{0};
    }
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];
    auto count = off_end - off_start;
    auto norm  = count > 0 ? DataT{1} / static_cast<DataT>(count) : DataT{-1};
    // aggregate over neighboring in-edges
    for (auto neigh = off_start; neigh < off_end; ++neigh) {
      auto eid     = map_csc_to_coo == nullptr ? neigh : map_csc_to_coo[neigh];
      auto eid_off = eid;
      for (size_t d = 0; d < dim_edge; ++d) {
        auto val = edge_feat[eid_off * dim_edge + d];
        if (use_mean) {
          acc[d] += val / count;
        } else {
          acc[d] += val;
        }
      }
    }
    // output over neighboring out-edges, and remove in-edge features again
    for (auto neigh = off_start; neigh < off_end; ++neigh) {
      auto eid        = map_csc_to_coo == nullptr ? neigh : map_csc_to_coo[neigh];
      auto eid_off    = eid;
      auto rev_neigh  = rev_map_csc_to_coo[neigh];
      auto rev_id     = map_csc_to_coo == nullptr ? rev_neigh : map_csc_to_coo[rev_neigh];
      auto rev_id_off = rev_id;
      for (size_t d = 0; d < dim_edge; ++d) {
        auto val                         = edge_feat[eid_off * dim_edge + d];
        out[rev_id_off * out_stride + d] = acc[d] - val;
        // concatenate own features if necessary
        if (concat_own)
          out[rev_id_off * out_stride + dim_edge + d] = edge_feat[rev_id_off * dim_edge + d];
      }
    }
  }
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_kernel(DataT* d_edge_feat,
                          const DataT* dout,
                          size_t dim_edge,
                          const IdxT* offsets,
                          const IdxT* /*indices*/,
                          const IdxT* map_csc_to_coo,
                          const IdxT* rev_map_csc_to_coo,
                          size_t n_nodes,
                          bool use_mean,
                          bool concat_own)
{
  auto in_stride = concat_own ? dim_edge * 2 : dim_edge;
  // this is the aggregation
  auto acc = std::make_unique<DataT[]>(dim_edge);
  for (auto dst_node_id = 0; dst_node_id < n_nodes; ++dst_node_id) {
    // reset accumulators
    for (size_t d = 0; d < dim_edge; ++d) {
      acc[d] = DataT{0};
    }
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];
    auto count = off_end - off_start;
    auto norm  = count > 0 ? DataT{1} / static_cast<DataT>(count) : DataT{-1};
    // aggregate over neighboring out-edges
    for (auto neigh = off_start; neigh < off_end; ++neigh) {
      auto rev_neigh  = rev_map_csc_to_coo[neigh];
      auto rev_id     = map_csc_to_coo == nullptr ? rev_neigh : map_csc_to_coo[rev_neigh];
      auto rev_id_off = rev_id;
      for (size_t d = 0; d < dim_edge; ++d) {
        auto val = dout[rev_id_off * in_stride + d];
        if (use_mean) {
          acc[d] += val / count;
        } else {
          acc[d] += val;
        }
      }
    }
    // output over neighboring in-edges, and remove out-edge features again
    for (auto neigh = off_start; neigh < off_end; ++neigh) {
      auto eid        = map_csc_to_coo == nullptr ? neigh : map_csc_to_coo[neigh];
      auto eid_off    = eid;
      auto rev_neigh  = rev_map_csc_to_coo[neigh];
      auto rev_id     = map_csc_to_coo == nullptr ? rev_neigh : map_csc_to_coo[rev_neigh];
      auto rev_id_off = rev_id;
      for (size_t d = 0; d < dim_edge; ++d) {
        auto val     = dout[rev_id_off * in_stride + d];
        auto d_final = acc[d] - val;
        // add gradient from "own" out-edge if necessary
        if (concat_own) d_final += dout[eid_off * in_stride + dim_edge + d];
        // output gradient to in-edge
        d_edge_feat[eid_off * dim_edge + d] = d_final;
      }
    }
  }
}

}  // namespace cugraph::ops::dmpnn
