#include "exchange_plan.hpp"

#include <array>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>

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

using namespace poplar;

namespace {
constexpr unsigned guardWords = 16;
constexpr int guardValue = static_cast<int>(0x6badf00dU);

int sourceValue(unsigned index) {
  return static_cast<int>(0x13579bdfU ^ (index * 0x9e3779b9U));
}

unsigned arg(char **argv, int index) {
  return static_cast<unsigned>(std::strtoul(argv[index], nullptr, 0));
}
} // namespace

int main(int argc, char **argv) {
  if (argc < 5 || argc > 7) {
    std::cerr << "usage: one_to_many_exchange SENDER RECEIVER_A RECEIVER_B "
                 "COUNT [ITERATIONS] [serialized]\n";
    return 2;
  }
  const unsigned senderTile = arg(argv, 1);
  const std::array<unsigned, 2> receivers = {arg(argv, 2), arg(argv, 3)};
  const unsigned count = arg(argv, 4);
  const unsigned iterations = argc >= 6 ? arg(argv, 5) : 4;
  const bool serialized = argc == 7 && std::string(argv[6]) == "serialized";
  if (senderTile == receivers[0] || senderTile == receivers[1] ||
      receivers[0] == receivers[1] || count == 0 || iterations == 0) {
    std::cerr << "tiles must be distinct and counts must be positive\n";
    return 2;
  }

  xcom::FanOutPlan fanOut;
  try {
    const std::vector<unsigned> receiverList(receivers.begin(), receivers.end());
    fanOut = serialized
                 ? xcom::assembleSerializedFanOut(senderTile, receiverList, count)
                 : xcom::assembleSingleSendFanOutCandidate(
                       senderTile, receiverList, count);
  } catch (const std::exception &error) {
    std::cerr << "cannot assemble "
              << (serialized ? "serialized" : "single-send")
              << " fan-out: " << error.what() << '\n';
    return 2;
  }

  DeviceManager manager;
  auto devices = manager.getDevices(TargetType::IPU, 1);
  if (devices.empty() || !devices[0].attach()) {
    std::cerr << "failed to attach an IPU\n";
    return 3;
  }
  Device &device = devices[0];
  Graph graph(device.getTarget());
  graph.addCodelets(std::vector<std::string>{"no_blob_codelets.cpp"}, "-O2",
                    "ipu2");

  const unsigned sourceWords = count + guardWords * 2 + iterations;
  std::vector<int> sourceHost(sourceWords);
  for (unsigned i = 0; i < sourceWords; ++i) {
    sourceHost[i] = sourceValue(i);
  }
  Tensor source = graph.addVariable(INT, {sourceWords}, "source");
  graph.setInitialValue<int>(source, sourceHost);
  graph.setTileMapping(source, senderTile);

  Tensor plans = graph.addVariable(UNSIGNED_INT, {3, 9}, "plans");
  Tensor dummy = graph.addVariable(UNSIGNED_INT, {3, 1}, "dummy");
  graph.setTileMapping(plans[0], senderTile);
  graph.setTileMapping(dummy[0], senderTile);
  for (unsigned i = 0; i < 2; ++i) {
    graph.setTileMapping(plans[i + 1], receivers[i]);
    graph.setTileMapping(dummy[i + 1], receivers[i]);
  }

  Tensor elementSelector = graph.addVariable(INT, {}, "elementSelector");
  graph.setTileMapping(elementSelector, senderTile);
  Tensor tileSelectors = graph.addVariable(INT, {2}, "tileSelectors");
  Tensor results = graph.addVariable(
      INT, {2, count + guardWords * 2}, "resultStorage");
  std::vector<int> resultInitial(results.numElements(), guardValue);
  graph.setInitialValue<int>(results, resultInitial);

  const unsigned numTiles = graph.getTarget().getNumTiles();
  std::vector<unsigned short> lut(numTiles);
  for (unsigned tile = 0; tile < numTiles; ++tile) {
    lut[tile] = static_cast<unsigned short>(xcom::logicalToPhysical(tile));
  }
  std::vector<unsigned short> lutCopies;
  lutCopies.reserve(numTiles * 2);
  lutCopies.insert(lutCopies.end(), lut.begin(), lut.end());
  lutCopies.insert(lutCopies.end(), lut.begin(), lut.end());
  Tensor physicalLuts = graph.addConstant(UNSIGNED_SHORT, {2, numTiles},
                                          lutCopies.data(), "physicalLuts");
  for (unsigned i = 0; i < 2; ++i) {
    graph.setTileMapping(tileSelectors[i], receivers[i]);
    graph.setTileMapping(results[i], receivers[i]);
    graph.setTileMapping(physicalLuts[i], receivers[i]);
  }

  ComputeSet exchange = graph.addComputeSet("oneToManyExchange");
  auto send = graph.addVertex(exchange, "NoBlobSend",
                              {{"planBuf", plans[0]},
                               {"nonexecutableDummy", dummy[0]},
                               {"elementSelector", elementSelector},
                               {"data", source}});
  graph.setTileMapping(send, senderTile);
  for (unsigned i = 0; i < 2; ++i) {
    auto result = results[i].slice(guardWords, guardWords + count);
    auto recv = graph.addVertex(exchange, "NoBlobRecv",
                                {{"planBuf", plans[i + 1]},
                                 {"nonexecutableDummy", dummy[i + 1]},
                                 {"tileSelector", tileSelectors[i]},
                                 {"logicalToPhysicalTile", physicalLuts[i]},
                                 {"result", result}});
    graph.setTileMapping(recv, receivers[i]);
  }
  for (unsigned tile = 0; tile < numTiles; ++tile) {
    if (tile == senderTile || tile == receivers[0] || tile == receivers[1]) {
      continue;
    }
    auto inactive = graph.addVertex(exchange, "NoBlobNonParticipation");
    graph.setTileMapping(inactive, tile);
  }

  graph.createHostWrite("plansWrite", plans.flatten());
  graph.createHostWrite("elementSelectorWrite", elementSelector);
  graph.createHostWrite("tileSelectorsWrite", tileSelectors);
  graph.createHostRead("resultsRead", results.flatten());
  graph.createHostRead("sourceRead", source);

  Engine engine(graph, program::Sequence({program::Sync(SyncType::INTERNAL),
                                          program::Execute(exchange),
                                          program::Sync(SyncType::INTERNAL)}));
  engine.load(device);
  std::vector<std::uint32_t> planWords;
  planWords.insert(planWords.end(), fanOut.sender.begin(), fanOut.sender.end());
  for (const auto &receiverPlan : fanOut.receivers) {
    planWords.insert(planWords.end(), receiverPlan.begin(), receiverPlan.end());
  }
  engine.writeTensor("plansWrite", planWords.data(),
                     planWords.data() + planWords.size());
  const std::array<int, 2> selectors = {static_cast<int>(senderTile),
                                        static_cast<int>(senderTile)};
  engine.writeTensor("tileSelectorsWrite", selectors.data(),
                     selectors.data() + selectors.size());

  bool valuesCorrect = true;
  std::array<bool, 2> receiverValues = {true, true};
  bool guardsPreserved = true;
  for (unsigned iteration = 0; iteration < iterations; ++iteration) {
    const int offset = static_cast<int>(guardWords + iteration);
    engine.writeTensor("elementSelectorWrite", &offset, &offset + 1);
    engine.run(0);
    std::vector<int> output(results.numElements());
    engine.readTensor("resultsRead", output.data(), output.data() + output.size());
    const unsigned stride = count + guardWords * 2;
    for (unsigned receiver = 0; receiver < 2; ++receiver) {
      const unsigned base = receiver * stride;
      for (unsigned i = 0; i < guardWords; ++i) {
        guardsPreserved &= output[base + i] == guardValue;
        guardsPreserved &= output[base + guardWords + count + i] == guardValue;
      }
      for (unsigned i = 0; i < count; ++i) {
        receiverValues[receiver] &=
            output[base + guardWords + i] ==
            sourceHost[static_cast<unsigned>(offset) + i];
      }
    }
  }
  std::vector<int> sourceAfter(sourceWords);
  engine.readTensor("sourceRead", sourceAfter.data(),
                    sourceAfter.data() + sourceAfter.size());
  const bool sourcePreserved = sourceAfter == sourceHost;
  valuesCorrect = receiverValues[0] && receiverValues[1];
  const bool ok = valuesCorrect && guardsPreserved && sourcePreserved;
  std::cout << "sender=" << senderTile << " receivers=" << receivers[0] << ','
            << receivers[1] << " count=" << count
            << " mode=" << (serialized ? "serialized" : "single-send")
            << " valuesCorrect="
            << (valuesCorrect ? "yes" : "NO") << " receiverValues="
            << (receiverValues[0] ? "yes" : "NO") << ','
            << (receiverValues[1] ? "yes" : "NO") << " sourcePreserved="
            << (sourcePreserved ? "yes" : "NO") << " guardsPreserved="
            << (guardsPreserved ? "yes" : "NO") << '\n'
            << (ok ? "PASS\n" : "FAIL\n");
  return ok ? 0 : 1;
}