#pragma once

#include "mlir/IR/Operation.h"
#include "mlir/Support/LogicalResult.h"

#include "llvm/ADT/SmallVector.h"

#include <memory>
#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"

namespace onnx_mlir {

enum class CompiledCoreOpKind : uint8_t {
  Load,
  Store,
  Lmv,
  VMV,
  Receive,
  Send,
  Wait,
  Sync,
  Concat,
  Vmm,
  VVAdd,
  VVSub,
  VVMul,
  VVMax,
  VVDMul,
  VAvg,
  VRelu,
  VTanh,
  VSigm,
  VSoftmax
};

struct CompiledCoreNode {
  enum class Kind : uint8_t { Op, Loop, If, IndexSwitch };

  Kind kind = Kind::Op;
  mlir::Operation* op = nullptr;
  CompiledCoreOpKind opKind = CompiledCoreOpKind::Load;
  CompiledIndexExpr lowerBound;
  CompiledIndexExpr upperBound;
  CompiledIndexExpr step;
  CompiledIndexExpr condition;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
  llvm::SmallVector<int64_t> caseValues;
  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
};

struct CompiledCoreProgram {
  llvm::SmallVector<CompiledCoreNode, 32> nodes;
};

mlir::LogicalResult compileCoreProgram(mlir::Operation* coreLikeOp, CompiledCoreProgram& program);

} // namespace onnx_mlir