#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
#include "llvm/Support/JSON.h"
#include "TestCheck.hpp"
#include <limits>
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
#include "src/Accelerators/PIM/Compiler/PimPipelineSelection.hpp"
#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.hpp"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/IR/Verifier.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
static WholeScheduleCost scheduleCost(double interval) {
  WholeScheduleCost cost;
  cost.primary = cost.alternative = interval;
  cost.selectionMargin = 0.021;
  cost.loweringSupport = StaticLoweringSupport::Supported;
  for (auto &recurrence : cost.recurrence)
    recurrence = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
  return cost;
}

static void testEmptyGraphPreparation() {
  using namespace mlir;
  MLIRContext context;
  context.loadDialect<func::FuncDialect, SpatialDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto function = func::FuncOp::create(builder, module->getLoc(), "main_graph",
    builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  func::ReturnOp::create(builder, module->getLoc());
  SchedulingTarget target;
  target.processorCount = target.residentWeightCapacity = 1;
  target.interProcessorLatencyNs = {0};
  for (bool collectCandidates : {false, true}) {
    auto prepared = prepareSpatialSchedules(*module, target, 1, collectCandidates);
    TEST_CHECK(succeeded(prepared) && prepared->candidates.size() == 1);
    TEST_CHECK(prepared->candidates.front().dominanceOrderCompute.empty());
  }
}

static void testSelectedLoweringFailure() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, pim::PimDialect, func::FuncDialect,
    arith::ArithDialect, tensor::TensorDialect, memref::MemRefDialect,
    scf::SCFDialect, linalg::LinalgDialect, affine::AffineDialect,
    bufferization::BufferizationDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto type = RankedTensorType::get({1, 536870912}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({type}, {type}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto compute = SpatGraphCompute::create(builder, loc, TypeRange{type}, ValueRange{},
    ValueRange{function.getArgument(0)});
  compute.getBody().push_back(new Block);
  Value input = compute.getBody().front().addArgument(type, loc);
  builder.setInsertionPointToStart(&compute.getBody().front());
  Value result = SpatReluOp::create(builder, loc, type, input);
  SpatYieldOp::create(builder, loc, ValueRange{result});
  builder.setInsertionPointAfter(compute);
  func::ReturnOp::create(builder, loc, compute.getResults());
  TEST_CHECK(succeeded(verify(*module)));
  SchedulingTarget target;
  target.processorCount = 2;
  target.residentWeightCapacity = 1;
  target.synchronizationRegisterCount = 64;
  target.interProcessorLatencyNs = {0, 2, 2, 0};
  target.averageInterProcessorLatencyNs = 2;
  target.estimateRealizedScheduleCost = [](ModuleOp, llvm::json::Object &) {
    TEST_CHECK(false && "single-candidate compilation must skip costs");
    return WholeScheduleCost{};
  };
  SpatialTargetResources resources{{128, 128}, 1, 2, 16};
  std::string diagnostics;
  ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
    diagnostics += diagnostic.str(); return success();
  });
  ScopedCompilerOutputSuppression silence;
  // Match the compiler's preloaded backend registry before nested pipelines run.
  PassManager backend(&context);
  addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
  DialectRegistry registry;
  backend.getDependentDialects(registry);
  pim::registerOpBufferizationInterfaces(registry);
  arith::registerBufferizableOpInterfaceExternalModels(registry);
  bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
  context.appendDialectRegistry(registry);
  context.loadAllAvailableDialects();
  for (bool collectCandidates : {false, true}) {
    OwningOpRef<ModuleOp> checkpoint = cast<ModuleOp>((*module)->clone());
    auto prepared = prepareSpatialSchedules(*checkpoint, target, 2, collectCandidates);
    TEST_CHECK(succeeded(prepared) && !prepared->candidates.empty());
  }
  PassManager manager(&context);
  manager.addPass(createPimPipelineSelectionPass(target, resources, 2, SpatialDataflowExportStage::None));
  TEST_CHECK(failed(manager.run(*module)));
  TEST_CHECK(diagnostics.find("backend selection failed") != std::string::npos);
  TEST_CHECK(diagnostics.find("outside representable range") != std::string::npos);
}

