/*
 * 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.
 */

#pragma once

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

#include <cstddef>

namespace cugraph::ops::aggregation_hg {

template <typename DataT, typename IdxT>
void cpu_fwd_basis_reduce(DataT* out,
                          const DataT* in,
                          const DataT* weights,
                          const DataT* etype_counts,
                          size_t dim,
                          size_t dim_b,
                          size_t ld,
                          int n_bases,
                          size_t out_node,
                          IdxT n_id,
                          int etype,
                          bool has_weights,
                          bool do_norm)
{
  auto n_etype = do_norm ? etype_counts[etype] : DataT{1.};
  for (int b = 0; b < n_bases; ++b) {
    if (!has_weights && b != etype) continue;
    auto in_off = n_id * dim + b * dim_b;
    auto factor = has_weights ? weights[etype * n_bases + b] : DataT{1};
    for (size_t d = 0; d < dim_b; ++d) {
      auto out_off = out_node * ld + d;
      auto in_val  = in[in_off + d];
      out[out_off] += factor * in_val / n_etype;
    }
  }
}

template <typename DataT, typename IdxT>
void cpu_bwd_basis_reduce(DataT* din,
                          DataT* dw,
                          const DataT* dout,
                          const DataT* in,
                          const DataT* weights,
                          const DataT* etype_counts,
                          size_t dim,
                          size_t dim_b,
                          size_t ld,
                          int n_bases,
                          size_t out_node,
                          IdxT n_id,
                          int etype,
                          bool has_weights,
                          bool do_norm)
{
  auto n_etype = do_norm ? etype_counts[etype] : DataT{1.};
  for (int b = 0; b < n_bases; ++b) {
    if (!has_weights && b != etype) continue;
    auto in_off = n_id * dim + b * dim_b;
    auto w_off  = etype * n_bases + b;
    auto factor = has_weights ? weights[w_off] : DataT{1};
    for (size_t d = 0; d < dim_b; ++d) {
      auto out_off = out_node * ld + d;
      auto grad    = dout[out_off] / n_etype;
      din[in_off + d] += grad * factor;
      if (has_weights) dw[w_off] += grad * in[in_off + d];
    }
  }
}

}  // namespace cugraph::ops::aggregation_hg
