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

namespace cugraph::ops::agg_naive_csc {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_kernel(DataT* out,
                          IdxT* out_pos,
                          const DataT* node_feat,
                          const DataT* concat_feat,
                          const DataT* edge_feat,
                          size_t dim_node,
                          size_t dim_concat,
                          size_t dim_edge,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          size_t n_dst_nodes,
                          AggOpT agg_op)
{
  // if dim_node > 0, we expect node_feat implicitly and use them for the aggregation
  // if dim_edge > 0, we expect edge_feat implicitly and use them for the aggregation
  // if dim_concat > 0, we expect concat_feat implicitly and concatenate these
  // with the result of the aggregation
  DataT identity;
  if (agg_op == AggOpT::kMean || agg_op == AggOpT::kSum) {
    identity = DataT{0};
  } else if (agg_op == AggOpT::kMax) {
    identity = -std::numeric_limits<DataT>::infinity();
  } else if (agg_op == AggOpT::kMin) {
    identity = std::numeric_limits<DataT>::infinity();
  }
  auto feat_stride = dim_node + dim_edge;
  auto out_stride  = feat_stride + dim_concat;
  // this is the aggregation
  auto acc = std::make_unique<DataT[]>(feat_stride);
  auto pos = std::make_unique<IdxT[]>(feat_stride);

  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];
    auto count = off_end - off_start;

    // concatenate own features
    if (dim_concat > 0) {
      for (size_t d = 0; d < dim_concat; ++d) {
        out[dst_node_id * out_stride + feat_stride + d] = concat_feat[dst_node_id * dim_concat + d];
      }
    }

    // reset accumulators
    for (size_t d = 0; d < feat_stride; ++d) {
      acc[d] = identity;
      pos[d] = graph::INVALID_ID<IdxT>;
    }

    // aggregate over neighbors
    for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
      // aggregate node features
      if (dim_node > 0) {
        auto neigh = indices[neigh_off];
        for (size_t d = 0; d < dim_node; ++d) {
          auto val = node_feat[neigh * dim_node + d];
          if (agg_op == AggOpT::kSum) {
            acc[d] += val;
          } else if (agg_op == AggOpT::kMean && count > 0) {
            acc[d] += val / static_cast<DataT>(count);
          } else if (agg_op == AggOpT::kMax) {
            if (val > acc[d]) {
              acc[d] = val;
              pos[d] = neigh;
            }
          } else if (agg_op == AggOpT::kMin) {
            if (val < acc[d]) {
              acc[d] = val;
              pos[d] = neigh;
            }
          }
        }
      }

