#include "src/Accelerators/PIM/Common/IR/CopyAccessPlan.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
#include "TestCheck.hpp"
#include <limits>
#include <random>

using namespace onnx_mlir;

int main(int argc, char **argv) {
  if (argc > 1 && std::string(argv[1]) == "--stride-overflow") {
    computeRowMajorStrides({2, std::numeric_limits<int64_t>::max(), 2});
    return 0;
  }
  if (argc > 1) TEST_CHECK(false);
  auto dense = CopyAccessPlan::build({3, 4}, {4, 1}, {4, 1});
  TEST_CHECK(mlir::succeeded(dense) && dense->getRunCount() == 1 && dense->getRunElements() == 12);
  auto strided = CopyAccessPlan::build({3, 4}, {8, 1}, {4, 1});
  TEST_CHECK(mlir::succeeded(strided) && strided->getRunCount() == 3 && strided->getRunElements() == 4);
  TEST_CHECK(*strided->getRunOffsets(2) == std::make_pair(int64_t{16}, int64_t{8}));
  TEST_CHECK(mlir::failed(strided->getRunOffsets(-1)) && mlir::failed(strided->getRunOffsets(3)));
  auto large = CopyAccessPlan::build({1000000000, 16}, {32, 1}, {16, 1});
  TEST_CHECK(mlir::succeeded(large) && large->getOuterShape().size() == 1);
  TEST_CHECK(large->getRunCount() == 1000000000 && large->getRunElements() == 16);
  TEST_CHECK(*large->getRunOffsets(999999999) == std::make_pair(int64_t{31999999968}, int64_t{15999999984}));
  auto scalar = CopyAccessPlan::build({}, {}, {});
  TEST_CHECK(mlir::succeeded(scalar) && scalar->getRunCount() == 1 && scalar->getRunElements() == 1);
  // Observed input planes and channel-wise publication slices, without PIM IR.
  auto input = CopyAccessPlan::build({3, 1, 224}, {50176, 224, 1}, {224, 224, 1});
  TEST_CHECK(mlir::succeeded(input) && input->getRunCount() == 3 && input->getRunElements() == 224);
  TEST_CHECK(*input->getRunOffsets(1) == std::make_pair(int64_t{50176}, int64_t{224}));
  for (int64_t channels : {32, 96}) {
    auto publication = CopyAccessPlan::build({1, channels, 1, 80},
      {channels * 6400, 6400, 80, 1}, {channels * 80, 80, 80, 1});
    TEST_CHECK(mlir::succeeded(publication) && publication->getRunCount() == channels);
    TEST_CHECK(publication->getRunElements() == 80);
    TEST_CHECK(*publication->getRunOffsets(1) == std::make_pair(int64_t{6400}, int64_t{80}));
  }
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({2}, {}, {1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({0}, {1}, {1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({-1}, {1}, {1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({2}, {-1}, {1})));
  int64_t limit = std::numeric_limits<int64_t>::max();
  TEST_CHECK(mlir::failed(computeRowMajorStridesChecked({2, limit, 2})));
  auto limitStrides = computeRowMajorStridesChecked({2, limit});
  TEST_CHECK(mlir::succeeded(limitStrides) && (*limitStrides)[0] == limit);
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({limit, 2}, {2, 1}, {2, 1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({limit, 2}, {3, 1}, {3, 1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({3}, {limit}, {1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({3}, {1}, {limit})));
  // A run start can fit even when its final element does not.
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({2, 2}, {limit, 1}, {2, 1})));
  TEST_CHECK(mlir::failed(CopyAccessPlan::build({2, 2}, {2, 1}, {limit, 1})));
  auto boundary = CopyAccessPlan::build({2}, {limit}, {1});
  TEST_CHECK(mlir::succeeded(boundary) && boundary->getRunOffsets(1)->first == limit);
  auto unit = CopyAccessPlan::build({1, 2}, {limit, 1}, {-1, 1});
  TEST_CHECK(mlir::succeeded(unit) && unit->getRunElements() == 2);

  // Independent element-by-element reference checks ordered offsets, not totals.
  std::mt19937 random(123);
  for (unsigned trial = 0; trial < 2000; ++trial) {
    size_t rank = random() % 4 + 1;
    llvm::SmallVector<int64_t> shape(rank), source(rank), target(rank);
    for (auto &size : shape) size = random() % 4 + 1;
    auto rowMajor = computeRowMajorStrides(shape);
    for (size_t dim = 0; dim < rank; ++dim) {
      source[dim] = random() % 2 ? rowMajor[dim] : random() % 20;
      target[dim] = random() % 2 ? rowMajor[dim] : random() % 20;
    }
    auto plan = CopyAccessPlan::build(shape, source, target);
    TEST_CHECK(mlir::succeeded(plan));
    int64_t linear = 0;
    for (int64_t run = 0; run < plan->getRunCount(); ++run) {
      auto offsets = plan->getRunOffsets(run);
      TEST_CHECK(mlir::succeeded(offsets));
      for (int64_t element = 0; element < plan->getRunElements(); ++element, ++linear) {
        auto indices = delinearizeIndex(linear, shape, rowMajor);
        TEST_CHECK(offsets->first + element == linearizeIndex(indices, source));
        TEST_CHECK(offsets->second + element == linearizeIndex(indices, target));
      }
    }
    TEST_CHECK(linear == getNumElements(shape));
  }

  // The cost caller and lowering use the same geometry; target timing stays separate.
  for (int64_t width : {8, 16}) {
    spatial::SchedulingTarget target;
    target.sharedMemoryWidthBytes = width;
    target.sharedMemoryReadLatencyNs = width == 8 ? 3 : 8;
    auto access = CopyAccessPlan::build({3, 5}, {8, 1}, {5, 1});
    TEST_CHECK(mlir::succeeded(access));
    auto demand = target.getSharedMemoryDemand(access->getRunElements() * 4, access->getRunCount(), false);
    TEST_CHECK(demand.readBytes == 60 && demand.readRequests == 3);
    TEST_CHECK(demand.serviceNs == 3 * ((20 + width - 1) / width) * target.sharedMemoryReadLatencyNs);
  }
}