/*
 * Copyright (c) 2022-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 "cpu_agg_hg_basis_common.hpp"

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

#include <cstdlib>
#include <cstring>

namespace cugraph::ops::aggregation_hg {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_pre_kernel(DataT* out,
                              const DataT* in,
                              const DataT* weights_comb,
                              size_t in_dim,
                              size_t out_ld,
                              int n_bases,
                              const IdxT* offsets,
                              const IdxT* indices,
                              const int32_t* edge_types,
                              IdxT n_dst_nodes,
                              int32_t n_edge_types,
                              bool concat_own,
                              bool norm_by_out_degree)
{
  std::memset(out, 0, n_dst_nodes * out_ld * sizeof(DataT));
  auto has_weights = (weights_comb != nullptr);
  if (!has_weights) n_bases = n_edge_types;
  size_t dim_b = in_dim / (concat_own ? n_bases + 1 : n_bases);
  // norm values
  auto* etype_counts = reinterpret_cast<DataT*>(std::malloc(n_edge_types * sizeof(DataT)));
  for (size_t out_node = 0; out_node < n_dst_nodes; ++out_node) {
    std::memset(etype_counts, 0, n_edge_types * sizeof(DataT));
    IdxT n_neighbors = offsets[out_node + 1] - offsets[out_node];
    for (IdxT n = 0; n < n_neighbors; ++n) {
      auto n_off = offsets[out_node] + n;
      auto etype = edge_types[n_off];
      etype_counts[etype] += DataT{1.};
    }

    for (IdxT n = 0; n < n_neighbors; ++n) {
      size_t n_off = offsets[out_node] + n;
      IdxT n_id    = indices[n_off];
      auto etype   = edge_types[n_off];
      aggregation_hg::cpu_fwd_basis_reduce(out,
                                           in,
                                           weights_comb,
                                           etype_counts,
                                           in_dim,
                                           dim_b,
                                           out_ld,
                                           n_bases,
                                           out_node,
                                           n_id,
                                           etype,
                                           has_weights,
                                           norm_by_out_degree);
    }
    if (concat_own) {
      auto out_off = out_node * out_ld;
      auto in_off  = out_node * in_dim + n_bases * dim_b;
      for (size_t d = 0; d < dim_b; ++d)
        out[out_off + d] += in[in_off + d];
    }
  }
  std::free(etype_counts);
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_pre_kernel(DataT* din,
                              DataT* dw,
                              const DataT* dout,
                              const DataT* in,
                              const DataT* weights_comb,
                              size_t in_dim,
                              size_t out_ld,
                              int n_bases,
                              const IdxT* offsets,
                              const IdxT* indices,
                              const int32_t* edge_types,
                              IdxT n_src_nodes,
                              IdxT n_dst_nodes,
                              int32_t n_edge_types,
                              bool concat_own,
                              bool norm_by_out_degree)
{
  // set gradients to 0 so we can accumulate onto them
  std::memset(din, 0, n_src_nodes * in_dim * sizeof(DataT));
  auto has_weights = weights_comb != nullptr && dw != nullptr;
  if (has_weights) {
    std::memset(dw, 0, static_cast<size_t>(n_edge_types * n_bases) * sizeof(DataT));
  } else {
    n_bases = n_edge_types;
  }
  size_t dim_b = in_dim / (concat_own ? n_bases + 1 : n_bases);
  // norm values
  auto* etype_counts =
    reinterpret_cast<DataT*>(std::malloc(static_cast<size_t>(n_edge_types) * sizeof(DataT)));

  for (size_t out_node = 0; out_node < n_dst_nodes; ++out_node) {
    auto n_neighbors = offsets[out_node + 1] - offsets[out_node];
    std::memset(etype_counts, 0, n_edge_types * sizeof(DataT));
    for (IdxT n = 0; n < n_neighbors; ++n) {
      auto n_off = offsets[out_node] + n;
      auto etype = edge_types[n_off];
      etype_counts[etype] += DataT{1.};
    }

    for (IdxT n = 0; n < n_neighbors; ++n) {
      size_t n_off = offsets[out_node] + n;
      IdxT n_id    = indices[n_off];
      auto etype   = edge_types[n_off];
      aggregation_hg::cpu_bwd_basis_reduce(din,
                                           dw,
                                           dout,
                                           in,
                                           weights_comb,
                                           etype_counts,
                                           in_dim,
                                           dim_b,
                                           out_ld,
                                           n_bases,
                                           out_node,
                                           n_id,
                                           etype,
                                           has_weights,
                                           norm_by_out_degree);
    }
    if (concat_own) {
      auto out_off = out_node * out_ld;
      auto in_off  = out_node * in_dim + n_bases * dim_b;
      for (size_t d = 0; d < dim_b; ++d)
        din[in_off + d] += dout[out_off + d];
    }
  }
  std::free(etype_counts);
}

}  // namespace cugraph::ops::aggregation_hg
