#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Value.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"

#include <cassert>
#include <utility>

#include "Common/IR/ShapeUtils.hpp"
#include "Common/IR/ConstantUtils.hpp"
#include "Common/PimCommon.hpp"
#include "Common/Support/CheckedArithmetic.hpp"
#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
#include "Conversion/ONNXToSpatial/Common/Common.hpp"
#include "Conversion/SpatialToPim/Common.hpp"
#include "Conversion/SpatialToPim/Patterns.hpp"
#include "Dialect/Pim/PimOps.hpp"
#include "Dialect/Spatial/SpatialOps.hpp"
#include "Pass/PIMPasses.h"
#include "SpatialToPimPass.hpp"

using namespace mlir;
using namespace onnx_mlir;
using namespace pim;

namespace onnx_mlir {

static FailureOr<Value>
createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTensorType resultType) {
  auto sourceType = cast<RankedTensorType>(value.getType());
  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> highPads;
  highPads.reserve(sourceType.getRank());
  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));

  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
  auto* padBlock = new Block();
  for (int64_t i = 0; i < sourceType.getRank(); ++i)
    padBlock->addArgument(rewriter.getIndexType(), loc);
  padOp.getRegion().push_back(padBlock);
  rewriter.setInsertionPointToStart(padBlock);
  auto zero = getOrCreateConstant(
    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
  tensor::YieldOp::create(rewriter, loc, zero);
  rewriter.setInsertionPointAfter(padOp);
  return padOp.getResult();
}

