#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "OpBufferizationInterfaces.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace mlir;
using namespace bufferization;

namespace onnx_mlir {
namespace pim {

static Value materializeContiguousMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
  if (succeeded(resolveContiguousAddress(memrefValue)))
    return memrefValue;

  auto shapedType = cast<ShapedType>(memrefValue.getType());
  auto contiguousType = MemRefType::get(shapedType.getShape(), shapedType.getElementType());
  Value contiguousBuffer = memref::AllocOp::create(rewriter, loc, contiguousType);
  auto sizeInBytes = shapedType.getNumElements() * shapedType.getElementTypeBitWidth() / 8;

  return PimMemCopyOp::create(rewriter,
                              loc,
                              contiguousType,
                              contiguousBuffer,
                              memrefValue,
                              rewriter.getI32IntegerAttr(0),
                              rewriter.getI32IntegerAttr(0),
                              rewriter.getI32IntegerAttr(sizeInBytes))
    .getOutput();
}

struct MemCopyHostToDevOpInterface
: DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto memCopyHostToDevOp = cast<PimMemCopyHostToDevOp>(op);
    auto deviceTarget = memCopyHostToDevOp.getDeviceTarget();
    auto hostSource = memCopyHostToDevOp.getHostSource();

    auto deviceTargetOpt = getBuffer(rewriter, deviceTarget, options, state);
    if (failed(deviceTargetOpt))
      return failure();
    auto deviceTargetMemRef = *deviceTargetOpt;

    auto hostSourceOpt = getBuffer(rewriter, hostSource, options, state);
    if (failed(hostSourceOpt))
      return failure();
    auto hostSourceMemRef = *hostSourceOpt;

    replaceOpWithNewBufferizedOp<PimMemCopyHostToDevOp>(rewriter,
                                                        memCopyHostToDevOp,
                                                        deviceTargetMemRef.getType(),
                                                        deviceTargetMemRef,
                                                        hostSourceMemRef,
                                                        memCopyHostToDevOp.getDeviceTargetOffsetAttr(),
                                                        memCopyHostToDevOp.getHostSourceOffsetAttr(),
                                                        memCopyHostToDevOp.getSizeAttr());
    return success();
  }
};

struct MemCopyDevToHostOpInterface
: DstBufferizableOpInterfaceExternalModel<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto memCopyDevToHostOp = cast<PimMemCopyDevToHostOp>(op);

    auto hostTarget = memCopyDevToHostOp.getHostTarget();
    auto hostTargetOpt = getBuffer(rewriter, hostTarget, options, state);
    if (failed(hostTargetOpt))
      return failure();
    auto hostTargetMemRef = *hostTargetOpt;

    auto deviceSource = memCopyDevToHostOp.getDeviceSource();
    auto deviceSourceOpt = getBuffer(rewriter, deviceSource, options, state);
    if (failed(deviceSourceOpt))
      return failure();
    auto deviceSourceMemRef = *deviceSourceOpt;

    replaceOpWithNewBufferizedOp<PimMemCopyDevToHostOp>(rewriter,
                                                        memCopyDevToHostOp,
                                                        hostTargetMemRef.getType(),
                                                        hostTargetMemRef,
                                                        deviceSourceMemRef,
                                                        memCopyDevToHostOp.getHostTargetOffsetAttr(),
                                                        memCopyDevToHostOp.getDeviceSourceOffsetAttr(),
                                                        memCopyDevToHostOp.getSizeAttr());
    return success();
  }
};

struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeOpInterface, PimTransposeOp> {
  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
  }

  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto transposeOp = cast<PimTransposeOp>(op);

    auto inputOpt = getBuffer(rewriter, transposeOp.getInput(), options, state);
    if (failed(inputOpt))
      return failure();

    auto outputBufferOpt = getBuffer(rewriter, transposeOp.getOutputBuffer(), options, state);
    if (failed(outputBufferOpt))
      return failure();

    replaceOpWithNewBufferizedOp<PimTransposeOp>(
      rewriter, op, outputBufferOpt->getType(), *inputOpt, transposeOp.getPermutation(), *outputBufferOpt);
    return success();
  }
};

struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface, PimVMMOp> {
  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
  }

  bool isNotConflicting(Operation* op, OpOperand* uRead, OpOperand* uWrite, const AnalysisState& state) const {
    auto vmmOp = cast<PimVMMOp>(op);
    Value readVal = uRead->get();
    Value writeVal = uWrite->get();
    if (writeVal != vmmOp.getOutputBuffer())
      return false;
    if (readVal == vmmOp.getInput())
      if (state.areEquivalentBufferizedValues(readVal, writeVal))
        return true;
    return false;
  }

  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto vmmOp = cast<PimVMMOp>(op);

    auto inputOpt = getBuffer(rewriter, vmmOp.getInput(), options, state);
    if (failed(inputOpt))
      return failure();

    auto outputBufferOpt = getBuffer(rewriter, vmmOp.getOutputBuffer(), options, state);
    if (failed(outputBufferOpt))
      return failure();

    replaceOpWithNewBufferizedOp<PimVMMOp>(
      rewriter, op, outputBufferOpt->getType(), vmmOp.getWeightIndexAttr(), *inputOpt, *outputBufferOpt);
    return success();
  }
};

