<commit_before>#include "schedule.hpp"

#include <stdexcept>
#include <cassert>
#include <cmath>
#include <numeric>
#include <algorithm>
#include <iterator>
#include <vector>

void to_difficulty(std::vector<double>& perstage_avg)
{
  auto min = *std::min_element(perstage_avg.begin(), perstage_avg.end());
  for_each(perstage_avg.begin(), perstage_avg.end(), 
    [min](double& x){x /= min;}
  );
}


std::vector<unsigned int> solve(unsigned int stages, 
                                unsigned int nodes, 
                                const std::vector<double>& difficulty_per_stage)
{
  assert(difficulty_per_stage.size() == stages);
  if (nodes < stages) {
    throw std::runtime_error{"Must have more or equal number of nodes than stages"};
  }

  std::vector<unsigned int> nodes_per_stage(stages);
  
  auto x1 = std::accumulate(difficulty_per_stage.begin(), difficulty_per_stage.end(), 0.0);
  x1 = static_cast<double>(nodes) / x1;

  for (size_t i = 0; i < stages; ++i) {
    nodes_per_stage[i] = ceil(difficulty_per_stage[i] * x1);
  }

  int off_by;

  while ( (off_by = std::accumulate(nodes_per_stage.begin(), nodes_per_stage.end(), 0) - nodes) ) {
    auto max_stage = std::max_element(nodes_per_stage.begin(), nodes_per_stage.end());
    if (off_by < 0) {
      *max_stage += 1;
    } else {
      *max_stage -= 1;
    }
  }

  return nodes_per_stage;
}

void assign(const int local_rank,
            std::vector<unsigned int>& nodes_per_stage, 
            schedule_type& rank_assignm,
            int* local_stage)
{
  rank_assignm.clear();
  int rank = 0;
  for (size_t stage = 0; stage < nodes_per_stage.size(); ++stage) {
    auto nodes = nodes_per_stage[stage];
    rank_assignm.push_back(std::vector<int>(nodes));
    for (size_t j = 0; j < nodes; ++j) {
      rank_assignm.back()[j] = rank;
      if (local_rank == rank) {
        *local_stage = stage;
      }
      rank++;<commit_msg>dirty not at all permanent solution to constraining the first and last stage to one mpi rank :squirrel:<commit_after>#include "schedule.hpp"

#include <stdexcept>
#include <cassert>
#include <cmath>
#include <numeric>
#include <algorithm>
#include <iterator>
#include <vector>

void to_difficulty(std::vector<double>& perstage_avg)
{
  auto min = *std::min_element(perstage_avg.begin(), perstage_avg.end());
  for_each(perstage_avg.begin(), perstage_avg.end(), 
    [min](double& x){x /= min;}
  );
}


std::vector<unsigned int> solve(unsigned int stages, 
                                unsigned int nodes, 
                                const std::vector<double>& difficulty_per_stage)
{
  assert(difficulty_per_stage.size() == stages);
  if (nodes < stages) {
    throw std::runtime_error{"Must have more or equal number of nodes than stages"};
  }

  std::vector<unsigned int> nodes_per_stage(stages);

  // TODO for the record the following is ugly and I hate it
  auto constrained_begin = std::begin(difficulty_per_stage);
  // std::advance(constrained_begin, 1);
  auto constrained_end = std::end(difficulty_per_stage);
  // std::advance(constrained_end, -1);
  
  auto x1 = std::accumulate(constrained_begin, constrained_end, 0.0);
  x1 = static_cast<double>(nodes) / x1;

  for (size_t i = 0; i < stages; ++i) {
    if (i == 0 or i == stages - 1) {
      nodes_per_stage[i] = 1;
    } else {
      nodes_per_stage[i] = ceil(difficulty_per_stage[i] * x1);
    }
  }

  int off_by = 0;

  while ( (off_by = std::accumulate(nodes_per_stage.begin(), nodes_per_stage.end(), 0) - nodes) ) {
    auto max_stage = std::max_element(nodes_per_stage.begin(), nodes_per_stage.end());
    if (off_by < 0) {
      *max_stage += 1;
    } else {
      *max_stage -= 1;
    }
  }

  return nodes_per_stage;
}

void assign(const int local_rank,
            std::vector<unsigned int>& nodes_per_stage, 
            schedule_type& rank_assignm,
            int* local_stage)
{
  rank_assignm.clear();
  int rank = 0;
  for (size_t stage = 0; stage < nodes_per_stage.size(); ++stage) {
    auto nodes = nodes_per_stage[stage];
    rank_assignm.push_back(std::vector<int>(nodes));
    for (size_t j = 0; j < nodes; ++j) {
      rank_assignm.back()[j] = rank;
      if (local_rank == rank) {
        *local_stage = stage;
      }
      rank++;
{'fraction_non_alphanumeric': 0.14117369785529107, 'fraction_numerical': 0.008013198208814519, 'mean_word_length': 3.286868686868687, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7414272', 'n_tokens_mistral': 1521, 'n_tokens_neox': 1446, 'n_words': 327}