#include <array>
#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <random>
#include <regex>
#include <sstream>
#include <string>
#include <vector>

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

#include "exchange_plan.hpp"

using namespace poplar;

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

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

std::vector<std::uint32_t> makePlan(unsigned senderTile, unsigned receiverTile,
                                    unsigned count) {
  const auto plan = xcom::assemble(senderTile, receiverTile, count);
  std::vector<std::uint32_t> words(plan.sender.begin(), plan.sender.end());
  words.insert(words.end(), plan.receiver.begin(), plan.receiver.end());
  return words;
}

unsigned short logicalToPhysicalTile(unsigned logicalTile) {
  return static_cast<unsigned short>(xcom::logicalToPhysical(logicalTile));
}

std::vector<unsigned short> makeLogicalToPhysicalLut(unsigned numTiles) {
  std::vector<unsigned short> lut(numTiles);
  for (unsigned tile = 0; tile < numTiles; ++tile) {
    lut[tile] = logicalToPhysicalTile(tile);
  }
  return lut;
}

std::vector<std::uint32_t> parseExactPlanFile(const std::string &path,
                                              unsigned senderTile,
                                              unsigned receiverTile) {
  std::ifstream in(path);
  if (!in) {
    throw std::runtime_error("failed to open plan file: " + path);
  }

  std::vector<std::uint32_t> senderRow;
  std::vector<std::uint32_t> receiverRow;
  std::regex hexRegex("0x([0-9a-fA-F]{1,8})");

  for (std::string line; std::getline(in, line);) {
    const std::string senderPrefix =
        "sender " + std::to_string(senderTile) + ":";
    const std::string receiverPrefix =
        "receiver " + std::to_string(receiverTile) + ":";
    const bool isSenderLine =
        line.rfind(senderPrefix, 0) == 0 || line.rfind("row0:", 0) == 0;
    const bool isReceiverLine =
        line.rfind(receiverPrefix, 0) == 0 || line.rfind("row1:", 0) == 0;
    if (!isSenderLine && !isReceiverLine) {
      continue;
    }
    std::vector<std::uint32_t> row;
    for (std::sregex_iterator it(line.begin(), line.end(), hexRegex), end;
         it != end; ++it) {
      row.push_back(
          static_cast<std::uint32_t>(std::stoul((*it)[1].str(), nullptr, 16)));
    }
    if (row.size() != 9) {
      throw std::runtime_error("expected 9 hex words in plan row: " + line);
    }
    if (isSenderLine) {
      senderRow = row;
    } else {
      receiverRow = row;
    }
  }

  if (senderRow.empty() || receiverRow.empty()) {
    throw std::runtime_error(
        "plan file must contain matching 'sender N:' and 'receiver N:' rows");
  }
  senderRow.insert(senderRow.end(), receiverRow.begin(), receiverRow.end());
  return senderRow;
}
} // namespace

