#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/DialectConversion.h"

#include "llvm/ADT/SmallVector.h"

#include <algorithm>
#include <optional>
#include <type_traits>

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

struct ReduceMeanSemantics {
  SmallVector<int64_t> axes;
  int64_t keepdims = 1;
  bool isIdentity = false;
};

static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }

static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
  if (!denseAttr)
    return failure();
  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
}

static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
  SmallVector<int64_t> normalizedAxes;
  normalizedAxes.reserve(axes.size());
  for (int64_t axis : axes) {
    auto normalizedAxis = normalizeAxisChecked(axis, rank);
    if (failed(normalizedAxis))
      return failure();
    normalizedAxes.push_back(*normalizedAxis);
  }
  llvm::sort(normalizedAxes);
  normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
  return normalizedAxes;
}

template <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
static FailureOr<ReduceMeanSemantics>
getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
  ReduceMeanSemantics semantics;
  semantics.keepdims = reduceMeanOp.getKeepdims();

  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputRank);
    if (failed(axes))
      return failure();
    semantics.axes = std::move(*axes);
    return semantics;
  }
  else {
    if (isNoneValueLike(adaptor.getAxes())) {
      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
        semantics.isIdentity = true;
        return semantics;
      }

      semantics.axes.reserve(inputRank);
      for (int64_t axis = 0; axis < inputRank; ++axis)
        semantics.axes.push_back(axis);
      return semantics;
    }

    auto axes = getConstantIntValues(adaptor.getAxes());
    if (failed(axes))
      return failure();

    if (axes->empty()) {
      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
        semantics.isIdentity = true;
        return semantics;
      }

      semantics.axes.reserve(inputRank);
      for (int64_t axis = 0; axis < inputRank; ++axis)
        semantics.axes.push_back(axis);
      return semantics;
    }

    auto normalizedAxes = normalizeAxesChecked(*axes, inputRank);
    if (failed(normalizedAxes))
      return failure();
    semantics.axes = std::move(*normalizedAxes);
    return semantics;
  }
}

static SmallVector<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
  SmallVector<bool> reducedAxes(rank, false);
  for (int64_t axis : axes) {
    if (axis < 0 || axis >= rank)
      return {};
    reducedAxes[axis] = true;
  }
  return reducedAxes;
}

static RankedTensorType getAllOnesType(RankedTensorType inputType, Type elementType) {
  return RankedTensorType::get(SmallVector<int64_t>(inputType.getRank(), 1), elementType);
}

static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
    shape.push_back(isReduced ? 1 : dim);
  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
}

static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
    shape.push_back(isReduced ? dim : 1);
  return RankedTensorType::get(shape, inputType.getElementType(), inputType.getEncoding());
}

static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
  return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
}

static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> keptAxes;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes))
    if (!isReduced)
      keptAxes.push_back(static_cast<int64_t>(axis));
  return keptAxes;
}

static Value
computeLaneIndex(Value lane, int64_t stride, int64_t dimSize, ConversionPatternRewriter& rewriter, Location loc) {
  if (dimSize == 1)
    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);

  MLIRContext* context = rewriter.getContext();
  AffineExpr d0 = getAffineDimExpr(0, context);
  AffineExpr expr = d0;
  if (stride != 1)
    expr = expr.floorDiv(stride);
  if (dimSize != 1)
    expr = expr % dimSize;
  return createOrFoldAffineApply(rewriter, loc, expr, ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
}

