#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.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::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
  if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
  if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
  if (isa<pim::PimWaitOp>(op)) return CompiledCoreOpKind::Wait;
  if (isa<pim::PimSyncOp>(op)) return CompiledCoreOpKind::Sync;
  if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
  if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
  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 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;
    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