pe_dpe

#ifndef PE_DPE_
#define PE_DPE_

#include "pe_base"

namespace pe {
#if ENABLE_ZMQ
struct DpeCommand {
  std::string name;
  std::vector<std::vector<int64>> args;
  std::string ToString() const {
    std::stringstream ss;
    ss << name;
    ss << "(";
    int arg_idx = 0;
    for (const auto &arg : args) {
      if (arg_idx > 0) ss << ",";
      if (std::size(arg) == 1) {
        ss << arg[0];
      } else {
        ss << "{";
        int val_idx = 0;
        for (const int64 val : arg) {
          if (val_idx > 0) ss << ",";
          ss << val;
          ++val_idx;
        }
        ss << "}";
      }
      ++arg_idx;
    }
    ss << ")";
    return ss.str();
  }
};

SL std::ostream &operator<<(std::ostream &o, const DpeCommand &cmd) {
  return o << cmd.ToString();
}
namespace internal {

SL std::optional<std::string> ParseDpeToken(std::string_view cmd, int &i) {
  while (i < std::size(cmd) && std::isspace(cmd[i])) ++i;
  if (i == std::size(cmd)) return "$";

  if (cmd[i] == '(' || cmd[i] == ')' || cmd[i] == '{' || cmd[i] == '}' ||
      cmd[i] == ',') {
    return std::string(cmd.substr(i++, 1));
  } else if (std::isalpha(cmd[i])) {
    int j = i;
    while (i < std::size(cmd) && std::isalpha(cmd[i])) {
      ++i;
    }
    return std::string(cmd.substr(j, i - j));
  } else if (std::isdigit(cmd[i]) || cmd[i] == '+' || cmd[i] == '-') {
    int j = i;
    while (i < std::size(cmd) && (cmd[i] == '+' || cmd[i] == '-')) {
      ++i;
    }
    if (i == std::size(cmd) || !std::isdigit(cmd[i])) {
      return std::nullopt;
    }
    while (i < std::size(cmd) && std::isdigit(cmd[i])) {
      ++i;
    }
    return std::string(cmd.substr(j, i - j));
  }

  return std::nullopt;
}

SL std::optional<int64> ParseNumberToken(std::optional<std::string> s) {
  if (!s.has_value() || std::empty(s.value())) {
    return std::nullopt;
  }
  if (std::isdigit(s.value()[0]) || s.value()[0] == '+' ||
      s.value()[0] == '-') {
    std::string t = *s;
    const int size = std::size(t);
    int sgn = 1;
    int j = 0;
    while (j < size && (t[j] == '+' || t[j] == '-')) {
      if (t[j] == '-') {
        sgn = -sgn;
      }
      ++j;
    }
    int64 value = 0;
    while (j < size && std::isdigit(t[j])) {
      value = value * 10 + t[j] - '0';
      ++j;
    }
    return sgn < 0 ? -value : value;
  }
  return std::nullopt;
}

SL std::optional<std::vector<int64>> ParseDpeNumberList(std::string_view cmd,
                                                        int &i) {
  std::optional<std::string> token = ParseDpeToken(cmd, i);
  if (!token.has_value()) {
    return std::nullopt;
  }

  if (token == "{") {
    std::vector<int64> ret;
    for (;;) {
      std::optional<std::string> token = ParseDpeToken(cmd, i);
      if (!token.has_value()) {
        return std::nullopt;
      }
      if (token == "}") {
        break;
      }
      if (token == ",") {
        continue;
      }
      if (std::optional<int64> number = ParseNumberToken(token);
          number.has_value()) {
        ret.push_back(*number);
      } else {
        return std::nullopt;
      }
    }
    return ret;
  } else if (std::optional<int64> number = ParseNumberToken(token);
             number.has_value()) {
    return std::vector<int64>{*number};
  }

  return std::nullopt;
}

SL std::optional<std::vector<std::vector<int64>>> ParseDpeArgs(
    std::string_view cmd, int &i) {
  std::optional<std::string> start = ParseDpeToken(cmd, i);

  if (!start.has_value()) {
    return std::nullopt;
  }

  if (start == "$") {
    return std::vector<std::vector<int64>>{};
  }

  if (start != "(") {
    return std::nullopt;
  }

  std::vector<std::vector<int64>> ret;
  for (;;) {
    std::optional<std::string> token = ParseDpeToken(cmd, i);
    if (!token.has_value()) {
      return std::nullopt;
    }
    if (token == ")") {
      break;
    }
    if (token == ",") {
      continue;
    }
    std::string s = *std::move(token);
    if (s[0] == '{' || s[0] == '+' || s[0] == '-' || std::isdigit(s[0])) {
      i -= std::size(s);
      std::optional<std::vector<int64>> number_list =
          ParseDpeNumberList(cmd, i);
      if (!number_list.has_value()) {
        return std::nullopt;
      }
      ret.push_back(*std::move(number_list));
    } else {
      return std::nullopt;
    }
  }
  return ret;
}

SL std::optional<DpeCommand> ParseDpeCommand(std::string_view cmd, int &i) {
  std::optional<std::string> name = ParseDpeToken(cmd, i);
  if (!name.has_value() || std::empty(name.value()) ||
      !std::isalpha(name.value()[0])) {
    return std::nullopt;
  }

  std::optional<std::vector<std::vector<int64>>> args = ParseDpeArgs(cmd, i);
  if (!args.has_value()) {
    return std::nullopt;
  }

  return DpeCommand{.name = *std::move(name), .args = *std::move(args)};
}

SL std::optional<DpeCommand> ParseDpeCommand(std::string_view cmd) {
  int i = 0;
  return ParseDpeCommand(cmd, i);
}
}  // namespace internal

// Distributred PE
template <typename Derived>
class DpeServer {
 public:
  DpeServer(std::string_view server_address, std::string_view cache_path)
      : server_address(server_address), cache_path(cache_path) {}

