#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/SmallVector.h"

#include <cassert>

#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

struct ConvToGemm : OpConversionPattern<ONNXConvOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(ONNXConvOp convOp,
                                ONNXConvOpAdaptor convOpAdaptor,
                                ConversionPatternRewriter& rewriter) const override;
};

} // namespace

LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
                                          ONNXConvOpAdaptor convOpAdaptor,
                                          ConversionPatternRewriter& rewriter) const {
  Location loc = convOp.getLoc();
  Value x = convOpAdaptor.getX();
  Value w = convOpAdaptor.getW();
  Value b = convOpAdaptor.getB();

  auto xType = cast<RankedTensorType>(x.getType());
  auto wType = cast<RankedTensorType>(w.getType());
  auto outType = cast<RankedTensorType>(convOp.getY().getType());

  assert("Only support static shapes" && xType.hasStaticShape() && wType.hasStaticShape() && outType.hasStaticShape());
  assert("Only support 2D convolution" && xType.getRank() == 4);

  // We need to understand what is group
  assert("Only support group=1" && convOp.getGroup() == 1);

  const int64_t batchSize = xType.getDimSize(0);
  const int64_t numChannelsIn = xType.getDimSize(1);
  const int64_t xHeight = xType.getDimSize(2);
  const int64_t xWidth = xType.getDimSize(3);
  const int64_t numChannelsOut = wType.getDimSize(0);
  const int64_t wHeight = wType.getDimSize(2);
  const int64_t wWidth = wType.getDimSize(3);
  const int64_t outHeight = outType.getDimSize(2);
  const int64_t outWidth = outType.getDimSize(3);

  // Read optional conv attributes (ONNX defaults: stride=1, dilation=1, pad=0)
  auto getI64 = [](ArrayAttr arr, size_t idx) -> int64_t { return cast<IntegerAttr>(arr[idx]).getInt(); };

  const auto stridesAttr = convOp.getStrides();
  const auto dilationsAttr = convOp.getDilations();
  const auto padsAttr = convOp.getPads();

  const int64_t strideHeight = stridesAttr ? getI64(*stridesAttr, 0) : 1;
  const int64_t strideWidth = stridesAttr ? getI64(*stridesAttr, 1) : 1;
  const int64_t dilationHeight = dilationsAttr ? getI64(*dilationsAttr, 0) : 1;
  const int64_t dilationWidth = dilationsAttr ? getI64(*dilationsAttr, 1) : 1;

  int64_t padHeightBegin = 0;
  int64_t padHeightEnd = 0;
  int64_t padWidthBegin = 0;
  int64_t padWidthEnd = 0;

  if (padsAttr) {
    padHeightBegin = getI64(*padsAttr, 0);
    padWidthBegin = getI64(*padsAttr, 1);
    padHeightEnd = getI64(*padsAttr, 2);
    padWidthEnd = getI64(*padsAttr, 3);
  }
  else {
    // Compute padding from auto_pad attribute
    const auto autoPad = convOp.getAutoPad();
    if (autoPad == "SAME_UPPER" || autoPad == "SAME_LOWER") {
      const int64_t effectiveKernelH = (wHeight - 1) * dilationHeight + 1;
      const int64_t effectiveKernelW = (wWidth - 1) * dilationWidth + 1;
      const int64_t totalPadH =
        std::max(static_cast<int64_t>(0), (outHeight - 1) * strideHeight + effectiveKernelH - xHeight);
      const int64_t totalPadW =
        std::max(static_cast<int64_t>(0), (outWidth - 1) * strideWidth + effectiveKernelW - xWidth);

      if (autoPad == "SAME_UPPER") {
        padHeightBegin = totalPadH / 2;
        padHeightEnd = totalPadH - padHeightBegin;
        padWidthBegin = totalPadW / 2;
        padWidthEnd = totalPadW - padWidthBegin;
      }
      else { // SAME_LOWER
        padHeightEnd = totalPadH / 2;
        padHeightBegin = totalPadH - padHeightEnd;
        padWidthEnd = totalPadW / 2;
        padWidthBegin = totalPadW - padWidthEnd;
      }
    }
    // "NOTSET" or "VALID" -> all pads stay 0
  }

  // im2col layout (flipped with respect to the standard, so filters sit in B = crossbar):
  //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
  //   B (weights): [patchSize, cOut]        -- W^T, stored in crossbar columns
  //   Gemm output: [numPatches, cOut]
  const int64_t patchSize = numChannelsIn * wHeight * wWidth;
  const int64_t numPatchesPerBatch = outHeight * outWidth;
  const int64_t numPatches = batchSize * numPatchesPerBatch;

  auto elemType = xType.getElementType();
  auto im2colType = RankedTensorType::get({numPatches, patchSize}, elemType);
  auto rowType = RankedTensorType::get({1, patchSize}, elemType);
  auto wFlatType = RankedTensorType::get({numChannelsOut, patchSize}, wType.getElementType());
  auto wTransType = RankedTensorType::get({patchSize, numChannelsOut}, wType.getElementType());
  auto gemmOutType = RankedTensorType::get({numPatches, numChannelsOut}, outType.getElementType());
  auto nhwcType = RankedTensorType::get({batchSize, outHeight, outWidth, numChannelsOut}, outType.getElementType());

  // Prepare weight matrix W for crossbar storage:
  // W: [numChannelsOut, numChannelsIn, wHeight, wWidth] -> [numChannelsOut, patchSize] -> [patchSize, numChannelsOut]
  Value wFlat = tensor::CollapseShapeOp::create(rewriter,
                                                loc,
                                                wFlatType,
                                                w,
                                                SmallVector<ReassociationIndices> {
                                                  {0},
                                                  {1, 2, 3}
  });
  Value wTrans = ONNXTransposeOp::create(rewriter, loc, wTransType, wFlat, rewriter.getI64ArrayAttr({1, 0}));

  // Pass bias through directly; Gemm handles rank-1 C canonicalization.
  bool hasB = !isa<ONNXNoneOp>(b.getDefiningOp());
  Value gemmC;
  if (hasB)
    gemmC = b;
  else
    gemmC = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());

  auto im2colComputeOp =
    spatial::SpatWeightedCompute::create(rewriter, loc, im2colType, SmallVector<Value>(), ValueRange {x});

  auto* im2colBlock = new Block();
  im2colBlock->addArgument(x.getType(), loc);
  im2colComputeOp.getBody().push_back(im2colBlock);
  rewriter.setInsertionPointToStart(im2colBlock);

  Value paddedInput = im2colBlock->getArgument(0);

  // Pad input with zeros if needed:
  // [1, numChannelsIn, xHeight, xWidth] -> [1, numChannelsIn, xHeight+padHeight, xWidth+padWidth]
  if (padHeightBegin || padHeightEnd || padWidthBegin || padWidthEnd) {
    const int64_t paddedHeight = xHeight + padHeightBegin + padHeightEnd;
    const int64_t paddedWidth = xWidth + padWidthBegin + padWidthEnd;
    auto paddedType = RankedTensorType::get({batchSize, numChannelsIn, paddedHeight, paddedWidth}, elemType);
    SmallVector<OpFoldResult> lowPads = {rewriter.getIndexAttr(0),
                                         rewriter.getIndexAttr(0),
                                         rewriter.getIndexAttr(padHeightBegin),
                                         rewriter.getIndexAttr(padWidthBegin)};
    SmallVector<OpFoldResult> highPads = {rewriter.getIndexAttr(0),
                                          rewriter.getIndexAttr(0),
                                          rewriter.getIndexAttr(padHeightEnd),
                                          rewriter.getIndexAttr(padWidthEnd)};
    auto padOp = tensor::PadOp::create(rewriter, loc, paddedType, paddedInput, lowPads, highPads);
    auto* padBlock = new Block();
    for (int i = 0; i < 4; i++)
      padBlock->addArgument(rewriter.getIndexType(), loc);
    padOp.getRegion().push_back(padBlock);
    rewriter.setInsertionPointToStart(padBlock);
    auto zero = arith::ConstantOp::create(rewriter, loc, elemType, rewriter.getFloatAttr(elemType, 0.0));
    tensor::YieldOp::create(rewriter, loc, zero.getResult());
    rewriter.setInsertionPointAfter(padOp);
    paddedInput = padOp.getResult();
  }

  // Build im2col [numPatches, patchSize]:
  // For each batch/output position (n, oh, ow), extract the patch from x
  SmallVector<Value> im2colRows;
  im2colRows.reserve(numPatches);
  for (int64_t n = 0; n < batchSize; n++) {
    for (int64_t oh = 0; oh < outHeight; oh++) {
      for (int64_t ow = 0; ow < outWidth; ow++) {
        SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(n),
                                             rewriter.getIndexAttr(0),
                                             rewriter.getIndexAttr(oh * strideHeight),
                                             rewriter.getIndexAttr(ow * strideWidth)};
        SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
                                           rewriter.getIndexAttr(numChannelsIn),
                                           rewriter.getIndexAttr(wHeight),
                                           rewriter.getIndexAttr(wWidth)};
        SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1),
                                             rewriter.getIndexAttr(1),
                                             rewriter.getIndexAttr(dilationHeight),
                                             rewriter.getIndexAttr(dilationWidth)};
        auto patchType = RankedTensorType::get({1, numChannelsIn, wHeight, wWidth}, elemType);
        Value patch = tensor::ExtractSliceOp::create(rewriter, loc, patchType, paddedInput, offsets, sizes, strides);

        // Flatten [1, numChannelsIn, wHeight, wWidth] -> [1, patchSize]
        Value row = tensor::CollapseShapeOp::create(rewriter,
                                                    loc,
                                                    rowType,
                                                    patch,
                                                    SmallVector<ReassociationIndices> {
                                                      {0},
                                                      {1, 2, 3}
        });
        im2colRows.push_back(row);
      }
    }
  }

  // Concatenate all rows: [numPatches, patchSize]
  Value im2col = tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, im2colRows);
  spatial::SpatYieldOp::create(rewriter, loc, im2col);

  rewriter.setInsertionPointAfter(im2colComputeOp);

  // Gemm: A @ B + C = im2col @ W^T + b
  // [numPatches, patchSize] @ [patchSize, numChannelsOut] + [1, numChannelsOut] -> [numPatches, numChannelsOut]
  auto gemmOp = ONNXGemmOp::create(rewriter,
                                   loc,
                                   gemmOutType,
                                   im2colComputeOp.getResult(0),
                                   wTrans,
                                   gemmC,
                                   rewriter.getF32FloatAttr(1.0f),
                                   rewriter.getF32FloatAttr(1.0f),
                                   rewriter.getBoolAttr(false),
                                   rewriter.getBoolAttr(false));
  Value gemmOut = gemmOp.getY();

  auto collectComputeOp =
    spatial::SpatWeightedCompute::create(rewriter, loc, convOp.getType(), SmallVector<Value>(), ValueRange {gemmOut});

  auto* collectBlock = new Block();
  collectBlock->addArgument(gemmOut.getType(), loc);
  collectComputeOp.getBody().push_back(collectBlock);
  rewriter.setInsertionPointToStart(collectBlock);

  auto gemmOutArg = collectBlock->getArguments().front();

  // Restore to NCHW layout:
  // [numPatches, numChannelsOut]
  // -> [1, outHeight, outWidth, numChannelsOut]
  // -> [1, numChannelsOut, outHeight, outWidth]
  Value nhwcOut = tensor::ExpandShapeOp::create(rewriter,
                                                loc,
                                                nhwcType,
                                                gemmOutArg,
                                                SmallVector<ReassociationIndices> {
                                                  {0, 1, 2},
                                                  {3}
  });
  Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2}));

  spatial::SpatYieldOp::create(rewriter, loc, nchwOut);

  rewriter.replaceOp(convOp, collectComputeOp.getResult(0));
  return success();
}

void populateConvPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert<ConvToGemm>(ctx); }

} // namespace onnx_mlir
