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

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

using namespace poplar;

int main() {
  auto devManager = DeviceManager();
  auto devs = devManager.getDevices(TargetType::IPU, 1);
  if (devs.empty()) {
    std::cerr << "no IPU device available\n";
    return 2;
  }
  Device &device = devs[0];
  if (!device.attach()) {
    std::cerr << "failed to attach device\n";
    return 3;
  }

  Graph graph(device.getTarget());
  graph.addCodelets(std::vector<std::string>{"tile_id_codelets.cpp"}, "-O2", "ipu2");

  const unsigned numTiles = graph.getTarget().getNumTiles();
  Tensor physicalIds = graph.addVariable(UNSIGNED_SHORT, {numTiles}, "physicalIds");

  ComputeSet cs = graph.addComputeSet("readPhysicalTileIds");
  for (unsigned tile = 0; tile < numTiles; ++tile) {
    auto v = graph.addVertex(cs, "ReadPhysicalTileId", {{"out", physicalIds[tile]}});
    graph.setTileMapping(v, tile);
    graph.setTileMapping(physicalIds[tile], tile);
  }

  graph.createHostRead("physicalIdsRead", physicalIds);
  Engine engine(graph, program::Execute(cs));
  engine.load(device);
  engine.run(0);

  std::vector<unsigned short> ids(numTiles);
  engine.readTensor("physicalIdsRead", ids.data(), ids.data() + ids.size());

  for (unsigned tile = 0; tile < numTiles; ++tile) {
    std::cout << tile << " " << ids[tile] << "\n";
  }
}