/*
 * Copyright (c) 2020-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

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

#include <cstdint>

namespace cugraph::ops::gnn {

/**
 * @brief Type of aggregation operation
 */
enum class AggOpT {
  /** mean across the neighbor embeddings */
  kMean = 0,
  /** sum across the neighbor embeddings */
  kSum,
  /** max across the neighbor embeddings */
  kMax,
  /** min across the neighbor embeddings */
  kMin
};  // enum class AggOpT

struct mha_params {
  mha_params() = default;
  int num_heads{1};
  activation_params activation;
  bool add_own_node{false};
};

struct self_loop_type_conv_t {
  int32_t factor{0}, offset{0};

  [[nodiscard]] __host__ __device__ inline int32_t to_etype(int32_t ntype) const
  {
    return factor * ntype + offset;
  }
};

}  // namespace cugraph::ops::gnn