static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
  auto vectorType = cast<RankedTensorType>(vector.getType());
  ArrayRef<int64_t> shape = vectorType.getShape();
  assert(isHVectorShape(shape) && "expected a horizontal vector");
  assert(shape[1] <= static_cast<int64_t>(crossbarSize) && "vector width must fit in one crossbar");

  if (shape[1] == static_cast<int64_t>(crossbarSize))
    return vector;

  auto paddedType = RankedTensorType::get(
    {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
  return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
}

void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
  outputTensors.clear();
  operationsToRemove.clear();
  ModuleOp moduleOp = getOperation();
  MLIRContext* ctx = moduleOp.getContext();

  auto entryFunc = getPimEntryFunc(moduleOp);
  if (failed(entryFunc)) {
    moduleOp.emitError("failed to locate the PIM entry function during Spatial-to-PIM lowering");
    signalPassFailure();
    return;
  }
  func::FuncOp funcOp = *entryFunc;
  if (failed(verifyScheduledSpatialInvariants(funcOp))) {
    funcOp.emitOpError(
      "scheduled Spatial verification failed at the start of SpatialToPim");
    signalPassFailure();
    return;
  }

  IRRewriter rewriter(&getContext());
  OperationFolder constantFolder(&getContext());

  ConversionTarget target(*ctx);
  target.addLegalDialect<affine::AffineDialect,
                         PimDialect,
                         tensor::TensorDialect,
                         arith::ArithDialect,
                         bufferization::BufferizationDialect,
                         func::FuncDialect,
                         memref::MemRefDialect,
                         scf::SCFDialect,
                         BuiltinDialect>();
  target.addLegalOp<spatial::SpatConcatOp,
                    spatial::SpatChannelReceiveOp,
                    spatial::SpatChannelSendOp,
                    spatial::SpatExtractRowsOp>();

  RewritePatternSet initialPatterns(ctx);
  populateInitialPatterns(initialPatterns);
  if (failed(applyPartialConversion(moduleOp, target, std::move(initialPatterns)))) {
    moduleOp.emitError("failed to lower required Spatial ops to the initial PIM form");
    signalPassFailure();
    return;
  }

  RewritePatternSet globalTensorPatterns(ctx);
  populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
  walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));

  auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
  addReturnOutputBuffers(returnOp, rewriter);
  if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
    funcOp.emitOpError("failed to allocate or initialize core-local tensors during Spatial-to-PIM lowering");
    signalPassFailure();
    return;
  }

  for (auto computeOp : funcOp.getOps<spatial::SpatScheduledCompute>()) {
    markOpToRemove(computeOp);
    if (failed(lowerComputeOp(computeOp, rewriter, constantFolder))) {
      computeOp.emitOpError("failed to lower spat.scheduled_compute to pim.core");
      signalPassFailure();
      return;
    }
  }

  for (auto computeBatchOp : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
    markOpToRemove(computeBatchOp);
    if (failed(lowerComputeBatchOp(computeBatchOp, rewriter))) {
      computeBatchOp.emitOpError("failed to lower spat.scheduled_compute_batch to pim.core_batch");
      signalPassFailure();
      return;
    }
  }

  RewritePatternSet initialTensorPackingPatterns(ctx);
  populateTensorPackingPatterns(initialTensorPackingPatterns);
  walkAndApplyPatterns(funcOp, std::move(initialTensorPackingPatterns));
  eraseUnusedTensorPackingOps(funcOp, rewriter);

  SmallVector<spatial::SpatChannelReceiveOp> receiveOps;
  for (auto op : funcOp.getOps<spatial::SpatChannelReceiveOp>())
    receiveOps.push_back(op);
  for (auto receiveOp : receiveOps) {
    bool onlyPendingRemovalUsers = llvm::all_of(
      receiveOp->getUsers(), [&](Operation* user) { return llvm::is_contained(operationsToRemove, user); });
    if (onlyPendingRemovalUsers) {
      markOpToRemove(receiveOp);
      continue;
    }
  }

  RewritePatternSet coreBodyPatterns(ctx);
  populateCoreBodyPatterns(coreBodyPatterns);
  populateAffineToStdConversionPatterns(coreBodyPatterns);
  FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));

  ConversionTarget coreBodyTarget(*ctx);
  coreBodyTarget.addLegalDialect<PimDialect,
                                 tensor::TensorDialect,
                                 arith::ArithDialect,
                                 bufferization::BufferizationDialect,
                                 func::FuncDialect,
                                 memref::MemRefDialect,
                                 scf::SCFDialect,
                                 BuiltinDialect>();
  coreBodyTarget.addLegalOp<spatial::SpatConcatOp,
                            spatial::SpatChannelReceiveOp,
                            spatial::SpatChannelSendOp,
                            spatial::SpatExtractRowsOp>();

  SmallVector<pim::PimCoreOp> coreOps;
  funcOp.walk([&](pim::PimCoreOp coreOp) { coreOps.push_back(coreOp); });
  for (auto coreOp : coreOps) {
    if (failed(applyFullConversion(coreOp.getOperation(), coreBodyTarget, frozenCoreBodyPatterns))) {
      coreOp.emitOpError("failed to convert nested Spatial ops inside pim.core");
      signalPassFailure();
      return;
    }
  }

  SmallVector<pim::PimCoreBatchOp> coreBatchOps;
  funcOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { coreBatchOps.push_back(coreBatchOp); });
  for (auto coreBatchOp : coreBatchOps) {
    if (failed(applyFullConversion(coreBatchOp.getOperation(), coreBodyTarget, frozenCoreBodyPatterns))) {
      coreBatchOp.emitOpError("failed to convert nested Spatial ops inside pim.core_batch");
      signalPassFailure();
      return;
    }
  }

  if (failed(enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter))) {
    funcOp.emitOpError("failed to enlarge VMM output tensors to crossbar size");
    signalPassFailure();
    return;
  }
  replaceReturnWithOutputBuffers(returnOp, rewriter);
  eraseOpsToRemove();

  RewritePatternSet finalTensorPackingPatterns(ctx);
  populateTensorPackingPatterns(finalTensorPackingPatterns);
  walkAndApplyPatterns(funcOp, std::move(finalTensorPackingPatterns));
  eraseUnusedTensorPackingOps(funcOp, rewriter);

  ConversionTarget communicationTarget(*ctx);
  communicationTarget.addLegalDialect<PimDialect,
                                      tensor::TensorDialect,
                                      arith::ArithDialect,
                                      bufferization::BufferizationDialect,
                                      func::FuncDialect,
                                      memref::MemRefDialect,
                                      scf::SCFDialect,
                                      BuiltinDialect>();
  communicationTarget.addLegalOp<ModuleOp>();
  communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                   spatial::SpatChannelReceiveOp,
                                   spatial::SpatChannelSendOp,
                                   spatial::SpatExtractRowsOp>();

  RewritePatternSet communicationPatterns(ctx);
  populateChannelLoweringPatterns(communicationPatterns);
  if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
    funcOp.emitOpError("failed to lower Spatial communication ops to PIM communication ops");
    signalPassFailure();
    return;
  }
  hoistAndUniquifyIndexConstants(funcOp, rewriter);

  // Dump to file for debug
  dumpModule(moduleOp, "pim0");
}

LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
  bool hasFailure = false;
  funcOp.walk([&](PimVMMOp vmmOp) {
    auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
    ArrayRef<int64_t> outputShape = outputType.getShape();
    assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
    assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");

    rewriter.setInsertionPoint(vmmOp);
    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
    if (failed(paddedInput)) {
      hasFailure = true;
      return WalkResult::interrupt();
    }
    auto paddedOutputType = RankedTensorType::get(
      {outputShape[0], static_cast<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
    Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(crossbarSize)
                               ? vmmOp.getOutputBuffer()
                               : createEmptyTensorFromShaped(rewriter, vmmOp.getLoc(), paddedOutputType).getResult();
    vmmOp.getInputMutable().assign(*paddedInput);
    vmmOp.getOutputBufferMutable().assign(paddedOutputBuffer);

    vmmOp.getOutput().setType(paddedOutputType);

    if (outputShape[1] == static_cast<int64_t>(crossbarSize))
      return WalkResult::advance();

    SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
    SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(outputShape[0]), rewriter.getIndexAttr(outputShape[1])};
    SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
    rewriter.setInsertionPointAfter(vmmOp);
    auto sliceOp =
      tensor::ExtractSliceOp::create(rewriter, vmmOp.getLoc(), outputType, vmmOp.getOutput(), offsets, sizes, strides);
    SmallPtrSet<Operation*, 2> exceptions = {vmmOp, sliceOp};
    vmmOp.getOutput().replaceAllUsesExcept(sliceOp.getResult(), exceptions);
    return WalkResult::advance();
  });
  return success(!hasFailure);
}

LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp,
                                                                                IRRewriter& rewriter) {
  Location loc = funcOp.getLoc();
  OperationFolder constantFolder(funcOp.getContext());
  bool hasFailure = false;

  auto insertMemCopyHostToDev = [&](Value inputTensor, int64_t elementsOffset) {
    auto tensorType = cast<ShapedType>(inputTensor.getType());
    Type elementType = tensorType.getElementType();
    if (!hasByteSizedElementType(elementType))
      return;
    size_t elementByteSize = getElementTypeSizeInBytes(elementType);
    rewriter.setInsertionPointAfter(inputTensor.getDefiningOp());

    auto deviceTensor = tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), elementType);
    auto offsetBytes = pim::checkedMul(
      static_cast<size_t>(elementsOffset), elementByteSize, deviceTensor.getOperation(), "host input byte offset");
    auto byteSize =
      pim::getCheckedShapedTypeSizeInBytes(tensorType, deviceTensor.getOperation(), "host input copy byte size");
    auto sizeAttr =
      succeeded(byteSize)
        ? pim::getCheckedI32Attr(rewriter, deviceTensor.getOperation(), *byteSize, "host input copy byte size")
        : FailureOr<IntegerAttr>(failure());
    if (failed(offsetBytes) || failed(sizeAttr)) {
      hasFailure = true;
      return;
    }

    auto memCopyHostToDevOp = PimMemCopyHostToDevOp::create(
      rewriter,
      loc,
      tensorType,
      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), static_cast<int64_t>(*offsetBytes)),
      deviceTensor,
      inputTensor,
      *sizeAttr);

    rewriter.replaceAllUsesExcept(inputTensor, memCopyHostToDevOp.getResult(), {memCopyHostToDevOp});
  };

  for (auto& op : funcOp.getBody().getOps())
    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
      if (!computeOp.getInputs().empty() || computeOp.getBody().front().getNumArguments() != 0)
        continue;
      for (auto getGlobal : computeOp.getOps<memref::GetGlobalOp>()) {
        if (getGlobal.getName().starts_with("arg") || getGlobal.getName().starts_with("const_")) {
          assert(getGlobal->hasOneUse() && "global must have a single entry point in the compute");
          auto toTensorOpValue = *getGlobal->getUsers().begin()->getResults().begin();
          insertMemCopyHostToDev(toTensorOpValue, 0);
        }
      }
    }

  return success(!hasFailure);
}

void raptor::SpatialToPimPass::markOpToRemove(Operation* op) {
  if (!llvm::is_contained(operationsToRemove, op))
    operationsToRemove.push_back(op);
}

void raptor::SpatialToPimPass::eraseOpsToRemove() {
  for (Operation* op : operationsToRemove) {
    op->dropAllUses();
    op->erase();
  }
}

std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<raptor::SpatialToPimPass>(); }

} // namespace onnx_mlir