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

#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) {
    std::cerr << "usage: sdk_two_pass_exchange_oracle OUTPUT\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");

  poplar::Graph graph(devices.front().getTarget());
  auto source = graph.addVariable(poplar::UNSIGNED_INT, {words}, "source");
  auto relay = graph.addVariable(poplar::UNSIGNED_INT, {words}, "relay");
  auto output = graph.addVariable(poplar::UNSIGNED_INT, {words}, "output");
  graph.setTileMapping(source, 0);
  graph.setTileMapping(relay, 1);
  graph.setTileMapping(output, 0);
  graph.createHostWrite("source-write", source);
  graph.createHostRead("output-read", output);

  poplar::program::Sequence program;
  program.add(poplar::program::Copy(source, relay));
  program.add(poplar::program::Copy(relay, output));
  poplar::OptionFlags options;
  options.set("debug.printHostExchangeSchedule", "0,1");
  poplar::Engine engine(graph, program, options);
  std::ofstream executable(argv[1], std::ios::binary);
  if (!executable)
    throw std::runtime_error("cannot create output");
  engine.serializeExecutable(executable);
  executable.close();

  std::vector<std::uint32_t> input(words);
  std::vector<std::uint32_t> result(words);
  for (std::size_t index = 0; index < words; ++index)
    input[index] = 0x51a70000u ^ static_cast<std::uint32_t>(index * 0x9e3779b9u);
  std::cerr << "BEGIN load\n";
  engine.load(devices.front());
  std::cerr << "END load\nBEGIN writeTensor\n";
  engine.writeTensor("source-write", input.data(), input.data() + input.size());
  std::cerr << "END writeTensor\nBEGIN run\n";
  engine.run();
  std::cerr << "END run\nBEGIN readTensor\n";
  engine.readTensor("output-read", result.data(), result.data() + result.size());
  std::cerr << "END readTensor\n";
  if (result != input)
    throw std::runtime_error("two-pass exchange produced incorrect data");
  std::cout << "two-pass host/exchange round trip passed\n";
  return 0;
} catch (const std::exception &error) {
  std::cerr << "sdk_two_pass_exchange_oracle: " << error.what() << '\n';
  return 1;
}