static void testResourceRetries() {
  using namespace mlir;
  auto previousCores = coresCount.getValue();
  auto previousSize = crossbarSize.getValue();
  auto previousCrossbars = crossbarCountInCore.getValue();
  coresCount = 2; crossbarSize = 128; crossbarCountInCore = 1;
  MLIRContext context;
  context.loadDialect<pim::PimDialect, memref::MemRefDialect, func::FuncDialect,
    arith::ArithDialect>();
  OwningOpRef<ModuleOp> checkpoint = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(checkpoint->getBodyRegion());
  auto loc = checkpoint->getLoc();
  auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({}, {}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto core = pim::PimCoreOp::create(builder, loc, ValueRange{}, uint32_t{0});
  core.getBody().push_back(new Block);
  builder.setInsertionPointToStart(&core.getBody().front());
  auto allocation = memref::AllocOp::create(builder, loc,
    MemRefType::get({int64_t{1} << 31}, builder.getI8Type()));
  auto smallType = MemRefType::get({1}, builder.getI8Type());
  auto destination = memref::AllocOp::create(builder, loc, smallType);
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  pim::PimMemCopyOp::create(builder, loc, smallType, zero, zero, destination,
    allocation, builder.getI32IntegerAttr(1));
  pim::PimHaltOp::create(builder, loc);
  builder.setInsertionPointAfter(core);
  func::ReturnOp::create(builder, loc);
  TEST_CHECK(succeeded(verify(*checkpoint)));
  std::string before;
  llvm::raw_string_ostream stream(before); checkpoint->print(stream, OpPrintingFlags().printGenericOpForm());
  auto original = scheduleCost(100);
  auto faster = scheduleCost(90);
  std::vector<WholeScheduleCost> costs{original, faster};
  std::string reason;
  SpatialTargetResources resources{{128, 128}, 1, 2, 16};
  PassManager backend(&context);
  addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
  DialectRegistry registry;
  backend.getDependentDialects(registry);
  pim::registerOpBufferizationInterfaces(registry);
  arith::registerBufferizableOpInterfaceExternalModels(registry);
  bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
  context.appendDialectRegistry(registry);
  context.loadAllAvailableDialects();
  auto emit = [&](ModuleOp module) {
    llvm::SmallString<128> temporary;
    TEST_CHECK(!llvm::sys::fs::createUniqueDirectory("raptor-retry-program", temporary));
    std::string output = temporary.str().str();
    ScopedCompilerOutputSuppression silence;
    TEST_CHECK(compileToPimCode(module, output) == CompilerSuccess);
    auto binary = llvm::MemoryBuffer::getFile(output + "/core_0.pim");
    TEST_CHECK(binary && (*binary)->getBufferSize() > 12);
    std::string program = (*binary)->getBuffer().str();
    TEST_CHECK(!llvm::sys::fs::remove_directories(temporary));
    return program;
  };
  std::string emitted;
  for (bool allFail : {false, true}) {
    std::vector<size_t> attempted;
    OwningOpRef<ModuleOp> winner;
    std::vector<llvm::json::Object> summaries;
    auto selected = lowerSelectedPimCandidate(costs, [&](size_t index) {
      attempted.push_back(index);
      IRMapping mapping;
      OwningOpRef<ModuleOp> candidate = cast<ModuleOp>((*checkpoint)->clone(mapping));
      if (!index && !allFail)
        cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
          .getResult().setType(MemRefType::get({16}, builder.getI8Type()));
      llvm::json::Object summary;
      auto result = lowerPimPipelineCandidate(*candidate, resources, summary);
      if (result == PimCandidateResult::Success) winner = std::move(candidate);
      summaries.push_back(std::move(summary));
      return result;
    }, reason);
    TEST_CHECK(attempted == std::vector<size_t>({1, 0}));
    TEST_CHECK(allFail ? !selected : selected == 0);
    TEST_CHECK(summaries.front().getString("status") == "resource-limit");
    TEST_CHECK(summaries.front().getString("reason")->contains("signed int32 address range"));
    if (selected) emitted = emit(*winner);
    if (allFail) TEST_CHECK(reason.find("all candidate schedules failed") != std::string::npos);
    std::string after;
    llvm::raw_string_ostream out(after); checkpoint->print(out, OpPrintingFlags().printGenericOpForm());
    TEST_CHECK(before == after);
  }
  TEST_CHECK(!emitted.empty());
  IRMapping mapping;
  OwningOpRef<ModuleOp> alternative = cast<ModuleOp>((*checkpoint)->clone(mapping));
  cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
    .getResult().setType(MemRefType::get({16}, builder.getI8Type()));
  llvm::json::Object summary;
  TEST_CHECK(lowerPimPipelineCandidate(*alternative, resources, summary) == PimCandidateResult::Success);
  TEST_CHECK(emit(*alternative) == emitted);
  size_t attempts = 0;
  auto empty = lowerSelectedPimCandidate({}, [&](size_t) {
    ++attempts;
    return PimCandidateResult::Success;
  }, reason);
  TEST_CHECK(!empty && attempts == 0 && reason.find("at least one candidate") != std::string::npos);
  auto selected = lowerSelectedPimCandidate(costs, [&](size_t) {
    ++attempts;
    OwningOpRef<ModuleOp> invalid = cast<ModuleOp>((*alternative)->clone());
    invalid->walk([&](pim::PimMemCopyOp copy) {
      copy->setAttr("size", builder.getI64IntegerAttr(1));
    });
    llvm::json::Object failure;
    auto result = lowerPimPipelineCandidate(*invalid, resources, failure);
    TEST_CHECK(result == PimCandidateResult::Failure);
    TEST_CHECK(failure.getString("status") == "systemic-or-unclassified-error");
    return result;
  }, reason);
  TEST_CHECK(!selected && attempts == 1);
  TEST_CHECK(reason.find("recovery is not justified") != std::string::npos);
  for (unsigned unknown = 0; unknown < 4; ++unknown) {
    auto invalid = faster;
    if (unknown == 0) invalid.primary.reset();
    if (unknown == 1) invalid.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
    if (unknown == 2) invalid.loweringSupport = StaticLoweringSupport::Unknown;
    if (unknown == 3) invalid.loweringSupport = StaticLoweringSupport::Unsupported;
    attempts = 0;
    std::vector<WholeScheduleCost> incomplete{original, invalid};
    selected = lowerSelectedPimCandidate(incomplete, [&](size_t) {
      ++attempts;
      return PimCandidateResult::Success;
    }, reason);
    TEST_CHECK(!selected && attempts == 0);
  }
  coresCount = previousCores; crossbarSize = previousSize; crossbarCountInCore = previousCrossbars;
}

int main() {
  testEmptyGraphPreparation();
  testSelectedLoweringFailure();
  testResourceRetries();
  mlir::MLIRContext context;
  context.allowUnregisteredDialects();
  mlir::OwningOpRef<mlir::ModuleOp> checkpoint = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OperationState state(checkpoint->getLoc(), "test.compute");
  auto *op = mlir::Operation::create(state);
  checkpoint->getBody()->push_back(op);
  spatial::ComputeInstance instance{op, 3, 2};
  spatial::MergeScheduleResult schedule;
  schedule.dominanceOrderCompute.push_back(instance);
  schedule.computeToCpuMap[instance] = 7;
  schedule.computeToCpuSlotMap[instance] = 4;
  schedule.computeToAestMap[instance] = 42;
  schedule.isLastComputeOfCpu.insert(instance);
  schedule.cpuToLastComputeMap[7] = instance;
  for (int i = 0; i < 3; ++i) {
    mlir::IRMapping mapping;
    mlir::OwningOpRef<mlir::ModuleOp> trial = mlir::cast<mlir::ModuleOp>((*checkpoint)->clone(mapping));
    auto remapped = spatial::remapSpatialSchedule(schedule, mapping);
    auto copy = remapped.dominanceOrderCompute.front();
    TEST_CHECK(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
    TEST_CHECK(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
    TEST_CHECK(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
    TEST_CHECK(remapped.cpuToLastComputeMap.lookup(7) == copy);
    copy.op->erase();
    TEST_CHECK(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
  }

  auto score = [](double value, std::string) {
    return scheduleCost(value);
  };
  std::string reason;
  auto original = score(100, "0"), faster = score(90, "1");
  for (auto support : {StaticLoweringSupport::Unknown,
      StaticLoweringSupport::Supported, StaticLoweringSupport::Unsupported}) {
    auto cost = faster;
    cost.loweringSupport = support;
    TEST_CHECK(cost.available());
    TEST_CHECK(cost.eligible() == (support == StaticLoweringSupport::Supported));
    TEST_CHECK(cost.loweringSupported().has_value() == (support != StaticLoweringSupport::Unknown));
  }
  auto incompleteProof = faster;
  incompleteProof.recurrence[1].periodGenerations = 0;
  TEST_CHECK(incompleteProof.available() && !incompleteProof.eligible());
  incompleteProof.recurrence[1] = {};
  TEST_CHECK(incompleteProof.convergenceProved() == false && !incompleteProof.eligible());
  incompleteProof.recurrence[0] = {};
  TEST_CHECK(!incompleteProof.convergenceProved().has_value());
  auto select = [&](std::initializer_list<WholeScheduleCost> values) {
    return selectPipelineCandidate(llvm::ArrayRef<WholeScheduleCost>(values), reason);
  };
  TEST_CHECK(select({original, faster}) == 1);
  TEST_CHECK(select({original, score(97, "2")}) == 0);
  TEST_CHECK(select({original, score(105, "2")}) == 0);
  auto sensitive = faster; sensitive.alternative = 101;
  TEST_CHECK(select({original, sensitive}) == 0);
  TEST_CHECK(select({original, faster, score(91, "2")}) == 1);
  TEST_CHECK(select({original, faster, score(90, "1")}) == 1);
  auto unavailable = WholeScheduleCost{};
  TEST_CHECK(select({unavailable, faster}) == 0);
  TEST_CHECK(select({original, unavailable}) == 0);
  auto unknownUncertainty = faster;
  unknownUncertainty.selectionMargin.reset();
  TEST_CHECK(select({original, unknownUncertainty}) == 0);
  for (double value : {0., -1., std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
    TEST_CHECK(select({original, score(value, "2")}) == 0);
  for (int i = 0; i < 8; ++i) TEST_CHECK(select({original, faster}) == 1);

}