#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Pass/Pass.h"

#include "Common/PimCommon.hpp"
#include "Compiler/PimCodeGen.hpp"
#include "Dialect/Pim/PimOps.hpp"
#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
#include "src/Compiler/CompilerOptions.hpp"

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

namespace onnx_mlir {

namespace {

#include "Dialect/Pim/Transforms/Bufferization/PimBufferization.hpp.inc"

struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
  StringRef getArgument() const override { return "bufferize-pim"; }
  StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; }

  PimBufferizationPass() = default;
  PimBufferizationPass(const PimBufferizationPass& pass) {}

  void runOnOperation() final;

private:
  void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
};

} // namespace

void PimBufferizationPass::runOnOperation() {
  auto moduleOp = getOperation();

  // One-Shot-Bufferization
  bufferization::OneShotBufferizationOptions options;
  options.allowUnknownOps = true;
  bufferization::BufferizationState state;
  if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
    signalPassFailure();
  }

  MLIRContext* ctx = moduleOp.getContext();
  ConversionTarget target(*ctx);
  target.addLegalDialect<PimDialect>();

  RewritePatternSet patterns(ctx);
  populateWithGenerated(patterns);

  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
    signalPassFailure();
    return;
  }

  // Remove toTensor operations: leave memrefs instead
  moduleOp.walk([](bufferization::ToTensorOp toTensorOp) {
    toTensorOp.replaceAllUsesWith(toTensorOp.getBuffer());
    toTensorOp.erase();
  });

  // Change main function return types from tensors to memrefs
  func::FuncOp funcOp;
  for (Operation& op : moduleOp.getBody()->getOperations())
    if ((funcOp = dyn_cast<func::FuncOp>(&op)))
      break;
  auto oldFuncType = funcOp.getFunctionType();
  SmallVector<Type> newResults;
  bool changed = false;
  for (Type type : oldFuncType.getResults())
    if (auto tensorType = dyn_cast<RankedTensorType>(type)) {
      newResults.push_back(MemRefType::get(tensorType.getShape(), tensorType.getElementType()));
      changed = true;
    }
    else
      newResults.push_back(type);
  if (changed)
    funcOp.setType(FunctionType::get(funcOp.getContext(), oldFuncType.getInputs(), newResults));

  annotateWeightsMemrefs(moduleOp, funcOp);

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

void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
    bool isAlwaysWeight = !getGlobalOp->getUsers().empty()
                       && all_of(getGlobalOp->getUsers(), [](auto user) -> bool { return isa<PimCoreOp>(user); });
    if (isAlwaysWeight) {
      auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
      assert("Weights must be constants" && globalMemrefOp.getConstant());
      markWeightAlways(getGlobalOp);
      markWeightAlways(globalMemrefOp);
    }
  });
}

std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }

} // namespace onnx_mlir
