// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include "onnx_core/runtime/kernels/kernel_context.h"
#include "onnx_core/runtime/memory/simple_tensor.h"
#include "onnx_core/runtime/runtime_context.h"
#include "onnx_extensions/kernels/tuning/portable_parallel_tuning.h"

namespace ONNX_LIGHT_NAMESPACE::onnx_kernels {
// Re-exports the runtime types moved to ``onnx_core::runtime`` so
// kernel implementations below can keep referring to them
// unqualified, matching pre-move code.
using namespace ::onnx_light::core::runtime;

namespace kernel {
using ::onnx_light::core::runtime::DefaultOpset;
using ::onnx_light::core::runtime::KernelBase;
using ::onnx_light::core::runtime::KernelContext;
using ::onnx_light::core::runtime::OpsetId;

// ---------------------------------------------------------------------------
// Reference implementations of the ``logical`` backend test kernels.
//
// Each kernel is exposed as a small class whose constructor takes a
// :ref:`KernelContext` (carrying the opset against which the kernel must
// behave) and whose ``operator()`` performs the computation.
//
// Two flavors of ``operator()`` are provided:
//
//   * The returning overload (``Tensor operator()(...) const``) allocates a
//     fresh ``Tensor`` whose data buffer is owned by the returned value.
//   * The in-place overload (``void operator()(..., Tensor &output) const``)
//     writes results into a caller-supplied output tensor whose buffer has
//     already been allocated. The caller is responsible for setting
//     ``output.data_type``, ``output.shape`` and sizing ``output.data`` to
//     match the operator's expected output; the kernel validates these
//     attributes and throws ``std::invalid_argument`` on mismatch.
//
// And/Or/Xor operate on ``BOOL`` tensors (one byte per element, 0 or 1)
// and support multidirectional broadcasting per the standard NumPy/ONNX
// rules — mirroring the broadcasting behavior exercised elsewhere in the
// backend test library (see ``kernel::Add``).
//
// Each kernel class also exposes a ``static constexpr bool CanRunInPlace()``
// query indicating whether the output tensor's data buffer may alias one of
// the input tensors' buffers. All three element-wise logical kernels here
// support in-place execution (the output element depends only on the
// corresponding input elements at the same index).
// ---------------------------------------------------------------------------

/// Element-wise logical AND on BOOL tensors with multidirectional broadcasting.
class And : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:And";
  explicit And(const KernelContext &ctx);
  /// Registers the portable BOOL tuning schema.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise logical OR on BOOL tensors with multidirectional broadcasting.
class Or : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Or";
  explicit Or(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise logical XOR on BOOL tensors with multidirectional broadcasting.
class Xor : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Xor";
  explicit Xor(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise logical NOT on a BOOL tensor (opset 1). The output has the
/// same BOOL dtype and shape as the input. Mirrors the upstream ONNX
/// ``Not`` reference implementation (``np.logical_not``).
class Not : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Not";
  explicit Not(const KernelContext &ctx);
  /// Registers the portable BOOL tuning schema.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise ``IsNaN``: returns a BOOL tensor with the same shape as the
/// input, where each element is ``true`` iff the corresponding input value
/// is NaN. Mirrors the upstream ONNX ``IsNaN`` reference implementation
/// (``np.isnan``). Supports FLOAT, DOUBLE, FLOAT16 and BFLOAT16 inputs.
class IsNaN : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:IsNaN";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Input is FLOAT (4 bytes/elt) and output is BOOL (1 byte/elt), so the
  /// output buffer cannot alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``IsInf``: returns a BOOL tensor with the same shape as the
/// input, where each element is ``true`` iff the corresponding input value
/// is +/- infinity. The two boolean attributes ``detect_positive`` and
/// ``detect_negative`` (both default to 1) toggle whether +inf and -inf
/// are reported respectively. Mirrors the upstream ONNX ``IsInf`` reference
/// implementation (``np.isinf`` / ``np.isposinf`` / ``np.isneginf``). Only
/// the FLOAT input dtype is supported.
class IsInf : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:IsInf";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int64_t detect_positive = 1, int64_t detect_negative = 1,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int64_t detect_positive, int64_t detect_negative,
                  Tensor &output) const;

  /// Input is FLOAT (4 bytes/elt) and output is BOOL (1 byte/elt), so the
  /// output buffer cannot alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``Greater`` comparison with multidirectional broadcasting.
/// Inputs may be FLOAT, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32 or
/// UINT64 (both inputs must share the same dtype); the output is BOOL (one
/// byte per element, ``0`` or ``1``). Mirrors the upstream ONNX ``Greater``
/// reference implementation (``np.greater``).
class Greater : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Greater";
  explicit Greater(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``Less`` comparison with multidirectional broadcasting.
/// Inputs may be FLOAT, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32 or
/// UINT64 (both inputs must share the same dtype); the output is BOOL (one
/// byte per element, ``0`` or ``1``). Mirrors the upstream ONNX ``Less``
/// reference implementation (``np.less``).
class Less : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Less";
  explicit Less(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``GreaterOrEqual`` comparison with multidirectional
/// broadcasting. Inputs may be FLOAT, INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32 or UINT64 (both inputs must share the same dtype); the
/// output is BOOL (one byte per element, ``0`` or ``1``). Mirrors the upstream
/// ONNX ``GreaterOrEqual`` reference implementation (``np.greater_equal``).
class GreaterOrEqual : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:GreaterOrEqual";
  explicit GreaterOrEqual(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``LessOrEqual`` comparison with multidirectional
/// broadcasting. Inputs may be FLOAT, INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32 or UINT64 (both inputs must share the same dtype); the
/// output is BOOL (one byte per element, ``0`` or ``1``). Mirrors the upstream
/// ONNX ``LessOrEqual`` reference implementation (``np.less_equal``).
class LessOrEqual : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:LessOrEqual";
  explicit LessOrEqual(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``Equal`` comparison with multidirectional broadcasting.
/// Inputs may be BOOL, FLOAT, DOUBLE, INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32, UINT64 or STRING (both inputs must share the same dtype);
/// the output is BOOL (one byte per element, ``0`` or ``1``). Mirrors the
/// upstream ONNX ``Equal`` reference implementation (``np.equal``). STRING
/// support matches the ``Equal`` opset 19 type expansion and is restricted
/// to equal-shape inputs or scalar broadcasting.
class Equal : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Equal";
  explicit Equal(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ``Where`` selection with multidirectional broadcasting.
/// ``condition`` must be BOOL; ``x`` and ``y`` must share the same dtype and
/// may be BOOL, FLOAT, DOUBLE, INT8, INT16, INT32, INT64, UINT8, UINT16,
/// UINT32, UINT64 or STRING. Output dtype equals ``x``/``y`` dtype.
class Where : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Where";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &condition, const Tensor &x, const Tensor &y,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &condition, const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise bitwise AND on integer tensors with multidirectional
/// broadcasting (opset 18). Inputs may be INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32 or UINT64 (both inputs must share the same dtype); the
/// output has the same dtype. Mirrors the upstream ONNX ``BitwiseAnd``
/// reference implementation (``np.bitwise_and``).
class BitwiseAnd : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitwiseAnd";
  explicit BitwiseAnd(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise bitwise OR on integer tensors with multidirectional
/// broadcasting (opset 18). Inputs may be INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32 or UINT64 (both inputs must share the same dtype); the
/// output has the same dtype. Mirrors the upstream ONNX ``BitwiseOr``
/// reference implementation (``np.bitwise_or``).
class BitwiseOr : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitwiseOr";
  explicit BitwiseOr(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise bitwise XOR on integer tensors with multidirectional
/// broadcasting (opset 18). Inputs may be INT8, INT16, INT32, INT64, UINT8,
/// UINT16, UINT32 or UINT64 (both inputs must share the same dtype); the
/// output has the same dtype. Mirrors the upstream ONNX ``BitwiseXor``
/// reference implementation (``np.bitwise_xor``).
class BitwiseXor : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitwiseXor";
  explicit BitwiseXor(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise bitwise NOT on integer tensors (opset 18). Input may be
/// INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32 or UINT64; the output
/// has the same dtype and shape. Mirrors the upstream ONNX ``BitwiseNot``
/// reference implementation (``np.bitwise_not``).
class BitwiseNot : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitwiseNot";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise ``BitShift`` on unsigned integer tensors (opset 11). Both
/// inputs must share the same dtype (UINT8, UINT16, UINT32 or UINT64); the
/// output has the same dtype. Multidirectional (Numpy-style) broadcasting
/// is supported. The required ``direction`` attribute selects ``"LEFT"`` or
/// ``"RIGHT"`` and is passed to ``operator()``. Mirrors the upstream
/// ``np.left_shift`` / ``np.right_shift`` reference implementations.
class BitShift : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitShift";
  explicit BitShift(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  /// Direction of the bitwise shift.
  enum class Direction { kLeft, kRight };

  Tensor operator()(const Tensor &x, const Tensor &y, Direction direction,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Direction direction, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return true; }
};

} // namespace kernel
} // namespace ONNX_LIGHT_NAMESPACE::onnx_kernels
