#include "TestCheck.hpp"
#include <cmath>
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;

static void testZeroContent(OpBuilder &builder) {
  auto type = RankedTensorType::get({4}, builder.getF32Type());
  for (double value : {0.0, -0.0, 1.0, -1.0}) {
    auto attribute = DenseElementsAttr::get(type, builder.getF32FloatAttr(value));
    TEST_CHECK(isZeroBitPatternSplat(attribute) == (value == 0.0 && !std::signbit(value)));
  }
  auto integerType = RankedTensorType::get({4}, builder.getI32Type());
  TEST_CHECK(isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(0))));
  TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(1))));
  TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(type, ArrayRef<float>{0, 1, 0, 0})));
  TEST_CHECK(!isZeroBitPatternSplat(builder.getF32FloatAttr(0)));
  TEST_CHECK(!isZeroBitPatternSplat({}));
}

static void testInputPlan(MLIRContext &context, bool batched) {
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({type}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  Value constant = arith::ConstantOp::create(builder, loc,
    DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
  Value view = tensor::CastOp::create(builder, loc, type, constant);
  SmallVector<Value> inputs{constant, function.getArgument(0), constant, view, function.getArgument(0)};
  Operation *owner;
  if (batched) {
    auto batch = SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
      builder.getI32IntegerAttr(2), ValueRange{}, inputs);
    batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
    owner = batch;
  } else {
    auto single = SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, inputs);
    single->setAttr("coreId", builder.getI32IntegerAttr(0));
    owner = single;
  }
  auto *body = new Block;
  owner->getRegion(0).push_back(body);
  if (batched) body->addArgument(builder.getIndexType(), loc);
  for (Value input : inputs) body->addArgument(input.getType(), loc);
  builder.setInsertionPointToStart(body);
  unsigned first = batched ? 1 : 0;
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value limit = arith::ConstantIndexOp::create(builder, loc, 1000000000000);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  auto loop = scf::ForOp::create(builder, loc, zero, limit, one);
  builder.setInsertionPointToStart(loop.getBody());
  SpatVAddOp::create(builder, loc, type, body->getArgument(first), body->getArgument(first + 1));
  SpatVAddOp::create(builder, loc, type, body->getArgument(first + 3), body->getArgument(first));
  builder.setInsertionPointAfter(loop);
  auto lastUse = SpatVAddOp::create(builder, loc, type, body->getArgument(first + 2), body->getArgument(first));
  auto lastFunctionUse = SpatVAddOp::create(builder, loc, type,
    body->getArgument(first + 4), body->getArgument(first));
  SpatYieldOp::create(builder, loc, ValueRange{});
  builder.setInsertionPointAfter(owner);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(*module)));
  // Body lowering can remove an input's final use before input materialization.
  lastUse.erase();
  lastFunctionUse.erase();

  std::string before;
  llvm::raw_string_ostream stream(before);
  module->print(stream);
  for (unsigned repeat = 0; repeat < 2; ++repeat) {
    auto constantPlan = planComputeInputMaterialization(owner, 0);
    auto runtimePlan = planComputeInputMaterialization(owner, 1);
    auto unusedPlan = planComputeInputMaterialization(owner, 2);
    auto viewPlan = planComputeInputMaterialization(owner, 3);
    TEST_CHECK(succeeded(constantPlan) && succeeded(runtimePlan) && succeeded(unusedPlan));
    TEST_CHECK(*constantPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
                                        : ComputeInputMaterialization::ForwardConstant));
    TEST_CHECK(*runtimePlan == ComputeInputMaterialization::ForwardFunctionInput);
    auto unusedFunctionPlan = planComputeInputMaterialization(owner, 4);
    TEST_CHECK(succeeded(unusedFunctionPlan)
      && *unusedFunctionPlan == (batched ? ComputeInputMaterialization::ForwardFunctionInput
                                       : ComputeInputMaterialization::Unused));
    TEST_CHECK(*unusedPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
                                      : ComputeInputMaterialization::Unused));
    TEST_CHECK(succeeded(viewPlan) && *viewPlan == ComputeInputMaterialization::CopyAtEntry);
  }
  TEST_CHECK(failed(planComputeInputMaterialization(owner, 5)));
  TEST_CHECK(failed(planComputeInputMaterialization(function, 0)));
  std::string after;
  llvm::raw_string_ostream afterStream(after);
  module->print(afterStream);
  TEST_CHECK(before == after);
}

static void testOutputReuse(MLIRContext &context) {
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  IRRewriter builder(&context);
  builder.setInsertionPointToStart(module->getBody());
  auto loc = module->getLoc();
  auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "reuse", builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
  Value constant = arith::ConstantOp::create(builder, loc,
    DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
  auto first = SpatVAddOp::create(builder, loc, type, input, constant);
  TEST_CHECK(getReusableOutputTensor(first) == input);
  auto second = SpatVAddOp::create(builder, loc, type, input, constant);
  TEST_CHECK(!getReusableOutputTensor(first));
  TEST_CHECK(getReusableOutputTensor(second) == input);
  second->moveBefore(first);
  TEST_CHECK(getReusableOutputTensor(first) == input);
  TEST_CHECK(!getReusableOutputTensor(second));
  second->moveAfter(first);
  TEST_CHECK(!getReusableOutputTensor(first));
  TEST_CHECK(getReusableOutputTensor(second) == input);
  Value view = tensor::CastOp::create(builder, loc, type, input);
  auto fromView = SpatVAddOp::create(builder, loc, type, view, constant);
  TEST_CHECK(!getReusableOutputTensor(fromView));
  TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, fromView)
    .getDefiningOp<tensor::EmptyOp>());
  builder.setInsertionPointAfter(fromView);
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  Value limit = arith::ConstantIndexOp::create(builder, loc, 10);
  Value one = arith::ConstantIndexOp::create(builder, loc, 1);
  auto loop = scf::ForOp::create(builder, loc, zero, limit, one, ValueRange{input});
  builder.setInsertionPointToStart(loop.getBody());
  auto carried = SpatVAddOp::create(builder, loc, type, loop.getRegionIterArg(0), constant);
  scf::YieldOp::create(builder, loc, ValueRange{carried});
  TEST_CHECK(getReusableOutputTensor(carried) == loop.getRegionIterArg(0));
  TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, carried) == loop.getRegionIterArg(0));
  builder.setInsertionPointAfter(loop);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(*module)));
}

int main() {
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
    tensor::TensorDialect, SpatialDialect>();
  OpBuilder builder(&context);
  testZeroContent(builder);
  testInputPlan(context, false);
  testInputPlan(context, true);
  testOutputReuse(context);
}