  void RunReducer() {
    Initialize();

    void *context = zmq_ctx_new();
    void *responder = zmq_socket(context, ZMQ_REP);
    int rc = zmq_bind(responder, std::data(server_address));
    assert(rc == 0);

    constexpr int buffer_length = 1 << 20;
    char *buffer = new char[buffer_length + 1];
    for (; solved < total;) {
      int received = zmq_recv(responder, buffer, buffer_length, 0);
      buffer[received] = 0;
      std::string response = HandleRequest(std::string_view(buffer, received));
      zmq_send(responder, std::data(response), std::size(response), 0);
    }
    LogProgress();
    delete[] buffer;
    zmq_close(responder);
    zmq_ctx_destroy(context);
  }

 private:
  void Initialize() {
    std::vector<int64> tasks = AsDerived().GetTasks();

    cache.clear();
    const int64 size = std::size(tasks);
    for (int i = 0; i < size; ++i) {
      cache[tasks[i]] = std::nullopt;
    }

    answer = AsDerived().GetInit();
    total = std::size(cache);
    solved = 0;

    LoadCache();

    unsolved = std::queue<int64>();
    for (const auto [key, value] : cache) {
      if (value.has_value()) {
        ++solved;
        answer = AsDerived().Reduce(answer, key, *value);
      } else {
        unsolved.push(key);
      }
    }

    LogProgress();
  }

  void SaveCache() {
    if (std::empty(cache_path)) return;
    FILE *fp = fopen(std::data(cache_path), "w");
    if (!fp) return;

    for (const auto [key, value] : cache) {
      if (value.has_value()) {
        fprintf(fp, "%I64d %I64d\n", key, *value);
      }
    }
    fclose(fp);
  }

  void LoadCache() {
    if (std::empty(cache_path)) return;
    FILE *fp = fopen(std::data(cache_path), "r");
    if (!fp) return;

    int64 k, v;
    while (fscanf(fp, "%I64d %I64d", &k, &v) == 2) {
      auto where = cache.find(k);
      if (where != cache.end()) {
        where->second = v;
      }
    }
    fclose(fp);
  }

  std::string HandleRequest(std::string_view request) {
    std::optional<DpeCommand> command = internal::ParseDpeCommand(request);
    if (!command.has_value()) {
      return "error()";
    }

    if (command->name == "GetTasks") {
      if (std::size(command->args) != 1 || std::size(command->args[0]) != 1) {
        return "error()";
      }

      return AsDerived().HandleGetTasks(command->args[0][0]);
    }

    if (command->name == "SetResults") {
      std::vector<std::pair<int64, int64>> answers;

      for (const auto &iter : command->args) {
        if (std::size(iter) != 2) {
          return "error()";
        }
        answers.push_back({iter[0], iter[1]});
      }

      return AsDerived().HandleSetResults(answers);
    }
    return "error()";
  }

  std::string HandleGetTasks(int64 max_task_count) {
    std::vector<int64> tasks;
    for (int64 i = 0; i < max_task_count; ++i) {
      if (std::empty(unsolved)) {
        break;
      }
      tasks.push_back(unsolved.front());
      unsolved.pop();
    }

    DpeCommand raw_command;
    raw_command.name = "Tasks";
    raw_command.args.push_back(tasks);

    return raw_command.ToString();
  }

  std::string HandleSetResults(
      const std::vector<std::pair<int64, int64>> &answers) {
    for (const auto [key, value] : answers) {
      auto where = cache.find(key);
      if (where != cache.end()) {
        where->second = value;
        ++solved;
      }
      answer = AsDerived().Reduce(answer, key, value);
      if (solved % 1000 == 0) {
        SaveCache();
      }
      if (solved % 100 == 0) {
        LogProgress();
      }
    }
    return "ok()";
  }

