// Copyright (C) 2018-2021 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

///////////////////////////////////////////////////////////////////////////////////////////////////
#pragma once
#include "primitive.hpp"
#include <vector>

namespace cldnn {
/// @addtogroup cpp_api C++ API
/// @{
/// @addtogroup cpp_topology Network Topology
/// @{
/// @addtogroup cpp_primitives Primitives
/// @{

/// @brief Select mode for the @ref eltwise layer.
enum class eltwise_mode : int32_t {
    /// @brief Eltwise sum.
    sum,
    /// @brief Eltwise subtract.
    sub,
    /// @brief Eltwise max.
    max,
    /// @brief Eltwise product (Hadamard).
    prod,
    /// @brief Eltwise div.
    div,
    /// @brief Eltwise min.
    min,
    /// @brief Eltwise pow.
    pow,
    /// @brief Eltwise squared diff.
    squared_diff,
    /// @brief Eltwise mod.
    mod,
    /// @brief Eltwise equal.
    eq,
    /// @brief Eltwise not equal.
    ne,
    /// @brief Eltwise less.
    lt,
    /// @brief Eltwise less of equal.
    le,
    /// @brief Eltwise greater.
    gt,
    /// @brief Eltwise greater or equal.
    ge,
    /// @brief Eltwise and.
    logic_and,
    /// @brief Eltwise or.
    logic_or,
    /// @brief Eltwise XOR.
    logic_xor,
    /// @brief Eltwise floormod.
    floor_mod
};

/// @brief Performs elementwise operations (sum, subtract, max or product) on two input primitives
/// Also supports built-in Relu @ref activation available by setting it in arguments.
/// @notes
/// - both inputs have to have equal sizes in all dimensions or the input tensors are broadcastable
///   to the same shape in which the size of each dimention is a max. of input sizes on this dimension)
/// - format of both inputs has to be the same
/// - when using integer types, only following eltwise modes are supported: sum, sub, prod, div
struct eltwise : public primitive_base<eltwise> {
    CLDNN_DECLARE_PRIMITIVE(eltwise)

    /// @brief Constructs eltwise primitive.
    /// @param id This primitive id.
    /// @param input Input primitive id.
    /// @param input2 Second input primitive id with values needed for eltwise computation.
    /// @param mode Eltwise mode.
    /// @param with_activation Enables Relu activation.
    /// @param activation_slp Relu activation slope.
    eltwise(const primitive_id& id,
            const primitive_id& input,
            const primitive_id& input2,
            eltwise_mode mode,
            const primitive_id& ext_prim_id = "",
            const padding& output_padding = padding())
        : primitive_base(id, {input, input2}, ext_prim_id, output_padding),
          mode(mode),
          coefficients(std::vector<float>(0)),
          stride(std::vector<tensor>(0)) {}

    /// @brief Constructs eltwise primitive.
    /// @param id This primitive id.
    /// @param input Input primitive id.
    /// @param input2 Second input primitive id with values needed for eltwise computation.
    /// @param stride Defines shift in input buffers between adjacent calculations of output values.
    /// @param mode Eltwise mode.
    /// @param with_activation Enables Relu activation.
    /// @param activation_slp Relu activation slope.
    eltwise(const primitive_id& id,
            const primitive_id& input,
            const primitive_id& input2,
            std::vector<tensor> stride,
            eltwise_mode mode,
            const primitive_id& ext_prim_id = "",
            const padding& output_padding = padding())
        : primitive_base(id, {input, input2}, ext_prim_id, output_padding),
          mode(mode),
          coefficients(std::vector<float>(0)),
          stride(stride) {}

    /// @brief Constructs eltwise primitive.
    /// @param id This primitive id.
    /// @param inputs Input primitives ids.
    /// @param mode Eltwise mode.
    /// @param data_type Expected output data type.
    eltwise(const primitive_id& id,
            const std::vector<primitive_id>& inputs,
            eltwise_mode mode,
            data_types data_type,
            const primitive_id& ext_prim_id = "",
            const padding& output_padding = padding())
        : primitive_base(id, inputs, ext_prim_id, output_padding, optional_data_type{data_type}),
          mode(mode),
          coefficients(std::vector<float>(0)),
          stride(std::vector<tensor>(0)) {}

    /// @brief Constructs eltwise primitive.
    /// @param id This primitive id.
    /// @param inputs Input primitives ids.
    /// @param mode Eltwise mode.
    eltwise(const primitive_id& id,
            const std::vector<primitive_id>& inputs,
            eltwise_mode mode,
            const primitive_id& ext_prim_id = "",
            const padding& output_padding = padding())
        : primitive_base(id, inputs, ext_prim_id, output_padding),
          mode(mode),
          coefficients(std::vector<float>(0)),
          stride(std::vector<tensor>(0)) {}

    /// @brief Constructs eltwise primitive.
    /// @param id This primitive id.
    /// @param inputs Input primitives ids.
    /// @param coefficients Blob-wise coefficient for SUM operation
    /// @param mode Eltwise mode.
    eltwise(const primitive_id& id,
            const std::vector<primitive_id>& inputs,
            eltwise_mode mode,
            const std::vector<float>& coefficients,
            data_types data_type,
            const primitive_id& ext_prim_id = "",
            const padding& output_padding = padding())
        : primitive_base(id, inputs, ext_prim_id, output_padding, optional_data_type{data_type}),
          mode(mode),
          coefficients(coefficients),
          stride(std::vector<tensor>(0)) {
        if (mode == eltwise_mode::sum && !coefficients.empty() && coefficients.size() != inputs.size()) {
            throw std::invalid_argument("Invalid eltwise sum coefficients count (should be equal to 0 or input.size)");
        }
        if (mode != eltwise_mode::sum && !coefficients.empty()) {
            throw std::invalid_argument("Only eltwise sum operation supports blob-wise coefficients");
        }
    }

    bool needs_onednn_sum_post_op(layout input_layout) const {
        if (mode == eltwise_mode::sum &&
            (input_layout.size.spatial[0] > 1 || input_layout.size.spatial[1] > 1 || input_layout.size.batch[0] > 1)) {
            return true;
        }

        return false;
    }

    /// @param mode Eltwise mode.
    eltwise_mode mode;
    /// @param coefficients Blob-wise coefficient for SUM operation.
    std::vector<float> coefficients;
    /// @brief Defines shift in input buffers between adjacent calculations of output values.
    std::vector<tensor> stride;
};
/// @}
/// @}
/// @}
}  // namespace cldnn