static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                                     ArrayRef<bool> reducedAxes,
                                                     RankedTensorType batchType,
                                                     RankedTensorType leafType,
                                                     ConversionPatternRewriter& rewriter,
                                                     Location loc) {
  auto inputType = cast<RankedTensorType>(input.getType());
  auto sliceType = getReducedSliceType(inputType, reducedAxes);
  SmallVector<int64_t> keptAxes = getKeptAxes(reducedAxes);

  int64_t laneCount = 1;
  SmallVector<int64_t> keptAxisStrides(keptAxes.size(), 1);
  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1; index >= 0; --index) {
    keptAxisStrides[index] = laneCount;
    int64_t dimSize = inputType.getDimSize(keptAxes[index]);
    if (dimSize <= 0)
      return failure();
    if (laneCount > std::numeric_limits<int32_t>::max() / dimSize)
      return failure();
    laneCount *= dimSize;
  }

  SmallVector<OpFoldResult> sliceOffsets;
  SmallVector<OpFoldResult> sliceSizes;
  SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
  sliceOffsets.reserve(inputType.getRank());
  sliceSizes.reserve(inputType.getRank());

  auto batchOp =
    createSpatComputeBatch(rewriter,
                           loc,
                           TypeRange {batchType},
                           laneCount,
                           {},
                           ValueRange {input},
                           [&](detail::SpatComputeBatchBodyArgs args) {
                             size_t keptAxisIndex = 0;
                             sliceOffsets.clear();
                             sliceSizes.clear();
                             for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                               if (isReduced) {
                                 sliceOffsets.push_back(rewriter.getIndexAttr(0));
                                 sliceSizes.push_back(rewriter.getIndexAttr(inputType.getDimSize(axis)));
                                 continue;
                               }

                               Value axisIndex = computeLaneIndex(
                                 args.lane, keptAxisStrides[keptAxisIndex], inputType.getDimSize(axis), rewriter, loc);
                               ++keptAxisIndex;
                               sliceOffsets.push_back(axisIndex);
                               sliceSizes.push_back(rewriter.getIndexAttr(1));
                             }

                             Value slice = tensor::ExtractSliceOp::create(
                               rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                             Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
                             publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
                           });
  if (failed(batchOp))
    return failure();
  return (*batchOp).getResult(0);
}

static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
    Value batchValue, RankedTensorType keepdimsType,
    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
    Location loc) {
  auto batchType = dyn_cast<RankedTensorType>(batchValue.getType());
  int64_t rank = keepdimsType.getRank();
  if (!batchType || !batchType.hasStaticShape()
      || !keepdimsType.hasStaticShape()
      || static_cast<int64_t>(reducedAxes.size()) != rank
      || batchType.getRank() != rank + 1
      || batchType.getElementType() != keepdimsType.getElementType())
    return failure();

  int64_t laneCount = 1;
  SmallVector<int64_t> keptAxes;
  SmallVector<int64_t> keptAxisStrides;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
    int64_t dim = keepdimsType.getDimSize(axis);
    if (dim <= 0 || (isReduced && dim != 1))
      return failure();
    if (!isReduced)
      keptAxes.push_back(axis);
  }
  keptAxisStrides.resize(keptAxes.size(), 1);
  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1;
       index >= 0; --index) {
    keptAxisStrides[index] = laneCount;
    int64_t dim = keepdimsType.getDimSize(keptAxes[index]);
    if (laneCount > std::numeric_limits<int64_t>::max() / dim)
      return failure();
    laneCount *= dim;
  }
  if (batchType.getDimSize(0) != laneCount
      || llvm::any_of(batchType.getShape().drop_front(),
                      [](int64_t dim) { return dim != 1; }))
    return failure();

  SmallVector<int64_t> operandIndices(laneCount, 0);
  SmallVector<int64_t> sourceSlots;
  SmallVector<int64_t> sourceOffsets(laneCount, 0);
  SmallVector<int64_t> fragmentOffsets;
  sourceSlots.reserve(laneCount);
  fragmentOffsets.reserve(laneCount * rank);
  for (int64_t lane = 0; lane < laneCount; ++lane) {
    sourceSlots.push_back(lane);
    size_t keptAxisIndex = 0;
    for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
      if (isReduced) {
        fragmentOffsets.push_back(0);
        continue;
      }
      int64_t dim = keepdimsType.getDimSize(axis);
      fragmentOffsets.push_back(
        (lane / keptAxisStrides[keptAxisIndex]) % dim);
      ++keptAxisIndex;
    }
  }
  SmallVector<int64_t> fragmentSizes(fragmentOffsets.size(), 1);
  SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
  return spatial::SpatBlueprintOp::create(
           rewriter, loc, keepdimsType, batchValue, ValueRange {},
           rewriter.getStringAttr("nchw"),
           rewriter.getStringAttr("fragmented"),
           rewriter.getDenseI64ArrayAttr(fragmentOffsets),
           rewriter.getDenseI64ArrayAttr(fragmentSizes),
           rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
           rewriter.getStringAttr("fragment_assembly"),
           rewriter.getDenseI64ArrayAttr(operandIndices),
           rewriter.getDenseI64ArrayAttr(sourceSlots),
           rewriter.getDenseI64ArrayAttr(sourceOffsets),
           rewriter.getDenseI64ArrayAttr(fragmentStrides),
           rewriter.getStringAttr("disjoint"),
           rewriter.getStringAttr("complete"))
    .getOutput();
}

