// Dump setup-generated sender and receiver plans for a contiguous count range.

#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <vector>

#include <poplar/DeviceManager.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>

using namespace poplar;

namespace {
unsigned argOr(char **argv, int argc, int idx, unsigned fallback) {
  return idx < argc ? static_cast<unsigned>(std::strtoul(argv[idx], nullptr, 0))
                    : fallback;
}
} // namespace

int main(int argc, char **argv) {
  const unsigned senderTile = argOr(argv, argc, 1, 0);
  const unsigned receiverTile = argOr(argv, argc, 2, 1286);
  const unsigned firstCount = argOr(argv, argc, 3, 1);
  const unsigned lastCount = argOr(argv, argc, 4, 128);
  if (senderTile == receiverTile || firstCount == 0 || lastCount < firstCount) {
    std::cerr << "invalid sender, receiver, or count range\n";
    return 2;
  }

  DeviceManager devManager;
  auto devs = devManager.getDevices(TargetType::IPU, 1);
  if (devs.empty() || !devs[0].attach()) {
    std::cerr << "failed to attach an IPU\n";
    return 3;
  }
  Device &device = devs[0];
  Graph graph(device.getTarget());
  graph.addCodelets("JDL.gp", CodeletFileType::Auto, "", "ipu2");

  const unsigned numCounts = lastCount - firstCount + 1;
  Tensor plans = graph.addVariable(UNSIGNED_INT, {numCounts, 2, 9}, "plans");
  program::Sequence setupPrograms;

  for (unsigned i = 0; i < numCounts; ++i) {
    const unsigned count = firstCount + i;
    Tensor countConst =
        graph.addConstant<unsigned>(UNSIGNED_INT, {}, count, "count");
    Tensor receiverConst = graph.addConstant<int>(INT, {}, receiverTile, "receiver");
    graph.setTileMapping(countConst, receiverTile);
    graph.setTileMapping(receiverConst, senderTile);
    graph.setTileMapping(plans[i][0], senderTile);
    graph.setTileMapping(plans[i][1], receiverTile);

    ComputeSet cs = graph.addComputeSet("setup_" + std::to_string(count));
    auto send = graph.addVertex(cs, "JDLSetupSend",
                                {{"planBuf", plans[i][0]},
                                 {"recvTile", receiverConst},
                                 {"count", countConst}});
    auto recv = graph.addVertex(cs, "JDLSetupRecv",
                                {{"planBuf", plans[i][1]}, {"count", countConst}});
    graph.setTileMapping(send, senderTile);
    graph.setTileMapping(recv, receiverTile);
    setupPrograms.add(program::Execute(cs));
  }

  graph.createHostRead("plansRead", plans.flatten());
  Engine engine(graph, setupPrograms);
  engine.load(device);
  engine.run(0);

  std::vector<std::uint32_t> words(numCounts * 2 * 9);
  engine.readTensor("plansRead", words.data(), words.data() + words.size());
  for (unsigned i = 0; i < numCounts; ++i) {
    std::cout << "count " << firstCount + i << ":";
    for (unsigned row = 0; row < 2; ++row) {
      std::cout << (row == 0 ? " send" : " recv");
      for (unsigned col = 0; col < 9; ++col) {
        std::cout << " 0x" << std::hex << words[(i * 2 + row) * 9 + col]
                  << std::dec;
      }
    }
    std::cout << '\n';
  }
}