// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0
#include <string>

#include "onnx_lib/defs/schema.h"

namespace ONNX_LIGHT_NAMESPACE {

static std::function<void(OpSchema &)> BinaryLogicDocGenerator_opset12(const char *name) {
  return [=](OpSchema &schema) {
    std::string doc;
    POPULATE_OP_DOC_STR(doc = R"DOC(
Returns the tensor resulted from performing the `{name}` logical operation
elementwise on the input tensors `A` and `B` (with Numpy-style broadcasting support).

{broadcast_doc}
)DOC";
                        ReplaceAll(doc, "{name}", name);
                        ReplaceAll(doc, "{broadcast_doc}", GenerateBroadcastingDocMul().c_str()););
    schema.SetDoc(doc);
    schema.Input(0, "A", "First input operand for the logical operator.", "T");
    schema.Input(1, "B", "Second input operand for the logical operator.", "T");
    schema.Output(0, "C", "Result tensor.", "T1");
    schema.TypeAndShapeInferenceFunction([](InferenceContext &ctx) {
      // Type inference
      updateOutputElemType(ctx, 0, TensorProto::BOOL);
      // Shape inference
      if (hasNInputShapes(ctx, 2))
        bidirectionalBroadcastShapeInference(
            ctx.getInputType(0)->tensor_type().shape(), ctx.getInputType(1)->tensor_type().shape(),
            *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape());
    });
  };
}

ONNX_OPERATOR_SET_SCHEMA(Greater, 9,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset12("greater"))
                             .TypeConstraint("T", OpSchema::all_numeric_types(),
                                             "Constrain input types to all numeric tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Less, 9,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset12("less"))
                             .TypeConstraint("T", OpSchema::all_numeric_types(),
                                             "Constrain input types to all numeric tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(
    Equal, 11,
    OpSchema()
        .FillUsing(BinaryLogicDocGenerator_opset12("equal"))
        .TypeConstraint("T",
                        {"tensor(bool)", "tensor(uint8)", "tensor(uint16)", "tensor(uint32)",
                         "tensor(uint64)", "tensor(int8)", "tensor(int16)", "tensor(int32)",
                         "tensor(int64)", "tensor(float16)", "tensor(float)", "tensor(double)"},
                        "Constrain input types to all numeric tensors.")
        .TypeConstraint("T1", {"tensor(bool)"}, "Constrain output to boolean tensor."));

inline static void logicalOpInference_opset1(InferenceContext &ctx) {
  updateOutputElemType(ctx, 0, TensorProto::BOOL);
  if (hasInputShape(ctx, 0)) {
    propagateShapeFromInputToOutput(ctx, 0, 0);
  }
}

static std::function<void(OpSchema &)> BinaryLogicDocGenerator_opset1(const char *name) {
  return [=](OpSchema &schema) {
    std::string doc;
    POPULATE_OP_DOC_STR(doc = R"DOC(
Returns the tensor resulted from performing the `{name}` logical operation
elementwise on the input tensors `A` and `B`.

If broadcasting is enabled, the right-hand-side argument will be broadcasted
to match the shape of left-hand-side argument. See the doc of `Add` for a
detailed description of the broadcasting rules.
)DOC";
                        ReplaceAll(doc, "{name}", name););
    schema.SetDoc(doc);
    schema.Attr("broadcast", "Enable broadcasting", AttributeProto::INT, static_cast<int64_t>(0));
    schema.Attr("axis", "If set, defines the broadcast dimensions.", AttributeProto::INT,
                OPTIONAL_VALUE);
    schema.Input(0, "A", "Left input tensor for the logical operator.", "T");
    schema.Input(1, "B", "Right input tensor for the logical operator.", "T");
    schema.Output(0, "C", "Result tensor.", "T1");
    schema.TypeAndShapeInferenceFunction(logicalOpInference_opset1);
  };
}

