// 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 <cstdint>

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 ``reduction`` 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.
//
// ``ReduceSum`` mirrors the ONNX ``ReduceSum`` operator at opset 13, where
// ``axes`` is an optional second input tensor (int64) rather than an
// attribute. The kernel supports FLOAT tensors and the standard ``keepdims``
// and ``noop_with_empty_axes`` attributes. When ``axes`` is omitted (or
// empty), the kernel either reduces over all dimensions (the default,
// ``noop_with_empty_axes == false``) or performs an identity copy
// (``noop_with_empty_axes == true``).
//
// 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. ``ReduceSum``'s output shape generally differs
// from its input, so storage cannot be shared with an input.
// ---------------------------------------------------------------------------

/// Returns the indices of the maximum (``ArgMax``) or minimum (``ArgMin``)
/// values of a FLOAT input ``data`` along a single ``axis``.
///
/// The output is always an INT64 tensor whose shape matches ``data`` with
/// the reduced ``axis`` either dropped (``keepdims = false``) or replaced
/// by ``1`` (``keepdims = true``, the default).
///
/// ``axis`` follows ONNX semantics and accepts values in
/// ``[-rank(data), rank(data) - 1]``. ``select_last_index`` (default
/// ``false``, ONNX opset 12+) selects the index of the last occurrence of
/// the extremum when multiple positions are tied; otherwise the first
/// occurrence is selected.
class ArgReduce : public KernelBase {
public:
  enum class Mode { kMax, kMin };

  ArgReduce(const KernelContext &ctx, Mode mode) : KernelBase(ctx), mode_(mode) {}

  Tensor operator()(const Tensor &data, int64_t axis = 0, bool keepdims = true,
                    bool select_last_index = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, int64_t axis, bool keepdims, bool select_last_index,
                  Tensor &output) const;

  /// The INT64 output buffer never aliases the FLOAT input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }

private:
  Mode mode_;
};

/// ``ArgMax`` wrapper around :class:`ArgReduce` configured for the maximum.
class ArgMax : public ArgReduce {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ArgMax";
  void Run(RuntimeContext &rt) override;
  explicit ArgMax(const KernelContext &ctx) : ArgReduce(ctx, ArgReduce::Mode::kMax) {}
};

/// ``ArgMin`` wrapper around :class:`ArgReduce` configured for the minimum.
class ArgMin : public ArgReduce {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ArgMin";
  void Run(RuntimeContext &rt) override;
  explicit ArgMin(const KernelContext &ctx) : ArgReduce(ctx, ArgReduce::Mode::kMin) {}
};

/// Sum reduction of a FLOAT or DOUBLE input ``data`` along the dimensions listed in the
/// optional ``axes`` int64 tensor. If ``axes`` is omitted (or empty), the
/// kernel reduces over all dimensions unless ``noop_with_empty_axes`` is true
/// in which case it performs an identity copy.
class ReduceSum : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceSum";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;

  /// ``axes`` omitted: reduces over all dimensions of ``data`` (the default
  /// when ``noop_with_empty_axes`` is false) or returns a copy of ``data``
  /// when ``noop_with_empty_axes`` is true.
  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  /// Explicit ``axes``: an int64 tensor whose elements are the dimensions to
  /// reduce. Negative axis values are accepted (ONNX semantics: ``axis`` in
  /// ``[-rank(data), rank(data) - 1]``).
  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  /// Output shape generally differs from the input shape, so the output
  /// buffer cannot in general alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Shared FLOAT kernel for ``ReduceMax`` and ``ReduceMin``.
class ReduceMinMax : public KernelBase {
public:
  enum class Mode { kMax, kMin };

  ReduceMinMax(const KernelContext &ctx, Mode mode) : KernelBase(ctx), mode_(mode) {}

  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

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

private:
  Mode mode_;
};

/// ``ReduceMax`` wrapper around :class:`ReduceMinMax` configured for maximum.
class ReduceMax : public ReduceMinMax {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceMax";
  void Run(RuntimeContext &rt) override;
  explicit ReduceMax(const KernelContext &ctx) : ReduceMinMax(ctx, ReduceMinMax::Mode::kMax) {}
};

/// ``ReduceMin`` wrapper around :class:`ReduceMinMax` configured for minimum.
class ReduceMin : public ReduceMinMax {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceMin";
  void Run(RuntimeContext &rt) override;
  explicit ReduceMin(const KernelContext &ctx) : ReduceMinMax(ctx, ReduceMinMax::Mode::kMin) {}
};

/// Shared FLOAT kernel for ``ReduceL1``, ``ReduceL2`` and ``ReduceSumSquare``.
///
/// ``ReduceL1`` computes the sum of absolute values along the reduced axes:
/// ``y = sum(|x|, axes)``. ``ReduceL2`` computes the Euclidean (L2) norm:
/// ``y = sqrt(sum(x * x, axes))``. ``ReduceSumSquare`` computes the sum of
/// squared values: ``y = sum(x * x, axes)`` — i.e. the L2 norm without the
/// final square root. The empty-set identity for all three is ``0`` (so
/// reducing over a size-0 axis yields a zero-filled output).
///
/// The interface mirrors :class:`ReduceMinMax`: the ``axes`` may either be
/// omitted (reduce-all unless ``noop_with_empty_axes == true``) or supplied
/// as an int64 tensor (opset 18+). Negative axes follow ONNX semantics.
class ReduceL1L2 : public KernelBase {
public:
  enum class Mode { kL1, kL2, kSumSquare };

