#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/Support/Debug.h"

#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/SpatialBufferizableOpInterface.hpp"
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
#include "src/Accelerators/PIM/PimAccelerator.hpp"

#define DEBUG_TYPE "PimAccelerator"

namespace onnx_mlir {
namespace accel {

Accelerator* createPIM() { return PimAccelerator::getInstance(); }

PimAccelerator* PimAccelerator::instance = nullptr;

PimAccelerator* PimAccelerator::getInstance() {
  if (instance == nullptr)
    instance = new PimAccelerator();
  return instance;
}

PimAccelerator::PimAccelerator()
: Accelerator(Kind::PIM) {
  LLVM_DEBUG(llvm::dbgs() << "Creating a PIM accelerator\n");
  acceleratorTargets.push_back(this);
}

uint64_t PimAccelerator::getVersionNumber() const { return 0x000001; }

void PimAccelerator::addPasses(mlir::OwningOpRef<mlir::ModuleOp>& module,
                               mlir::PassManager& pm,
                               EmissionTargetType& emissionTarget,
                               std::string outputNameNoExt) const {
  LLVM_DEBUG(llvm::dbgs() << "Adding passes for PIM accelerator\n");
  addPassesPim(module, pm, emissionTarget, outputNameNoExt);
}

void PimAccelerator::registerDialects(mlir::DialectRegistry& registry) const {
  LLVM_DEBUG(llvm::dbgs() << "Registering dialects for PIM accelerator\n");
  registry.insert<mlir::tensor::TensorDialect>();
  registry.insert<mlir::tosa::TosaDialect>();
  registry.insert<mlir::bufferization::BufferizationDialect>();
  registry.insert<pim::PimDialect>();
  registry.insert<spatial::SpatialDialect>();
  mlir::tensor::registerBufferizableOpInterfaceExternalModels(registry);
  mlir::tensor::registerInferTypeOpInterfaceExternalModels(registry);
  mlir::arith::registerBufferizableOpInterfaceExternalModels(registry);
  mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
  spatial::registerBufferizableOpInterfaceExternalModels(registry);
  spatial::registerONNXBufferizableOpInterfaceExternalModels(registry);
  pim::registerOpBufferizationInterfaces(registry);
}

void PimAccelerator::registerPasses(int optLevel) const {
  LLVM_DEBUG(llvm::dbgs() << "Registering passes for PIM accelerator\n");
  registerPass(createONNXToSpatialPass);
  registerPass(createSpatialToGraphvizPass);
  registerPass(createSpatialToPimPass);
  registerPass(createBufferizePimPass);
  registerPass(createPimConstantFoldingPass);
  registerPass(createPimHostVerificationPass);
  registerPass(createEmitPimJsonPass);
}

void PimAccelerator::configurePasses() const {
  LLVM_DEBUG(llvm::dbgs() << "Configuring passes for PIM accelerator\n");
  // TODO: This does nothing for now.
}

mlir::MemRefType PimAccelerator::convertTensorTypeToMemRefType(const mlir::TensorType tensorType) const {
  // Do not convert tensor types to memref types.
  return nullptr;
}

void PimAccelerator::conversionTargetONNXToKrnl(mlir::ConversionTarget& target) const {
  target.addLegalDialect<pim::PimDialect>();
}

void PimAccelerator::rewritePatternONNXToKrnl(mlir::RewritePatternSet& patterns,
                                              mlir::TypeConverter& typeConverter,
                                              mlir::MLIRContext* ctx) const {
  // TODO: Add patterns for conversion
}

void PimAccelerator::conversionTargetKrnlToLLVM(mlir::ConversionTarget& target) const {}

void PimAccelerator::rewritePatternKrnlToLLVM(mlir::RewritePatternSet& patterns,
                                              mlir::LLVMTypeConverter& typeConverter,
                                              mlir::MLIRContext* ctx) const {
  // We should not need this, since we offload it all to PIM.
}

} // namespace accel
} // namespace onnx_mlir