/*!
 * \file tl/op/reduce.h
 * \brief Reduction operators for tensor computations
 */

#ifndef TVM_TL_OP_REDUCE_H_
#define TVM_TL_OP_REDUCE_H_

#include "operator.h"
#include "support/check.h"
#include <tvm/runtime/logging.h>

namespace tvm {

namespace tl {

using namespace tirx;
using namespace ffi;

/// Supported reduction operation types
enum class ReduceTypeEnum : uint8_t {
  kSum,    ///< Sum reduction
  kAbsSum, ///< Absolute sum reduction
  kMax,    ///< Maximum value reduction
  kMin,    ///< Minimum value reduction
  kAbsMax, ///< Maximum absolute value reduction
  kBitAnd, ///< Bitwise and reduction
  kBitOr,  ///< Bitwise or reduction
  kBitXor, ///< Bitwise xor reduction
};

/// Node class representing a reduction type
class ReduceTypeNode : public Object {
public:
  int type{-1}; ///< Internal type identifier
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.ReduceType", ReduceTypeNode, Object);

  static void RegisterReflection() {
    namespace refl = reflection;
    refl::ObjectDef<ReduceTypeNode>().def_ro("type", &ReduceTypeNode::type);
  }

  /// Type checking methods
  bool IsSum() const { return type == int(ReduceTypeEnum::kSum); }
  bool IsAbsSum() const { return type == int(ReduceTypeEnum::kAbsSum); }
  bool IsMax() const { return type == int(ReduceTypeEnum::kMax); }
  bool IsMin() const { return type == int(ReduceTypeEnum::kMin); }
  bool IsAbsMax() const { return type == int(ReduceTypeEnum::kAbsMax); }
  bool IsBitAnd() const { return type == int(ReduceTypeEnum::kBitAnd); }
  bool IsBitOr() const { return type == int(ReduceTypeEnum::kBitOr); }
  bool IsBitXor() const { return type == int(ReduceTypeEnum::kBitXor); }
};

/// Wrapper class for reduction type with string-based construction
class ReduceType : public ObjectRef {
public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ReduceType, ObjectRef,
                                             ReduceTypeNode);
  TVM_DLL ReduceType(std::string type) {
    auto node = make_object<ReduceTypeNode>();
    if (type == "sum") {
      node->type = int(ReduceTypeEnum::kSum);
    } else if (type == "abssum") {
      node->type = int(ReduceTypeEnum::kAbsSum);
    } else if (type == "max") {
      node->type = int(ReduceTypeEnum::kMax);
    } else if (type == "absmax") {
      node->type = int(ReduceTypeEnum::kAbsMax);
    } else if (type == "min") {
      node->type = int(ReduceTypeEnum::kMin);
    } else if (type == "bitand") {
      node->type = int(ReduceTypeEnum::kBitAnd);
    } else if (type == "bitor") {
      node->type = int(ReduceTypeEnum::kBitOr);
    } else if (type == "bitxor") {
      node->type = int(ReduceTypeEnum::kBitXor);
    } else {
      LOG(FATAL) << "Invalid reduce type: " << type;
    }
    data_ = std::move(node);
  }
};

/// Node class for reduction operations
class ReduceOpNode : public TileOperatorNode {
public:
  tirx::Buffer src, dst; ///< Source and destination buffers
  // Optional: keep the original regions used to construct this op
  BufferRegion srcRegion_, dstRegion_;
  int dim;         ///< Dimension to reduce along
  ReduceType type; ///< Type of reduction operation
  bool clear;      ///< Whether to clear destination before reduction
  int batch{1};    ///< Number of output elements per batched AllReduce
                   ///< call. Default 1 = scalar (current behaviour).
                   ///< When batch > 1, the compiler emits
                   ///< ceil(N/batch) batched AllReduce calls each
                   ///< sharing a single pair of barriers across batch
                   ///< elements. batch must evenly divide the
                   ///< per-thread output element count N derived from
                   ///< the fragment layout.
  bool nan_propagate{false}; ///< For fp16/bf16 max/min/absmax: propagate NaN
                             ///< (use __hmax_nan/__hmin_nan) instead of the
                             ///< default __hmax/__hmin which return the
                             ///< non-NaN operand.

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.ReduceOp", ReduceOpNode,
                                    TileOperatorNode);

  static void RegisterReflection() {
    namespace refl = reflection;
    refl::ObjectDef<ReduceOpNode>()
        .def_ro("src", &ReduceOpNode::src)
        .def_ro("dst", &ReduceOpNode::dst)
        .def_ro("srcRegion", &ReduceOpNode::srcRegion_)
        .def_ro("dstRegion", &ReduceOpNode::dstRegion_)
        .def_ro("dim", &ReduceOpNode::dim)
        .def_ro("type", &ReduceOpNode::type)
        .def_ro("clear", &ReduceOpNode::clear)
        .def_ro("batch", &ReduceOpNode::batch)
        .def_ro("nan_propagate", &ReduceOpNode::nan_propagate);
  }

  /// Lower the operator to TIR statements
  Stmt Lower(const LowerArgs &lower_args,
             arith::Analyzer *analyzer) const override;
  /// Infer memory layout for buffers
  LayoutMap InferLayout(const LayoutInferArgs &layout_args,
                        InferLevel level) const override;
  AccessRegions GetAccessRegions() const override;
  static const Op &Get();
  TileOperator Clone() const;
};

using ReduceTargetPredicate = bool (*)(Target target);

struct ReduceImpl {
  const char *name;
  ReduceTargetPredicate match_target;

  Stmt (*lower)(const ReduceOpNode &op, const LowerArgs &lower_args,
                arith::Analyzer *analyzer);
};

void RegisterReduceImpl(ReduceImpl impl);

/// Wrapper class for reduction operations
class ReduceOp : public TileOperator {
public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ReduceOp, TileOperator,
                                             ReduceOpNode);
  TVM_DLL
  ReduceOp(Array<PrimExpr> args,
           Map<String, ObjectRef> annotations = Map<String, ObjectRef>());
  static const Op &Get();
};

} // namespace tl
} // namespace tvm

#endif //  TVM_TL_OP_REDUCE_H_