struct MVMOpInterface : DstBufferizableOpInterfaceExternalModel<MVMOpInterface, PimMVMOp> {
  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
  }

  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto mvmOp = cast<PimMVMOp>(op);

    auto inputOpt = getBuffer(rewriter, mvmOp.getInput(), options, state);
    if (failed(inputOpt))
      return failure();

    auto outputBufferOpt = getBuffer(rewriter, mvmOp.getOutputBuffer(), options, state);
    if (failed(outputBufferOpt))
      return failure();

    replaceOpWithNewBufferizedOp<PimMVMOp>(
      rewriter, op, outputBufferOpt->getType(), mvmOp.getWeightIndexAttr(), *inputOpt, *outputBufferOpt);
    return success();
  }
};

template <typename OpTy>
struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstOpInterface<OpTy>, OpTy> {
  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
  }

  bool bufferizesToElementwiseAccess(Operation* tablegen_opaque_val,
                                     const AnalysisState& state,
                                     ArrayRef<OpOperand*> opOperands) const {
    return true;
  }

  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto binaryOp = cast<OpTy>(op);

    auto lhsOpt = getBuffer(rewriter, binaryOp.getLhs(), options, state);
    if (failed(lhsOpt))
      return failure();

    auto rhsOpt = getBuffer(rewriter, binaryOp.getRhs(), options, state);
    if (failed(rhsOpt))
      return failure();

    auto outputBufferOpt = getBuffer(rewriter, binaryOp.getOutputBuffer(), options, state);
    if (failed(outputBufferOpt))
      return failure();

    Value contiguousLhs = materializeContiguousMemRef(*lhsOpt, op->getLoc(), rewriter);
    Value contiguousRhs = materializeContiguousMemRef(*rhsOpt, op->getLoc(), rewriter);

    replaceOpWithNewBufferizedOp<OpTy>(
      rewriter, op, outputBufferOpt->getType(), contiguousLhs, contiguousRhs, *outputBufferOpt);
    return success();
  }
};

template <typename OpTy>
struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpInterface<OpTy>, OpTy> {
  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
  }

  bool bufferizesToElementwiseAccess(Operation* tablegen_opaque_val,
                                     const AnalysisState& state,
                                     ArrayRef<OpOperand*> opOperands) const {
    return true;
  }

  LogicalResult bufferize(Operation* op,
                          RewriterBase& rewriter,
                          const BufferizationOptions& options,
                          BufferizationState& state) const {
    auto unaryOp = cast<OpTy>(op);

    auto inputOpt = getBuffer(rewriter, unaryOp.getInput(), options, state);
    if (failed(inputOpt))
      return failure();

    auto outputBufferOpt = getBuffer(rewriter, unaryOp.getOutputBuffer(), options, state);
    if (failed(outputBufferOpt))
      return failure();

    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);

    replaceOpWithNewBufferizedOp<OpTy>(rewriter, op, outputBufferOpt->getType(), contiguousInput, *outputBufferOpt);
    return success();
  }
};

void registerOpBufferizationInterfaces(DialectRegistry& registry) {
  registry.addExtension(+[](MLIRContext* ctx, PimDialect* dialect) {
    PimMemCopyHostToDevOp::attachInterface<MemCopyHostToDevOpInterface>(*ctx);
    PimMemCopyDevToHostOp::attachInterface<MemCopyDevToHostOpInterface>(*ctx);
    PimTransposeOp::attachInterface<TransposeOpInterface>(*ctx);
    PimVMMOp::attachInterface<VMMOpInterface>(*ctx);
    PimMVMOp::attachInterface<MVMOpInterface>(*ctx);

    PimVVAddOp::attachInterface<BinaryDstOpInterface<PimVVAddOp>>(*ctx);
    PimVVSubOp::attachInterface<BinaryDstOpInterface<PimVVSubOp>>(*ctx);
    PimVVMulOp::attachInterface<BinaryDstOpInterface<PimVVMulOp>>(*ctx);
    PimVVMaxOp::attachInterface<BinaryDstOpInterface<PimVVMaxOp>>(*ctx);

    PimVAvgOp::attachInterface<UnaryDstOpInterface<PimVAvgOp>>(*ctx);
    PimVReluOp::attachInterface<UnaryDstOpInterface<PimVReluOp>>(*ctx);
    PimVTanhOp::attachInterface<UnaryDstOpInterface<PimVTanhOp>>(*ctx);
    PimVSigmOp::attachInterface<UnaryDstOpInterface<PimVSigmOp>>(*ctx);
  });
}

} // namespace pim
} // namespace onnx_mlir