static std::function<void(OpSchema &)> BinaryLogicDocGenerator_opset7(const char *name) {
  return [=](OpSchema &schema) {
    std::string doc;
    POPULATE_OP_DOC_STR(doc = R"DOC(
Returns the tensor resulted from performing the `{name}` logical operation
elementwise on the input tensors `A` and `B` (with Numpy-style broadcasting support).

{broadcast_doc}
)DOC";
                        ReplaceAll(doc, "{name}", name);
                        ReplaceAll(doc, "{broadcast_doc}", GenerateBroadcastingDocMul().c_str()););
    schema.SetDoc(doc);
    schema.Input(0, "A", "First input operand for the logical operator.", "T");
    schema.Input(1, "B", "Second input operand for the logical operator.", "T");
    schema.Output(0, "C", "Result tensor.", "T1");
    schema.TypeAndShapeInferenceFunction([](InferenceContext &ctx) {
      updateOutputElemType(ctx, 0, TensorProto::BOOL);
      if (hasNInputShapes(ctx, 2))
        bidirectionalBroadcastShapeInference(
            ctx.getInputType(0)->tensor_type().shape(), ctx.getInputType(1)->tensor_type().shape(),
            *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape());
    });
  };
}

ONNX_OPERATOR_SET_SCHEMA(And, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("and"))
                             .TypeConstraint("T", {"tensor(bool)"},
                                             "Constrain input to boolean tensor.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Or, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("or"))
                             .TypeConstraint("T", {"tensor(bool)"},
                                             "Constrain input to boolean tensor.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Xor, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("xor"))
                             .TypeConstraint("T", {"tensor(bool)"},
                                             "Constrain input to boolean tensor.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Greater, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("greater"))
                             .TypeConstraint("T",
                                             {"tensor(float16)", "tensor(float)", "tensor(double)"},
                                             "Constrain input to float tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Less, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("less"))
                             .TypeConstraint("T",
                                             {"tensor(float16)", "tensor(float)", "tensor(double)"},
                                             "Constrain input to float tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Equal, 1,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset1("equal"))
                             .TypeConstraint("T",
                                             {"tensor(bool)", "tensor(int32)", "tensor(int64)"},
                                             "Constrain input to integral tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Equal, 7,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset7("equal"))
                             .TypeConstraint("T",
                                             {"tensor(bool)", "tensor(int32)", "tensor(int64)"},
                                             "Constrain input to integral tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Greater, 7,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset7("greater"))
                             .TypeConstraint("T",
                                             {"tensor(float16)", "tensor(float)", "tensor(double)"},
                                             "Constrain input to float tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

ONNX_OPERATOR_SET_SCHEMA(Less, 7,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator_opset7("less"))
                             .TypeConstraint("T",
                                             {"tensor(float16)", "tensor(float)", "tensor(double)"},
                                             "Constrain input to float tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor."));

// Shares same doc generator as newer opset 16 version.
extern std::function<void(OpSchema &)> BinaryLogicDocGenerator(const char *name);

ONNX_OPERATOR_SET_SCHEMA(LessOrEqual, 12,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator("less_equal"))
                             .TypeConstraint("T", OpSchema::all_numeric_types(),
                                             "Constrain input types to all numeric tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor.")
                             .TypeAndShapeInferenceFunction(InferenceFunction())
                             .FunctionBody(R"ONNX(
        {
            O1 = Less (A, B)
            O2 = Equal (A, B)
            C = Or (O1, O2)
        }
        )ONNX"));

ONNX_OPERATOR_SET_SCHEMA(GreaterOrEqual, 12,
                         OpSchema()
                             .FillUsing(BinaryLogicDocGenerator("greater_equal"))
                             .TypeConstraint("T", OpSchema::all_numeric_types(),
                                             "Constrain input types to all numeric tensors.")
                             .TypeConstraint("T1", {"tensor(bool)"},
                                             "Constrain output to boolean tensor.")
                             .TypeAndShapeInferenceFunction(InferenceFunction())
                             .FunctionBody(R"ONNX(
        {
            O1 = Greater (A, B)
            O2 = Equal (A, B)
            C = Or (O1, O2)
        }
        )ONNX"));

ONNX_OPERATOR_SET_SCHEMA(
    Equal, 13,
    OpSchema()
        .FillUsing(BinaryLogicDocGenerator("equal"))
        .TypeConstraint("T",
                        {"tensor(bool)", "tensor(uint8)", "tensor(uint16)", "tensor(uint32)",
                         "tensor(uint64)", "tensor(int8)", "tensor(int16)", "tensor(int32)",
                         "tensor(int64)", "tensor(float16)", "tensor(float)", "tensor(double)",
                         "tensor(bfloat16)"},
                        "Constrain input types to all numeric tensors.")
        .TypeConstraint("T1", {"tensor(bool)"}, "Constrain output to boolean tensor."));

} // namespace ONNX_LIGHT_NAMESPACE