      if (dim_edge > 0) {
        auto eid = map_csc_to_coo == nullptr ? neigh_off : map_csc_to_coo[neigh_off];
        // aggregate edge features
        for (size_t d = 0; d < dim_edge; ++d) {
          auto val = edge_feat[eid * dim_edge + d];
          if (agg_op == AggOpT::kSum) {
            acc[dim_node + d] += val;
          } else if (agg_op == AggOpT::kMean && count > 0) {
            acc[dim_node + d] += val / static_cast<DataT>(count);
          } else if (agg_op == AggOpT::kMax) {
            if (val > acc[dim_node + d]) {
              acc[dim_node + d] = val;
              pos[dim_node + d] = eid;
            }
          } else if (agg_op == AggOpT::kMin) {
            if (val < acc[dim_node + d]) {
              acc[dim_node + d] = val;
              pos[dim_node + d] = eid;
            }
          }
        }
      }
    }

    // output aggregated values
    for (size_t d = 0; d < feat_stride; ++d) {
      out[dst_node_id * out_stride + d] = (acc[d] == identity) ? DataT{0} : acc[d];
      if (agg_op == AggOpT::kMin || agg_op == AggOpT::kMax) {
        // out_pos only feat_stride
        if (out_pos != nullptr) { out_pos[dst_node_id * feat_stride + d] = pos[d]; }
      }
    }
  }
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_kernel(DataT* d_node_feat,
                          DataT* d_concat_feat,
                          DataT* d_edge_feat,
                          const DataT* dout,
                          const IdxT* out_pos,
                          size_t dim_node,
                          size_t dim_concat,
                          size_t dim_edge,
                          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,
                          AggOpT agg_op)
{
  // if dim_node > 0, we expect node_feat implicitly and use them for the aggregation
  // if dim_edge > 0, we expect edge_feat implicitly and use them for the aggregation
  // if dim_concat > 0, we expect concat_feat implicitly and concatenate these
  // with the result of the aggregation

  // first set all the grads to 0
  if (d_node_feat != nullptr) {
    std::memset(d_node_feat, 0, n_src_nodes * dim_node * sizeof(DataT));
  }
  if (d_concat_feat != nullptr) {
    std::memset(d_concat_feat, 0, n_dst_nodes * dim_concat * sizeof(DataT));
  }
  if ((d_edge_feat != nullptr) && ((agg_op == AggOpT::kMax) || (agg_op == AggOpT::kMin))) {
    std::memset(d_edge_feat, 0, n_edges * dim_edge * sizeof(DataT));
  }

  auto feat_stride = dim_node + dim_edge;
  auto out_stride  = feat_stride + dim_concat;
  // iterate over output nodes in batch
  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];
    auto count = off_end - off_start;
    auto norm  = count > 0 ? DataT{1} / static_cast<DataT>(count) : DataT{-1};
    // iterate over dimension. for min/max, can directly output
    // for sum/mean, iterate over neighbors as well
    if (agg_op == AggOpT::kMax || agg_op == AggOpT::kMin) {
      // node features
      if (d_node_feat != nullptr) {
        for (size_t d = 0; d < dim_node; ++d) {
          auto grad = dout[dst_node_id * out_stride + d];
          // don't pass gradient on if node does not have incoming edges
          auto nid = out_pos[dst_node_id * feat_stride + d];
          if (nid != graph::INVALID_ID<IdxT>) d_node_feat[nid * dim_node + d] += grad;
        }
      }
      // edge features
      if (d_edge_feat != nullptr) {
        for (size_t d = 0; d < dim_edge; ++d) {
          auto grad = dout[dst_node_id * out_stride + dim_node + d];
          // don't pass gradient on if node does not have incoming edges
          auto eid = out_pos[dst_node_id * feat_stride + dim_node + d];
          if (eid != graph::INVALID_ID<IdxT>) d_edge_feat[eid * dim_edge + d] = grad;
        }
      }
    } else {
      // node features
      if (d_node_feat != nullptr) {
        for (size_t d = 0; d < dim_node; ++d) {
          auto grad = dout[dst_node_id * out_stride + d];
          if (agg_op == AggOpT::kMean) grad *= norm;
          for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
            auto nid = indices[neigh_off];
            d_node_feat[nid * dim_node + d] += grad;
          }
        }
      }
      // edge features
      if (d_edge_feat != nullptr) {
        for (size_t d = 0; d < dim_edge; ++d) {
          auto grad = dout[dst_node_id * out_stride + dim_node + d];
          if (agg_op == AggOpT::kMean) grad *= norm;
          for (auto neigh = off_start; neigh < off_end; ++neigh) {
            auto eid = map_csc_to_coo == nullptr ? neigh : map_csc_to_coo[neigh];
            d_edge_feat[eid * dim_edge + d] = grad;
          }
        }
      }
    }
    if ((dim_concat > 0) && (d_concat_feat != nullptr)) {
      for (size_t d = 0; d < dim_concat; ++d) {
        d_concat_feat[dst_node_id * dim_concat + d] +=
          dout[dst_node_id * out_stride + feat_stride + d];
      }
    }
  }
}

}  // namespace cugraph::ops::agg_naive_csc
