#include "TestCheck.hpp"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
#include <limits>

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

static void checkProjection(ArrayRef<int64_t> shape, ArrayRef<int64_t> sizes,
    ArrayRef<int64_t> strides, bool available, bool valid = true) {
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, tensor::TensorDialect, func::FuncDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto sourceType = RankedTensorType::get(shape, builder.getF32Type());
  auto resultType = RankedTensorType::get(sizes, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "projection",
    builder.getFunctionType({sourceType}, {resultType}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  auto compute = createEmptySpatGraphCompute(builder, loc, TypeRange{resultType},
    ValueRange{}, ValueRange{function.getArgument(0)});
  auto attrs = [&](ArrayRef<int64_t> values) {
    SmallVector<OpFoldResult> result;
    for (int64_t value : values) result.push_back(builder.getIndexAttr(value));
    return result;
  };
  auto slice = tensor::ExtractSliceOp::create(builder, loc, resultType,
    *compute.getInputArgument(0), attrs(SmallVector<int64_t>(shape.size(), 0)),
    attrs(sizes), attrs(strides));
  Value result = SpatReluOp::create(builder, loc, resultType, slice);
  SpatYieldOp::create(builder, loc, ValueRange{result});
  builder.setInsertionPointAfter(compute);
  func::ReturnOp::create(builder, loc, compute.getResults());
  std::string diagnostics;
  ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
    diagnostics += diagnostic.str();
    return success();
  });
  bool verified = succeeded(verify(*module));
  if (verified != valid) std::fprintf(stderr, "%s\n", diagnostics.c_str());
  TEST_CHECK(verified == valid);
  if (!valid) return;
  diagnostics.clear();
  for (Cost width : {1, 16}) {
    diagnostics.clear();
    SchedulingTarget target;
    target.processorCount = 1;
    target.sharedMemoryWidthBytes = width;
    target.sharedMemoryReadLatencyNs = width == 1 ? 7 : 3;
    auto graph = buildComputeGraph(function, target);
    TEST_CHECK(succeeded(graph));
    if (available) {
      TEST_CHECK(graph->baselineMemory.readRequests == 2);
      TEST_CHECK(graph->baselineMemory.readBytes == 4);
      TEST_CHECK(graph->baselineMemory.serviceNs == (width == 1 ? 28 : 6));
      TEST_CHECK(diagnostics.empty());
    } else {
      TEST_CHECK(graph->baselineMemory.readRequests == static_cast<Cost>(resultType.getNumElements()));
      TEST_CHECK(graph->baselineMemory.readBytes == static_cast<Cost>(resultType.getNumElements()));
      TEST_CHECK(graph->baselineMemory.serviceNs == resultType.getNumElements()
        * target.sharedMemoryReadLatencyNs);
      auto message = diagnostics.find("projection memory cost unavailable");
      TEST_CHECK(message != std::string::npos);
      TEST_CHECK(diagnostics.find("projection memory cost unavailable", message + 1) == std::string::npos);
    }
    // The optional estimate must not make otherwise verified Spatial IR fail
    // schedule preparation. Target lowering retains its own legality checks.
    target.residentWeightCapacity = 1;
    target.interProcessorLatencyNs = {0};
    function.setName("main_graph");
    auto schedules = prepareSpatialSchedules(*module, target, 1, false);
    TEST_CHECK(succeeded(schedules) && schedules->candidates.size() == 1);
    IRMapping mapping;
    OwningOpRef<ModuleOp> scheduled = cast<ModuleOp>(module->getOperation()->clone(mapping));
    auto schedule = remapSpatialSchedule(schedules->candidates.front(), mapping);
    auto realized = realizeSpatialSchedule(*scheduled, schedule, target, 1,
      SpatialDataflowExportStage::None);
    if (failed(realized)) std::fprintf(stderr, "%s\n", diagnostics.c_str());
    TEST_CHECK(succeeded(realized));
    TEST_CHECK(succeeded(verify(*scheduled)));
  }
}

int main() {
  int64_t limit = std::numeric_limits<int64_t>::max();
  checkProjection({3, 5}, {2, 2}, {1, 1}, true);
  // These tensor types and slices are valid Spatial IR. The cost query cannot
  // represent their flattened offsets or strides with signed 64-bit values.
  checkProjection({2, limit}, {2, 2}, {1, 1}, false);
  checkProjection({1, limit, 2}, {1, 1, 1}, {1, 1, 1}, false);
  checkProjection({2, limit}, {1, 1}, {2, 1}, false);
  checkProjection({3, 5}, {4, 2}, {1, 1}, false, false);
}