#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/DialectConversion.h"

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }

static SmallVector<int64_t> permuteShape(ArrayRef<int64_t> shape, ArrayRef<int64_t> permutation) {
  SmallVector<int64_t> permutedShape;
  permutedShape.reserve(permutation.size());
  for (int64_t axis : permutation)
    permutedShape.push_back(shape[axis]);
  return permutedShape;
}

static Value buildLoopSoftmaxSlice(Value input,
                                   Value accumulator,
                                   RankedTensorType inputType,
                                   ArrayRef<Value> outerIndices,
                                   ConversionPatternRewriter& rewriter,
                                   Location loc) {
  int64_t rank = inputType.getRank();
  SmallVector<int64_t> sliceShape(static_cast<size_t>(rank - 1), 1);
  sliceShape.push_back(inputType.getDimSize(rank - 1));
  auto sliceType = RankedTensorType::get(sliceShape, inputType.getElementType(), inputType.getEncoding());

  SmallVector<OpFoldResult> offsets;
  SmallVector<OpFoldResult> sizes;
  SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
  offsets.reserve(rank);
  sizes.reserve(rank);

  for (Value outerIndex : outerIndices) {
    offsets.push_back(outerIndex);
    sizes.push_back(rewriter.getIndexAttr(1));
  }
  offsets.push_back(rewriter.getIndexAttr(0));
  sizes.push_back(rewriter.getIndexAttr(inputType.getDimSize(rank - 1)));

  Value inputSlice = tensor::ExtractSliceOp::create(rewriter, loc, sliceType, input, offsets, sizes, strides);
  Value softmaxSlice = spatial::SpatSoftmaxOp::create(rewriter, loc, sliceType, inputSlice).getResult();
  return tensor::InsertSliceOp::create(rewriter, loc, softmaxSlice, accumulator, offsets, sizes, strides);
}

static Value buildLoopSoftmaxNest(Value input,
                                  Value accumulator,
                                  RankedTensorType inputType,
                                  int64_t axis,
                                  SmallVectorImpl<Value>& outerIndices,
                                  ConversionPatternRewriter& rewriter,
                                  Location loc) {
  if (axis == inputType.getRank() - 1)
    return buildLoopSoftmaxSlice(input, accumulator, inputType, outerIndices, rewriter, loc);

  Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
  Value c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
  Value cUpper = arith::ConstantIndexOp::create(rewriter, loc, inputType.getDimSize(axis));

  auto loop = scf::ForOp::create(rewriter, loc, c0, cUpper, c1, ValueRange {accumulator});
  rewriter.setInsertionPointToStart(loop.getBody());

  Value loopIndex = loop.getInductionVar();
  Value loopAccumulator = loop.getRegionIterArgs().front();
  outerIndices.push_back(loopIndex);
  Value updatedAccumulator =
    buildLoopSoftmaxNest(input, loopAccumulator, inputType, axis + 1, outerIndices, rewriter, loc);
  outerIndices.pop_back();

  scf::YieldOp::create(rewriter, loc, updatedAccumulator);
  rewriter.setInsertionPointAfter(loop);
  return loop.getResult(0);
}

static Value createLoopSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
  auto inputType = cast<RankedTensorType>(input.getType());
  constexpr size_t numInputs = 1;
  auto computeOp =
    createSpatCompute<numInputs>(rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) {
      if (inputType.getRank() == 1) {
        Value softmax = spatial::SpatSoftmaxOp::create(rewriter, loc, inputType, x).getResult();
        spatial::SpatYieldOp::create(rewriter, loc, softmax);
        return;
      }

      Value outputInit = tensor::EmptyOp::create(rewriter, loc, inputType.getShape(), inputType.getElementType());
      SmallVector<Value> outerIndices;
      Value result = buildLoopSoftmaxNest(x, outputInit, inputType, /*axis=*/0, outerIndices, rewriter, loc);
      spatial::SpatYieldOp::create(rewriter, loc, result);
    });
  return computeOp.getResult(0);
}

struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(ONNXSoftmaxOp softmaxOp,
                                ONNXSoftmaxOpAdaptor adaptor,
                                ConversionPatternRewriter& rewriter) const override {
    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
    if (!inputType || !inputType.hasStaticShape())
      return failure();

    int64_t axis = normalizeAxis(softmaxOp.getAxis(), inputType.getRank());
    if (axis < 0 || axis >= inputType.getRank())
      return failure();

    Value input = adaptor.getInput();
    Value result;
    if (axis == inputType.getRank() - 1) {
      result = createLoopSoftmaxCompute(input, rewriter, softmaxOp.getLoc());
    }
    else {
      SmallVector<int64_t> permutation;
      permutation.reserve(inputType.getRank());
      for (int64_t dim = 0; dim < inputType.getRank(); ++dim)
        if (dim != axis)
          permutation.push_back(dim);
      permutation.push_back(axis);

      SmallVector<int64_t> inversePermutation(inputType.getRank());
      for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
        inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);

      auto transposedType = RankedTensorType::get(
        permuteShape(inputType.getShape(), permutation), inputType.getElementType(), inputType.getEncoding());
      auto preTransposeCompute =
        createSpatCompute<1>(rewriter, softmaxOp.getLoc(), TypeRange {transposedType}, {}, input, [&](Value x) {
          Value transposed = ONNXTransposeOp::create(
            rewriter, softmaxOp.getLoc(), transposedType, x, rewriter.getI64ArrayAttr(permutation));
          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
        });
      Value transposedInput = preTransposeCompute.getResult(0);
      Value transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
      auto postTransposeCompute =
        createSpatCompute<1>(rewriter, softmaxOp.getLoc(), TypeRange {inputType}, {}, transposedResult, [&](Value x) {
          Value transposed = ONNXTransposeOp::create(
            rewriter, softmaxOp.getLoc(), inputType, x, rewriter.getI64ArrayAttr(inversePermutation));
          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
        });
      result = postTransposeCompute.getResult(0);
    }

    rewriter.replaceOp(softmaxOp, result);
    return success();
  }
};

} // namespace

void populateSoftmaxPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.add<SoftmaxToSpatialCompute>(ctx);
}

} // namespace onnx_mlir