#include <cstdint>
#include <fstream>
#include <iostream>
#include <string_view>

#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 != 2 && !(argc == 3 && std::string_view(argv[2]) == "--load")) {
    std::cerr << "usage: sdk_host_exchange_redirect_oracle OUTPUT [--load]\n";
    return 2;
  }
  constexpr std::size_t words = 16;
  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());
  auto allocation =
      graph.addVariable(poplar::UNSIGNED_INT, {2 * words}, "sources");
  auto destination =
      graph.addVariable(poplar::UNSIGNED_INT, {words}, "destination");
  graph.setTileMapping(allocation, 0);
  graph.setTileMapping(destination, 1);
  auto writeTarget = allocation.slice(0, words);
  auto redirectedTarget = allocation.slice(words, 2 * words);
  graph.createHostWrite("source-write", writeTarget);
  graph.createHostRead("destination-read", destination);

  poplar::program::Sequence program;
  program.add(poplar::program::Copy(redirectedTarget, destination));
  poplar::OptionFlags options;
  options.set("debug.printHostExchangeSchedule", "0,1");
  poplar::Engine engine(graph, program, options);
  std::ofstream output(argv[1], std::ios::binary);
  if (!output)
    throw std::runtime_error("cannot create output");
  engine.serializeExecutable(output);
  output.close();
  if (argc == 3)
    engine.load(device);
  return 0;
} catch (const std::exception &error) {
  std::cerr << "sdk_host_exchange_redirect_oracle: " << error.what() << '\n';
  return 1;
}