/*
 * Copyright (c) 2021-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/activation.hpp>
#include <cugraph-ops/utils/error.hpp>

#include <cmath>

namespace cugraph::ops {

template <typename DataT>
inline DataT activation_fwd_cpu(const activation_params& aparams, DataT in)
{
  static constexpr DataT ZERO = DataT{0};
  static constexpr DataT ONE  = DataT{1};
  switch (aparams.type) {
    case ActivationOpT::kLinear: return in;
    case ActivationOpT::kRelu: return in < ZERO ? ZERO : in;
    case ActivationOpT::kSigmoid: return ONE / (ONE + std::exp(-in));
    case ActivationOpT::kTanh: return std::tanh(in);
    case ActivationOpT::kELU: return in < ZERO ? std::exp(in) - ONE : in;
    case ActivationOpT::kScalar: return in * aparams.alpha;
    case ActivationOpT::kLeakyRelu: return in < ZERO ? in * aparams.alpha : in;
    default: ASSERT(false, "Invalid activation type pased '%d'!", static_cast<int>(aparams.type));
  }
}

template <typename DataT>
inline DataT activation_bwd_cpu(const activation_params& aparams,
                                DataT out,
                                DataT /*in*/ = DataT{0})
{
  static constexpr DataT ZERO = DataT{0};
  static constexpr DataT ONE  = DataT{1};
  switch (aparams.type) {
    case ActivationOpT::kLinear: return ONE;
    case ActivationOpT::kRelu: return out <= ZERO ? ZERO : ONE;
    case ActivationOpT::kSigmoid: return out * (ONE - out);
    case ActivationOpT::kTanh: return ONE - out * out;
    case ActivationOpT::kELU: return out < ZERO ? out + ONE : ONE;
    case ActivationOpT::kScalar: return aparams.alpha;
    case ActivationOpT::kLeakyRelu: return out < ZERO ? aparams.alpha : ONE;
    default: ASSERT(false, "Invalid activation type pased '%d'!", static_cast<int>(aparams.type));
  }
}

}  // namespace cugraph::ops