  ReduceL1L2(const KernelContext &ctx, Mode mode) : KernelBase(ctx), mode_(mode) {}

  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

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

private:
  Mode mode_;
};

/// ``ReduceL1`` wrapper around :class:`ReduceL1L2` configured for the L1 norm.
class ReduceL1 : public ReduceL1L2 {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceL1";
  void Run(RuntimeContext &rt) override;
  explicit ReduceL1(const KernelContext &ctx) : ReduceL1L2(ctx, ReduceL1L2::Mode::kL1) {}
};

/// ``ReduceL2`` wrapper around :class:`ReduceL1L2` configured for the L2 norm.
class ReduceL2 : public ReduceL1L2 {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceL2";
  void Run(RuntimeContext &rt) override;
  explicit ReduceL2(const KernelContext &ctx) : ReduceL1L2(ctx, ReduceL1L2::Mode::kL2) {}
};

/// ``ReduceSumSquare`` wrapper around :class:`ReduceL1L2` configured for the
/// sum of squared values (L2 norm without the final square root).
class ReduceSumSquare : public ReduceL1L2 {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceSumSquare";
  void Run(RuntimeContext &rt) override;
  explicit ReduceSumSquare(const KernelContext &ctx)
      : ReduceL1L2(ctx, ReduceL1L2::Mode::kSumSquare) {}
};

/// Shared FLOAT/DOUBLE kernel for ``ReduceLogSum`` and ``ReduceLogSumExp``.
///
/// ``ReduceLogSum`` computes ``y = log(sum(x, axes))``.
/// ``ReduceLogSumExp`` computes ``y = log(sum(exp(x), axes))`` using the
/// numerically-stable ``max-shift`` trick:
/// ``y = m + log(sum(exp(x - m), axes))`` where ``m = max(x, axes)``.
/// For both operators the empty-set identity is ``-inf`` (``log(0)``); when
/// reducing over a size-0 axis the result is ``-inf``.
///
/// The interface mirrors :class:`ReduceL1L2`: the ``axes`` may either be
/// omitted (reduce-all unless ``noop_with_empty_axes == true``) or supplied
/// as an int64 tensor (opset 18+). Negative axes follow ONNX semantics.
class ReduceLogSumOp : public KernelBase {
public:
  enum class Mode { kLogSum, kLogSumExp };

  ReduceLogSumOp(const KernelContext &ctx, Mode mode) : KernelBase(ctx), mode_(mode) {}

  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes, Tensor &output,
                  RawBufferAllocator *temporary_allocator = nullptr) const;

  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output, RawBufferAllocator *temporary_allocator = nullptr) const;

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

private:
  Mode mode_;
};

/// ``ReduceLogSum`` wrapper around :class:`ReduceLogSumOp` configured for
/// ``log(sum(x, axes))``.
class ReduceLogSum : public ReduceLogSumOp {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceLogSum";
  void Run(RuntimeContext &rt) override;
  explicit ReduceLogSum(const KernelContext &ctx)
      : ReduceLogSumOp(ctx, ReduceLogSumOp::Mode::kLogSum) {}
};

/// ``ReduceLogSumExp`` wrapper around :class:`ReduceLogSumOp` configured for
/// ``log(sum(exp(x), axes))``.
class ReduceLogSumExp : public ReduceLogSumOp {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceLogSumExp";
  void Run(RuntimeContext &rt) override;
  explicit ReduceLogSumExp(const KernelContext &ctx)
      : ReduceLogSumOp(ctx, ReduceLogSumOp::Mode::kLogSumExp) {}
};

/// Arithmetic mean reduction of a FLOAT input ``data`` along the dimensions
/// listed in the optional ``axes`` int64 tensor. ``ReduceMean`` computes
/// ``y = sum(x, axes) / N`` where ``N`` is the product of the reduced
/// dimension sizes. ONNX does not define a value for the empty-set case
/// (division by zero); this kernel returns ``0`` in that case for
/// consistency with the other reduction kernels.
///
/// The interface mirrors :class:`ReduceMinMax` / :class:`ReduceL1L2`: the
/// ``axes`` may either be omitted (reduce-all unless
/// ``noop_with_empty_axes == true``) or supplied as an int64 tensor (opset
/// 18+). Negative axes follow ONNX semantics.
class ReduceMean : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceMean";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

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

/// Product reduction of a FLOAT input ``data`` along the dimensions listed in
/// the optional ``axes`` int64 tensor. ``ReduceProd`` computes
/// ``y = prod(x, axes)``; the empty-set identity is ``1`` (so reducing over a
/// size-0 axis yields a one-filled output, and ``noop_with_empty_axes`` skips
/// the reduction entirely).
///
/// The interface mirrors :class:`ReduceMinMax` / :class:`ReduceL1L2`: the
/// ``axes`` may either be omitted (reduce-all unless
/// ``noop_with_empty_axes == true``) or supplied as an int64 tensor (opset
/// 18+). Negative axes follow ONNX semantics.
class ReduceProd : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReduceProd";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &data, bool keepdims = true, bool noop_with_empty_axes = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

  Tensor operator()(const Tensor &data, const Tensor &axes, bool keepdims = true,
                    bool noop_with_empty_axes = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &axes, bool keepdims, bool noop_with_empty_axes,
                  Tensor &output) const;

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

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