int main(int argc, char **argv) {
  const unsigned senderTile =
      argc > 1 ? static_cast<unsigned>(std::strtoul(argv[1], nullptr, 0)) : 0;
  const unsigned receiverTileId =
      argc > 2 ? static_cast<unsigned>(std::strtoul(argv[2], nullptr, 0))
               : 1286;
  const unsigned lookupSize =
      argc > 3 ? static_cast<unsigned>(std::strtoul(argv[3], nullptr, 0)) : 3;
  const std::string planPath = argc > 4 ? argv[4] : "";
  const unsigned iterations =
      argc > 5 ? static_cast<unsigned>(std::strtoul(argv[5], nullptr, 0)) : 4;
  if (senderTile == receiverTileId) {
    std::cerr << "sender and receiver must be different tiles\n";
    return 2;
  }
  if (lookupSize < 1 || iterations < 1) {
    std::cerr << "lookup size and iterations must be positive\n";
    return 2;
  }

  auto devManager = DeviceManager();
  auto devs = devManager.getDevices(TargetType::IPU, 1);
  if (devs.empty()) {
    std::cerr << "no IPU device available\n";
    return 3;
  }
  Device &device = devs[0];
  if (!device.attach()) {
    std::cerr << "failed to attach device\n";
    return 4;
  }

  Graph graph(device.getTarget());
  if (senderTile >= graph.getTarget().getNumTiles() ||
      receiverTileId >= graph.getTarget().getNumTiles()) {
    std::cerr << "sender/receiver tile out of range\n";
    return 2;
  }
  graph.addCodelets(std::vector<std::string>{"no_blob_codelets.cpp"}, "-O2",
                    "ipu2");

  const unsigned dataElements = lookupSize + guardWords * 3 + iterations;
  const std::uint32_t hostNonce = std::random_device{}();
  std::vector<int> data_h(dataElements);
  for (unsigned i = 0; i < data_h.size(); ++i) {
    data_h[i] = sourceValue(i) ^ static_cast<int>(hostNonce);
  }

  Tensor data = graph.addVariable(INT, {dataElements}, "data");
  graph.setInitialValue<int>(data, data_h);
  graph.setTileMapping(data, senderTile);

  Tensor planBuf = graph.addVariable(UNSIGNED_INT, {2, 9}, "planBuf");
  Tensor dummy = graph.addVariable(UNSIGNED_INT, {2, 1}, "dummy");
  graph.setTileMapping(planBuf[0], senderTile);
  graph.setTileMapping(planBuf[1], receiverTileId);
  graph.setTileMapping(dummy[0], senderTile);
  graph.setTileMapping(dummy[1], receiverTileId);

  Tensor tileSelector = graph.addVariable(INT, {}, "tileSelector");
  Tensor elementSelector = graph.addVariable(INT, {}, "elementSelector");
  Tensor resultStorage =
      graph.addVariable(INT, {lookupSize + guardWords * 2}, "resultStorage");
  std::vector<int> resultInitial(resultStorage.numElements(), resultGuardValue);
  graph.setInitialValue<int>(resultStorage, resultInitial);
  Tensor result = resultStorage.slice(guardWords, guardWords + lookupSize);
  graph.setTileMapping(tileSelector, receiverTileId);
  graph.setTileMapping(elementSelector, senderTile);
  graph.setTileMapping(resultStorage, receiverTileId);

  auto lut_h = makeLogicalToPhysicalLut(graph.getTarget().getNumTiles());
  Tensor logicalToPhysicalTile = graph.addConstant(
      UNSIGNED_SHORT, {lut_h.size()}, lut_h.data(), "logicalToPhysicalTile");
  graph.setTileMapping(logicalToPhysicalTile, receiverTileId);
  Tensor countConst =
      graph.addConstant<unsigned>(UNSIGNED_INT, {}, lookupSize, "countConst");
  Tensor recvTileConst = graph.addConstant<unsigned>(
      UNSIGNED_INT, {1}, receiverTileId, "recvTileConst");
  graph.setTileMapping(countConst, receiverTileId);
  graph.setTileMapping(recvTileConst, senderTile);

  ComputeSet setupShapeCS = graph.addComputeSet("unusedSetupShapeCS");
  auto setupRecv =
      graph.addVertex(setupShapeCS, "NoBlobSetupRecv",
                      {{"planBuf", planBuf[1]}, {"count", countConst}});
  auto setupSend = graph.addVertex(setupShapeCS, "NoBlobSetupSend",
                                   {{"planBuf", planBuf[0]},
                                    {"recvTile", recvTileConst},
                                    {"count", countConst}});
  graph.setTileMapping(setupRecv, receiverTileId);
  graph.setTileMapping(setupSend, senderTile);

  ComputeSet exchangeCS = graph.addComputeSet("noBlobExchangeCS");
  auto recvVtx =
      graph.addVertex(exchangeCS, "NoBlobRecv",
                      {{"planBuf", planBuf[1]},
                       {"nonexecutableDummy", dummy[1]},
                       {"tileSelector", tileSelector},
                       {"logicalToPhysicalTile", logicalToPhysicalTile},
                       {"result", result}});
  auto sendVtx = graph.addVertex(exchangeCS, "NoBlobSend",
                                 {{"planBuf", planBuf[0]},
                                  {"nonexecutableDummy", dummy[0]},
                                  {"elementSelector", elementSelector},
                                  {"data", data}});
  graph.setTileMapping(sendVtx, senderTile);
  graph.setTileMapping(recvVtx, receiverTileId);
  for (unsigned tile = 0; tile < graph.getTarget().getNumTiles(); ++tile) {
    if (tile == senderTile || tile == receiverTileId) {
      continue;
    }
    auto inactiveVtx = graph.addVertex(exchangeCS, "NoBlobNonParticipation");
    graph.setTileMapping(inactiveVtx, tile);
  }

  graph.createHostWrite("planWrite", planBuf.flatten());
  graph.createHostWrite("sourceWrite", data);
  graph.createHostWrite("tileSelectorWrite", tileSelector);
  graph.createHostWrite("elementSelectorWrite", elementSelector);
  graph.createHostRead("resultStorageRead", resultStorage);
  graph.createHostRead("sourceRead", data);

  program::Sequence program({program::Sync(SyncType::INTERNAL),
                             program::Execute(exchangeCS),
                             program::Sync(SyncType::INTERNAL)});
  Engine engine(graph, program);
  engine.load(device);
  engine.writeTensor("sourceWrite", data_h.data(),
                     data_h.data() + data_h.size());

  auto plan = planPath.empty()
                  ? makePlan(senderTile, receiverTileId, lookupSize)
                  : parseExactPlanFile(planPath, senderTile, receiverTileId);
  engine.writeTensor("planWrite", plan.data(), plan.data() + plan.size());
  int tileSel = static_cast<int>(senderTile);
  engine.writeTensor("tileSelectorWrite", &tileSel, &tileSel + 1);
  bool valuesOk = true;
  bool guardsOk = true;
  for (unsigned iteration = 0; iteration < iterations; ++iteration) {
    int elemSel = static_cast<int>(guardWords + iteration);
    engine.writeTensor("elementSelectorWrite", &elemSel, &elemSel + 1);
    engine.run(0);

    std::vector<int> result_h(resultStorage.numElements());
    engine.readTensor("resultStorageRead", result_h.data(),
                      result_h.data() + result_h.size());
    for (unsigned i = 0; i < guardWords; ++i) {
      guardsOk &= result_h[i] == resultGuardValue;
      guardsOk &= result_h[guardWords + lookupSize + i] == resultGuardValue;
    }
    for (unsigned i = 0; i < lookupSize; ++i) {
      valuesOk &= result_h[guardWords + i] == data_h[elemSel + i];
    }
  }

  std::vector<int> sourceAfter(data_h.size());
  engine.readTensor("sourceRead", sourceAfter.data(),
                    sourceAfter.data() + sourceAfter.size());
  const bool sourcePreserved = sourceAfter == data_h;
  const bool ok = valuesOk && guardsOk && sourcePreserved;

  std::cout << "sender=" << senderTile << " receiver=" << receiverTileId
            << " count=" << lookupSize << " iterations=" << iterations
            << " hostNonce=0x" << std::hex << hostNonce << std::dec
            << " hostUpload=yes"
            << " valuesCorrect=" << (valuesOk ? "yes" : "NO")
            << " sourcePreserved=" << (sourcePreserved ? "yes" : "NO")
            << " guardsPreserved=" << (guardsOk ? "yes" : "NO") << "\n"
            << (ok ? "PASS\n" : "FAIL\n");
  return ok ? 0 : 1;
}