static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
  SmallVector<ReassociationIndices> reassociation;
  ReassociationIndices currentGroup;

  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
    currentGroup.push_back(axis);
    if (!isReduced) {
      reassociation.push_back(currentGroup);
      currentGroup.clear();
    }
  }

  if (!currentGroup.empty()) {
    if (reassociation.empty())
      reassociation.push_back(std::move(currentGroup));
    else
      reassociation.back().append(currentGroup.begin(), currentGroup.end());
  }

  return reassociation;
}

static Value squeezeReducedAxes(Value keepdimsValue,
                                RankedTensorType resultType,
                                ArrayRef<bool> reducedAxes,
                                ConversionPatternRewriter& rewriter,
                                Location loc) {
  SmallVector<ReassociationIndices> reassociation =
    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(reducedAxes);
  if (isCompileTimeComputable(keepdimsValue))
    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, keepdimsValue, reassociation).getResult();

  auto squeezeCompute =
    createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {keepdimsValue}, [&](Value input) {
      Value collapsed = tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
      spatial::SpatYieldOp::create(rewriter, loc, collapsed);
    });
  return squeezeCompute.getResult(0);
}

template <typename ReduceMeanOp>
struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
  using OpConversionPattern<ReduceMeanOp>::OpConversionPattern;
  using Adaptor = typename ReduceMeanOp::Adaptor;

  LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp,
                                Adaptor adaptor,
                                ConversionPatternRewriter& rewriter) const override {
    auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
    auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
    if (!inputType || !resultType || !inputType.hasStaticShape() || !resultType.hasStaticShape())
      return failure();
    if (inputType.getRank() == 0) {
      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
      return success();
    }

    auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank());
    if (failed(semantics))
      return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes");
    if (semantics->isIdentity) {
      if (inputType != resultType)
        return rewriter.notifyMatchFailure(
          reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type");
      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
      return success();
    }

    SmallVector<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
    if (reducedAxes.empty() && inputType.getRank() != 0)
      return failure();

    Location loc = reduceMeanOp.getLoc();
    RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
    RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
    int64_t laneCount = 1;
    for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
      if (isReduced)
        continue;
      if (dim <= 0 || laneCount > std::numeric_limits<int32_t>::max() / dim)
        return rewriter.notifyMatchFailure(
          reduceMeanOp, "ReduceMean physical lane count is not representable");
      laneCount *= dim;
    }
    RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);

    auto lanePackedKeepdims =
      buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
    if (failed(lanePackedKeepdims))
      return failure();
    auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint(
      *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc);
    if (failed(reducedKeepdims))
      return rewriter.notifyMatchFailure(
        reduceMeanOp,
        "cannot build physical-fragment ReduceMean keepdims reconstruction");

    if (semantics->keepdims != 0) {
      rewriter.replaceOp(reduceMeanOp, *reducedKeepdims);
      return success();
    }

    Value reduced = squeezeReducedAxes(
      *reducedKeepdims, resultType, reducedAxes, rewriter, loc);
    rewriter.replaceOp(reduceMeanOp, reduced);
    return success();
  }
};

} // namespace

void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.add<ReduceMeanToSpatialCompute<ONNXReduceMeanV13Op>, ReduceMeanToSpatialCompute<ONNXReduceMeanOp>>(ctx);
}

} // namespace onnx_mlir