  void LogProgress() {
    TimeDelta te = tr.Elapsed();

    char buff[32];
    sprintf(buff, "%2.2f%%", 100. * solved / total);

    std::string eta;
    while (std::size(history_progress) > 100) history_progress.pop();
    if (std::empty(history_progress) ||
        history_progress.front().second == solved) {
      eta = "n/a";
    } else {
      auto his = history_progress.front();
      // (total - solved) / eta = (solved - his.second) / (te - his.first)
      // eta = (total - solved) / (solved - his.second) * (te - his.first)

      eta = TimeDelta(duration_t(static_cast<int64>(
                          1. * (total - solved) / (solved - his.second) *
                          (te.NativeTime() - his.first.NativeTime()))))
                .Format();
    }
    history_progress.push({te, solved});
    std::cerr << std::setw(6) << buff << " " << solved << " " << total << " "
              << answer << " " << te.Format() << " " << eta << std::endl;
  }

  int64 GetInit() { return 0; }

  Derived &AsDerived() { return static_cast<Derived &>(*this); }

 private:
  std::string server_address;
  std::string cache_path;

  int64 solved = 0;

  std::map<int64, std::optional<int64>> cache;
  std::queue<int64> unsolved;
  int64 answer = 0;
  int64 total = 0;
  TimeRecorder tr;

  std::queue<std::pair<TimeDelta, int64>> history_progress;
};

template <typename Derived>
class DpeClient {
 public:
  DpeClient(std::string_view target)
      : target(target), buffer_length(1 << 20), buffer(nullptr) {}
  ~DpeClient() {
    if (buffer != nullptr) delete[] buffer;
  }

  std::vector<int64> GetTasks(int64 task_count) {
    void *context = zmq_ctx_new();
    void *requester = zmq_socket(context, ZMQ_REQ);
    zmq_connect(requester, std::data(target));

    DpeCommand raw_command;
    raw_command.name = "GetTasks";
    raw_command.args.push_back({task_count});
    std::string cmd = raw_command.ToString();

    zmq_send(requester, std::data(cmd), std::size(cmd), 0);
    int cnt = zmq_recv(requester, buffer, buffer_length, 0);
    zmq_close(requester);
    zmq_ctx_destroy(context);
    buffer[cnt] = 0;

    std::optional<DpeCommand> result =
        internal::ParseDpeCommand(std::string_view(buffer, cnt));
    if (!result.has_value() || result->name != "Tasks") {
      return {};
    }

    if (std::size(result->args) != 1) {
      return {};
    }

    return result->args[0];
  }

  int64 SetResults(const std::vector<std::pair<int64, int64>> &answers) {
    void *context = zmq_ctx_new();
    void *requester = zmq_socket(context, ZMQ_REQ);
    zmq_connect(requester, std::data(target));

    DpeCommand raw_command;
    raw_command.name = "SetResults";

    for (auto [key, value] : answers) {
      raw_command.args.push_back({key, value});
    }

    std::string cmd = raw_command.ToString();
    zmq_send(requester, std::data(cmd), std::size(cmd), 0);
    zmq_recv(requester, buffer, buffer_length, 0);
    zmq_close(requester);
    zmq_ctx_destroy(context);

    return 0;
  }

  void RunWorker(int thread_count, int task_count) {
    if (buffer == nullptr) {
      buffer = new char[buffer_length];
    }
#if ENABLE_OPENMP
    OmpLock lock;
#endif
    for (;;) {
      std::vector<int64> tasks = GetTasks(task_count);
      std::cerr << std::endl;
      std::cerr << "Received " << std::size(tasks) << " tasks" << std::endl;
      if (std::empty(tasks)) {
        break;
      }

      std::cerr << "First task: " << tasks[0] << std::endl;
      std::cerr << "Last task: " << tasks.back() << std::endl;

      TimeRecorder tr;
      std::vector<std::pair<int64, int64>> answers;
#if ENABLE_OPENMP
#pragma omp parallel for schedule(dynamic, 1) num_threads(thread_count)
#endif
      for (auto key : tasks) {
        int64 value = static_cast<Derived &>(*this).Cal(key);
#if ENABLE_OPENMP
        OmpGuard guard(lock);
#endif
        answers.push_back({key, value});
      }
      std::cerr << tr.Elapsed().Format() << std::endl;
      SetResults(answers);
    }
  }

 private:
  std::string target;
  const int buffer_length;
  char *buffer;
};

template <typename Derived>
class DpeSolver : public DpeServer<Derived>, public DpeClient<Derived> {
 public:
  DpeSolver(std::string_view server_address, std::string_view cache_path = "")
      : DpeServer<Derived>(server_address, cache_path),
        DpeClient<Derived>(server_address) {}
};

#endif
}  // namespace pe
#endif