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

#define DEG_EPS 1.0e-15

namespace cugraph::ops::agg_naive_csc {

template <typename DataT, typename IdxT>
void weighted_cpu_fwd_kernel(DataT* out,
                             IdxT* out_pos,
                             DataT* degree,
                             const DataT* node_feat,
                             const DataT* concat_feat,
                             const DataT* edge_weight,
                             size_t dim_node,
                             size_t dim_concat,
                             const IdxT* offsets,
                             const IdxT* indices,
                             size_t n_dst_nodes,
                             AggOpT agg_op)
{
  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();
  }
  bool concat_own = (dim_concat > 0);
  auto out_stride = dim_node + dim_concat;
  // this is the aggregation
  auto acc = std::make_unique<DataT[]>(dim_node);
  auto pos = std::make_unique<IdxT[]>(dim_node);

  for (auto dst_node_id = 0; dst_node_id < n_dst_nodes; ++dst_node_id) {
    // reset accumulators
    for (size_t d = 0; d < dim_node; ++d) {
      acc[d] = identity;
      pos[d] = graph::INVALID_ID<IdxT>;
    }
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];
    DataT norm   = DataT{0};
    DataT weight = DataT{0};
    // calculate norm if required
    if (agg_op == AggOpT::kMean) {
      for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
        weight = edge_weight[neigh_off];
        norm += weight;
      }
      degree[dst_node_id] = norm;
    }
    // aggregate over neighbors
    for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
      auto neigh = indices[neigh_off];
      weight     = edge_weight[neigh_off];
      for (size_t d = 0; d < dim_node; ++d) {
        auto val = node_feat[neigh * dim_node + d];
        val *= weight;
        if (agg_op == AggOpT::kSum || agg_op == AggOpT::kMean) {
          acc[d] += val;
        } else if (agg_op == AggOpT::kMax) {
          if (val > acc[d]) {
            acc[d] = val;
            pos[d] = neigh_off;
          }
        } else if (agg_op == AggOpT::kMin) {
          if (val < acc[d]) {
            acc[d] = val;
            pos[d] = neigh_off;
          }
        }
      }
    }
    if (agg_op == AggOpT::kMean) {
      for (size_t d = 0; d < dim_node; ++d) {
        acc[d] *= (1.0 / (norm + DEG_EPS));
      }
    }
    // output aggregated values
    for (size_t d = 0; d < dim_node; ++d) {
      out[dst_node_id * out_stride + d] = (acc[d] == identity) ? DataT{0} : acc[d];
      if (agg_op == AggOpT::kMin || agg_op == AggOpT::kMax) {
        if (out_pos != nullptr) { out_pos[dst_node_id * dim_node + d] = pos[d]; }
      }
    }
    // concatenate own features
    if (concat_own) {
      for (size_t d = 0; d < dim_concat; ++d) {
        out[dst_node_id * out_stride + dim_node + d] = concat_feat[dst_node_id * dim_concat + d];
      }
    }
  }
}

template <typename DataT, typename IdxT>
void weighted_cpu_bwd_kernel(DataT* d_node_feat,
                             DataT* d_concat_feat,
                             DataT* d_edge_weight,
                             const DataT* dout,
                             const DataT* out,
                             const IdxT* out_pos,
                             const DataT* edge_weight,
                             const DataT* degree,
                             const DataT* node_feat,
                             size_t dim_node,
                             size_t dim_concat,
                             const IdxT* offsets,
                             const IdxT* indices,
                             size_t n_src_nodes,
                             size_t n_dst_nodes,
                             size_t n_edges,
                             AggOpT agg_op)
{
  bool need_ew_grad = (d_edge_weight != nullptr);
  // 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 (need_ew_grad) { std::memset(d_edge_weight, 0, n_edges * sizeof(DataT)); }
  bool concat_own = (dim_concat > 0);
  auto out_stride = dim_node + 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];
    DataT norm     = DataT{0};
    DataT weight   = DataT{0};
    DataT tmp_out  = DataT{0};
    DataT tmp_feat = DataT{0};
    // 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) {
      for (size_t d = 0; d < dim_node; ++d) {
        auto grad      = dout[dst_node_id * out_stride + d];
        auto neigh_off = out_pos[dst_node_id * dim_node + d];
        // grad only passed on if node actually has incoming nodes
        if (neigh_off != graph::INVALID_ID<IdxT>) {
          auto neigh = indices[neigh_off];
          weight     = edge_weight[neigh_off];
          if (d_node_feat != nullptr) d_node_feat[neigh * dim_node + d] += weight * grad;
          if (need_ew_grad) {
            tmp_feat = node_feat[neigh * dim_node + d];
            d_edge_weight[neigh_off] += tmp_feat * grad;
          }
        }
      }
    } else {
      if (agg_op == AggOpT::kMean) norm = DataT{1} / (degree[dst_node_id] + DEG_EPS);
      for (size_t d = 0; d < dim_node; ++d) {
        auto grad = dout[dst_node_id * out_stride + d];
        if (agg_op == AggOpT::kMean) {
          // tmp_out needed for grad later
          tmp_out = out[dst_node_id * out_stride + d];
          // scale out_grad already now by pre-computed norm
          grad *= norm;
        }
        for (auto neigh_off = off_start; neigh_off < off_end; ++neigh_off) {
          auto neigh = indices[neigh_off];
          weight     = edge_weight[neigh_off];
          // gradient to node features simply scaled with edge_weight
          if (d_node_feat != nullptr) { d_node_feat[neigh * dim_node + d] += weight * grad; }
          // gradient w.r.t. weights needs to differentiate between kSum and kMean
          if (need_ew_grad) {
            tmp_feat = node_feat[neigh * dim_node + d];
            if (agg_op == AggOpT::kSum) {
              // for kSum, grad_w_j_to_i = sum_d node_feat_j,d * out_grad_i,d
              d_edge_weight[neigh_off] += tmp_feat * grad;
            } else {
              // kMin/kMax already checked for, i.e. here else == kMean
              // grad_w_j_to_i = sum_d norm_i * out_grad_i,d * (node_feat_j,d - out_i,d)
              d_edge_weight[neigh_off] += (tmp_feat - tmp_out) * grad;
            }
          }
        }
      }
    }
    if (concat_own && (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 + dim_node + d];
      }
    }
  }
}

}  // namespace cugraph::ops::agg_naive_csc
