#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_consolidated_exchange_oracle OUTPUT\n";
    return 2;
  }

  constexpr std::size_t transferCount = 4;
  constexpr std::size_t words = 256;
  constexpr unsigned sourceTiles[transferCount] = {5, 6, 7, 8};
  constexpr unsigned destinationTiles[transferCount] = {0, 0, 0, 1};

  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,
                                  {transferCount, words}, "source");
  auto destination = graph.addVariable(poplar::UNSIGNED_INT,
                                       {transferCount, words}, "destination");
  for (std::size_t transfer = 0; transfer < transferCount; ++transfer) {
    graph.setTileMapping(source[transfer], sourceTiles[transfer]);
    graph.setTileMapping(destination[transfer], destinationTiles[transfer]);
  }
  graph.createHostWrite("source-write", source.flatten());
  graph.createHostRead("destination-read", destination.flatten());

  poplar::program::Copy copy(source, destination);
  poplar::program::Repeat program(2, copy);
  poplar::OptionFlags options;
  options.set("debug.dumpGlobalExchangePackets", "true");
  options.set("debug.dumpDirectory", "sdk_consolidated_dumps");
  options.set("debug.retainDebugInformation", "true");
  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(transferCount * words);
  std::vector<std::uint32_t> result(input.size());
  for (std::size_t index = 0; index < input.size(); ++index)
    input[index] = 0x51a70000u ^ static_cast<std::uint32_t>(index * 0x9e3779b9u);

  engine.load(devices.front());
  engine.writeTensor("source-write", input.data(), input.data() + input.size());
  engine.run();
  engine.readTensor("destination-read", result.data(), result.data() + result.size());
  if (result != input)
    throw std::runtime_error("consolidated exchange produced incorrect data");
  std::cout << "repeated consolidated exchange passed\n";
  return 0;
} catch (const std::exception &error) {
  std::cerr << "sdk_consolidated_exchange_oracle: " << error.what() << '\n';
  return 1;
}