#include <fstream>
#include <iostream>
#include <cstdint>
#include <string>

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

int main(int argc, char **argv) try {
  if (argc != 5 && argc != 6) {
    std::cerr << "usage: sdk_host_exchange_codegen_oracle WORDS TILE "
                 "PREFIX_WORDS OUTPUT [PRINT_TILES]\n";
    return 2;
  }
  const std::size_t words = std::stoul(argv[1], nullptr, 0);
  const unsigned tile = std::stoul(argv[2], nullptr, 0);
  const std::size_t prefix = std::stoul(argv[3], nullptr, 0);
  if (words == 0)
    throw std::invalid_argument("WORDS must be nonzero");

  poplar::DeviceManager manager;
  auto devices = manager.getDevices(poplar::TargetType::IPU, 1);
  if (devices.empty() || !devices.front().attach())
    throw std::runtime_error("failed to attach one IPU");
  auto &device = devices.front();

  poplar::Graph graph(device.getTarget());
  if (tile >= graph.getTarget().getNumTiles())
    throw std::out_of_range("TILE");
  auto allocation = graph.addVariable(poplar::UNSIGNED_INT, {prefix + words},
                                      "allocation");
  graph.setTileMapping(allocation, tile);
  auto tensor = allocation.slice(prefix, prefix + words);
  graph.createHostWrite("write", tensor);
  graph.createHostRead("read", tensor);

  poplar::OptionFlags options;
  options.set("debug.printHostExchangeSchedule",
              argc == 6 ? argv[5] : std::to_string(tile));
  options.set("debug.dumpHostExchangePackets", "true");
  poplar::program::Sequence program;
  poplar::Engine engine(graph, program, options);
  std::ofstream output(argv[4], std::ios::binary);
  if (!output)
    throw std::runtime_error("cannot create output");
  engine.serializeExecutable(output);
  std::cout << "words=" << words << " tile=" << tile
            << " prefixWords=" << prefix << '\n';
  return 0;
} catch (const std::exception &error) {
  std::cerr << "sdk_host_exchange_codegen_oracle: " << error.what() << '\n';
  return 1;
}