#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/STLExtras.h"

#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace llvm;
using namespace mlir;

namespace onnx_mlir {
namespace {

static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
  if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
  if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
  if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
  if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
  if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
  if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
  if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
  if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
  if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
  if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
  if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
  if (isa<pim::PimVVMaxOp>(op)) return CompiledCoreOpKind::VVMax;
  if (isa<pim::PimVVDMulOp>(op)) return CompiledCoreOpKind::VVDMul;
  if (isa<pim::PimVAvgOp>(op)) return CompiledCoreOpKind::VAvg;
  if (isa<pim::PimVReluOp>(op)) return CompiledCoreOpKind::VRelu;
  if (isa<pim::PimVTanhOp>(op)) return CompiledCoreOpKind::VTanh;
  if (isa<pim::PimVSigmOp>(op)) return CompiledCoreOpKind::VSigm;
  if (isa<pim::PimVSoftmaxOp>(op)) return CompiledCoreOpKind::VSoftmax;
  return failure();
}

static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
  SmallVector<unsigned> sourceNonUnitDims;
  SmallVector<unsigned> destinationSourceNonUnitDims;
  for (auto [dim, size] : llvm::enumerate(sourceShape))
    if (size != 1)
      sourceNonUnitDims.push_back(dim);
  for (int64_t sourceDim : permutation)
    if (sourceShape[sourceDim] != 1)
      destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
  return sourceNonUnitDims == destinationSourceNonUnitDims;
}

static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
  auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
  ArrayRef<int64_t> sourceShape = sourceType.getShape();
  size_t rank = sourceShape.size();
  CompiledTransposePlan plan;
  plan.source = transposeOp.getInput();
  plan.destination = transposeOp.getOutputBuffer();
  plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
  auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
  if (failed(totalElements)) return failure();
  plan.totalElements = *totalElements;
  auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
  if (failed(totalBytes)) return failure();
  plan.totalBytes = *totalBytes;

  SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
                                                   [](IntegerAttr attr) { return attr.getInt(); });
  if (permutation.size() != rank) {
    transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
    return failure();
  }

  SmallVector<size_t> destinationShape(rank);
  plan.destinationStrides.assign(rank, 1);
  plan.destinationDimensionForSource.assign(rank, 0);
  plan.destinationRewinds.assign(rank, 0);
  SmallVector<bool> seenSourceDimensions(rank, false);
  for (size_t dim = 0; dim < rank; ++dim) {
    auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
    if (failed(size)) return failure();
    plan.sourceShape.push_back(*size);
  }
  for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
    if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
      transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
      return failure();
    }
    seenSourceDimensions[sourceDim] = true;
    destinationShape[destinationDim] = plan.sourceShape[sourceDim];
    plan.destinationDimensionForSource[sourceDim] = destinationDim;
  }
  for (size_t dim = rank; dim > 1; --dim) {
    auto stride = pim::checkedMul(
      plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
    if (failed(stride)) return failure();
    plan.destinationStrides[dim - 2] = *stride;
  }
  for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
    auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
                                  plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
                                  transposeOp,
                                  "transpose destination rewind");
    if (failed(rewind)) return failure();
    plan.destinationRewinds[sourceDim] = *rewind;
  }
  plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
  return plan;
}

static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
  for (Operation& op : block) {
    if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
      continue;
    if (auto loadOp = dyn_cast<memref::LoadOp>(op); loadOp && succeeded(compileIndexExpr(loadOp.getResult())))
      continue;

    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
      auto lower = compileIndexExpr(forOp.getLowerBound());
      auto upper = compileIndexExpr(forOp.getUpperBound());
      auto step = compileIndexExpr(forOp.getStep());
      if (failed(lower) || failed(upper) || failed(step)) {
        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
        return failure();
      }
      CompiledCoreNode node;
      node.kind = CompiledCoreNode::Kind::Loop;
      node.op = forOp;
      node.lowerBound = *lower;
      node.upperBound = *upper;
      node.step = *step;
      node.loopBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
      if (failed(compileCoreEmissionPlan(forOp.getRegion().front(), *node.loopBody))) return failure();
      plan.push_back(std::move(node));
      continue;
    }
    if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
      auto condition = compileIndexExpr(ifOp.getCondition());
      if (failed(condition)) {
        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
        return failure();
      }
      CompiledCoreNode node;
      node.kind = CompiledCoreNode::Kind::If;
      node.op = ifOp;
      node.condition = *condition;
      node.thenBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
      node.elseBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), *node.thenBody))) return failure();
      if (!ifOp.getElseRegion().empty()
          && failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), *node.elseBody)))
        return failure();
      plan.push_back(std::move(node));
      continue;
    }
    if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(op)) {
      auto selector = compileIndexExpr(switchOp.getArg());
      if (failed(selector)) {
        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
        return failure();
      }
      CompiledCoreNode node;
      node.kind = CompiledCoreNode::Kind::IndexSwitch;
      node.op = switchOp;
      node.condition = *selector;
      llvm::append_range(node.caseValues, switchOp.getCases());
      for (Region& region : switchOp.getCaseRegions()) {
        auto body = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
        if (failed(compileCoreEmissionPlan(region.front(), *body))) return failure();
        node.caseBodies.push_back(std::move(body));
      }
      node.defaultBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
      if (failed(compileCoreEmissionPlan(switchOp.getDefaultRegion().front(), *node.defaultBody))) return failure();
      plan.push_back(std::move(node));
      continue;
    }

    auto opKind = classifyCompiledCoreOpKind(op);
    if (failed(opKind)) {
      InFlightDiagnostic diagnostic = op.emitError() << "unsupported codegen for op '" << op.getName() << "'";
      if (auto coreOp = op.getParentOfType<pim::PimCoreOp>())
        diagnostic << " inside pim.core " << coreOp.getCoreId();
      else if (auto batchOp = op.getParentOfType<pim::PimCoreBatchOp>())
        diagnostic << " inside pim.core_batch with laneCount " << batchOp.getLaneCount();
      return failure();
    }
    CompiledCoreNode node;
    node.op = &op;
    node.opKind = *opKind;
    if (*opKind == CompiledCoreOpKind::Transpose) {
      auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
      if (failed(transposePlan)) return failure();
      node.transposePlan = *transposePlan;
    }
    plan.push_back(std::move(node));
  }
  return success();
}

} // namespace

LogicalResult compileCoreProgram(Operation* coreLikeOp, CompiledCoreProgram& program) {
  Block& block = isa<pim::PimCoreOp>(coreLikeOp) ? cast<pim::PimCoreOp>(coreLikeOp).getBody().front()
                                                 : cast<pim::PimCoreBatchOp>(coreLikeOp).getBody().front();
  return compileCoreEmissionPlan(block, program.nodes);
}

} // namespace onnx_mlir