/*
 * 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 "cpu_agg_hg_basis_common.hpp"

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

#include <cstdlib>
#include <cstring>

namespace cugraph::ops::aggregation_hg {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_post_kernel(DataT* out,
                               const DataT* in,
                               const DataT* weights,
                               size_t dim,
                               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 do_norm)
{
  bool use_weights = weights != nullptr;
  auto out_dim     = dim * (use_weights ? n_bases : n_edge_types) + (concat_own ? dim : 0);

  // set outputs to 0 so we can accumulate onto them
  std::memset(out, 0, n_dst_nodes * out_dim * sizeof(DataT));
  auto has_weights = weights != nullptr;
  if (!has_weights) n_bases = n_edge_types;
  // norm values
  auto* etype_counts = reinterpret_cast<DataT*>(std::malloc(n_edge_types * sizeof(DataT)));
  for (IdxT out_node = 0; out_node < n_dst_nodes; ++out_node) {
    auto off_start = offsets[out_node];
    auto off_end   = offsets[out_node + 1];
    auto n_edges   = off_end - off_start;
    std::memset(etype_counts, 0, n_edge_types * sizeof(DataT));
    for (IdxT n = 0; n < n_edges; ++n)
      etype_counts[edge_types[n + off_start]] += DataT{1.};

    for (IdxT n = 0; n < n_edges; ++n) {
      auto n_id    = indices[n + off_start];
      auto etype   = edge_types[n + off_start];
      auto n_etype = do_norm ? etype_counts[etype] : DataT{1.};
      for (size_t d = 0; d < dim; ++d) {
        auto out_off = out_node * out_dim + d;
        auto in_val  = in[n_id * dim + d];
        if (has_weights) {
          for (int b = 0; b < n_bases; ++b) {
            out[out_off + b * dim] += weights[etype * n_bases + b] * in_val / n_etype;
          }
        } else {
          out[out_off + etype * dim] += in_val / n_etype;
        }
      }
    }
    if (concat_own) {
      auto out_off = out_node * out_dim + n_bases * dim;
      auto in_off  = out_node * dim;
      for (size_t d = 0; d < dim; ++d)
        out[out_off + d] += in[in_off + d];
    }
  }
  std::free(etype_counts);
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_post_kernel(DataT* din,
                               DataT* dw,
                               const DataT* dout,
                               const DataT* in,
                               const DataT* weights,
                               size_t dim,
                               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 do_norm)
{
  bool use_weights = weights != nullptr;
  auto out_dim     = dim * (use_weights ? n_bases : n_edge_types) + (concat_own ? dim : 0);

  // set gradients to 0 so we can accumulate onto them
  auto has_din = din != nullptr;
  if (has_din) std::memset(din, 0, n_src_nodes * dim * sizeof(DataT));
  auto has_weights = weights != nullptr;
  auto has_dw      = has_weights && (dw != nullptr);
  if (has_dw) { std::memset(dw, 0, static_cast<size_t>(n_edge_types * n_bases) * sizeof(DataT)); }
  if (!has_weights) n_bases = n_edge_types;

  // norm values
  auto* etype_counts =
    reinterpret_cast<DataT*>(std::malloc(static_cast<size_t>(n_edge_types) * sizeof(DataT)));

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

    for (IdxT n = 0; n < n_edges; ++n) {
      auto n_id    = indices[n + off_start];
      auto etype   = edge_types[n + off_start];
      auto n_etype = do_norm ? etype_counts[etype] : DataT{1.};
      for (size_t d = 0; d < dim; ++d) {
        auto out_off = out_node * out_dim + d;
        auto in_off  = n_id * dim + d;
        auto in_val  = in[in_off];
        auto grad_in = DataT{0};
        if (has_weights) {
          for (int b = 0; b < n_bases; ++b) {
            auto grad_out = dout[out_off + b * dim] / n_etype;
            if (has_dw) { dw[etype * n_bases + b] += grad_out * in_val; }
            grad_in += weights[etype * n_bases + b] * grad_out;
          }
        } else {
          // no grad on dw
          grad_in = dout[out_off + etype * dim] / n_etype;
        }
        if (has_din) { din[in_off] += grad_in; }
      }
    }
    if (concat_own && has_din) {
      auto out_off = out_node * out_dim + n_bases * dim;
      auto in_off  = out_node * dim;
      for (size_t d = 0; d < dim; ++d) {
        din[in_off + d] += dout[out_off + d];
      }
    }
  }
  std::free(etype_counts);
}

}  // namespace cugraph::ops::aggregation_hg
