<commit_before>#include <WorkerManager.h>
#include <BaseManager.h>
#include <RectangleArray.h>
#include "Util.h"
#include <algorithm>
using namespace BWAPI;
using namespace std;
using namespace Util;
WorkerManager::WorkerManager(Arbitrator::Arbitrator<Unit*,double>* arbitrator)
{
  this->arbitrator     = arbitrator;
  this->baseManager    = NULL;
  this->lastSCVBalance = 0;
  this->WorkersPerGas  = 3;
  this->mineralRate    = 0;
  this->gasRate        = 0;
}
void WorkerManager::setBaseManager(BaseManager* baseManager)
{
  this->baseManager = baseManager;
}
void WorkerManager::onOffer(set<Unit*> units)
{
  for(set<Unit*>::iterator u = units.begin(); u != units.end(); u++)
  {
    if ((*u)->getType().isWorker() && !this->mineralOrder.empty())
    {
      arbitrator->accept(this, *u);
      WorkerData temp;
      this->desiredWorkerCount[this->mineralOrder[this->mineralOrderIndex].first]++;
      this->currentWorkers[this->mineralOrder[this->mineralOrderIndex].first].insert(*u);
      temp.newResource = this->mineralOrder[this->mineralOrderIndex].first;
      this->mineralOrderIndex = (this->mineralOrderIndex+1) % mineralOrder.size();
      workers.insert(make_pair(*u,temp));
    }
    else
      arbitrator->decline(this, *u, 0);
  }
}
void WorkerManager::onRevoke(Unit* unit, double bid)
{
  this->onRemoveUnit(unit);
}

void WorkerManager::updateWorkerAssignments()
{
  //determine current worker assignments
  //also workers that are mining from resources that dont belong to any of our bases will be set to free

  for(map<Unit*,WorkerData >::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    if (w->second.newResource != NULL)
    {
      if (resourceBase.find(w->second.newResource) == resourceBase.end())
        w->second.newResource = NULL;
      else
        currentWorkers[w->second.newResource].insert(w->first);
    }
  }

  // get free workers
  set<Unit*> freeWorkers;
  for(map<Unit*,WorkerData>::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    if (w->second.newResource == NULL)
      freeWorkers.insert(w->first);
    else
    {
      // free workers that are too far away from their resources
      if (w->first->getDistance(w->second.newResource)>32*10)
      {
        freeWorkers.insert(w->first);
        //erase worker from resource's current workers set
        currentWorkers[w->second.newResource].erase(w->first);
      }
    }
  }

  // free workers from resources with too many workers
  for(map<Unit*,int>::iterator i = desiredWorkerCount.begin(); i != desiredWorkerCount.end(); i++)
    if (i->second < (int)currentWorkers[i->first].size())
    {
      // desired worker count is less than the current worker count for this resource, so lets remove some workers.
      int amountToRemove = currentWorkers[i->first].size() - i->second;
      for(int j = 0; j < amountToRemove; j++)
      {
        Unit* worker = *currentWorkers[i->first].begin();
        freeWorkers.insert(worker);
        workers[worker].newResource = NULL;
        currentWorkers[i->first].erase(worker);
      }
    }

  vector< Unit* > workerUnit;
  vector< Unit* > taskUnit;
  map<int,int> assignment;

  for(set<Unit*>::iterator i=freeWorkers.begin();i!=freeWorkers.end();i++)
    workerUnit.push_back(*i);

  // assign workers to resources that need more workers
  for(map<Unit*,int>::iterator i = desiredWorkerCount.begin(); i != desiredWorkerCount.end(); i++)
    if (i->second>(int)currentWorkers[i->first].size())
      for(int j=(int)currentWorkers[i->first].size();j<i->second;j++)
        taskUnit.push_back(i->first);

  //construct cost matrix
  //currently just uses euclidean distance, but walking distance would be more accurate
  RectangleArray<double> cost(workerUnit.size(),taskUnit.size());
  for(int w=0;w<(int)workerUnit.size();w++)
    for(int t=0;t<(int)taskUnit.size();t++)
      cost[w][t]=workerUnit[w]->getDistance(taskUnit[t]);

  //calculate assignment for workers and tasks (resources)
  assignment=computeAssignments(cost);

  //use assignment
  for(map<int,int>::iterator a=assignment.begin();a!=assignment.end();a++)
  {
    Unit* worker=workerUnit[a->first];
    Unit* resource=taskUnit[a->second];
    workers[worker].newResource = resource;
    currentWorkers[resource].insert(worker);
  }
}

bool mineralCompare (const pair<Unit*, int> i, const pair<Unit*, int> j) { return (i.second>j.second); }

void WorkerManager::rebalanceWorkers()
{
  mineralOrder.clear();
  desiredWorkerCount.clear();
  currentWorkers.clear();
  resourceBase.clear();
  int remainingWorkers = this->workers.size();
  
  for(map<Unit*,WorkerData >::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    currentWorkers[w->second.newResource].insert(w->first);
  }
  // iterate over all the resources of each active base
  for(set<Base*>::iterator b = this->basesCache.begin(); b != this->basesCache.end(); b++)
  {
    set<Unit*> baseMinerals = (*b)->getMinerals();
    for(set<Unit*>::iterator m = baseMinerals.begin(); m != baseMinerals.end(); m++)
    {
      resourceBase[*m] = *b;
      desiredWorkerCount[*m] = 0;
      mineralOrder.push_back(make_pair(*m, currentWorkers[*m].size()*400+(*m)->getResources() - 2*(int)(*m)->getPosition().getDistance((*b)->getBaseLocation()->getPosition())));
      currentWorkers[*m].clear();
    }
    set<Unit*> baseGeysers = (*b)->getGeysers();
    for(set<Unit*>::iterator g = baseGeysers.begin(); g != baseGeysers.end(); g++)
    {
      resourceBase[*g] = *b;
      desiredWorkerCount[*g]=0;
      if ((*g)->getType().isRefinery() && (*g)->getPlayer()==Broodwar->self() && (*g)->isCompleted())
      {
        for(int w=0;w<this->WorkersPerGas && remainingWorkers>0;w++)
        {
          desiredWorkerCount[*g]++;
          remainingWorkers--;
        }
      }
      currentWorkers[*g].clear();
    }<commit_msg>updated worker manager balancing algorithm<commit_after>#include <WorkerManager.h>
#include <BaseManager.h>
#include <RectangleArray.h>
#include "Util.h"
#include <algorithm>
using namespace BWAPI;
using namespace std;
using namespace Util;
WorkerManager::WorkerManager(Arbitrator::Arbitrator<Unit*,double>* arbitrator)
{
  this->arbitrator     = arbitrator;
  this->baseManager    = NULL;
  this->lastSCVBalance = 0;
  this->WorkersPerGas  = 3;
  this->mineralRate    = 0;
  this->gasRate        = 0;
}
void WorkerManager::setBaseManager(BaseManager* baseManager)
{
  this->baseManager = baseManager;
}
void WorkerManager::onOffer(set<Unit*> units)
{
  for(set<Unit*>::iterator u = units.begin(); u != units.end(); u++)
  {
    if ((*u)->getType().isWorker() && !this->mineralOrder.empty())
    {
      arbitrator->accept(this, *u);
      WorkerData temp;
      this->desiredWorkerCount[this->mineralOrder[this->mineralOrderIndex].first]++;
      this->currentWorkers[this->mineralOrder[this->mineralOrderIndex].first].insert(*u);
      temp.newResource = this->mineralOrder[this->mineralOrderIndex].first;
      this->mineralOrderIndex = (this->mineralOrderIndex+1) % mineralOrder.size();
      workers.insert(make_pair(*u,temp));
    }
    else
      arbitrator->decline(this, *u, 0);
  }
}
void WorkerManager::onRevoke(Unit* unit, double bid)
{
  this->onRemoveUnit(unit);
}

void WorkerManager::updateWorkerAssignments()
{
  //determine current worker assignments
  //also workers that are mining from resources that dont belong to any of our bases will be set to free

  for(map<Unit*,WorkerData >::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    if (w->second.newResource != NULL)
    {
      if (resourceBase.find(w->second.newResource) == resourceBase.end())
        w->second.newResource = NULL;
      else
        currentWorkers[w->second.newResource].insert(w->first);
    }
  }

  // get free workers
  set<Unit*> freeWorkers;
  for(map<Unit*,WorkerData>::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    if (w->second.newResource == NULL)
      freeWorkers.insert(w->first);
    else
    {
      // free workers that are too far away from their resources
      if (w->first->getDistance(w->second.newResource)>32*10)
      {
        freeWorkers.insert(w->first);
        //erase worker from resource's current workers set
        currentWorkers[w->second.newResource].erase(w->first);
      }
    }
  }

  // free workers from resources with too many workers
  for(map<Unit*,int>::iterator i = desiredWorkerCount.begin(); i != desiredWorkerCount.end(); i++)
    if (i->second < (int)currentWorkers[i->first].size())
    {
      // desired worker count is less than the current worker count for this resource, so lets remove some workers.
      int amountToRemove = currentWorkers[i->first].size() - i->second;
      for(int j = 0; j < amountToRemove; j++)
      {
        Unit* worker = *currentWorkers[i->first].begin();
        freeWorkers.insert(worker);
        workers[worker].newResource = NULL;
        currentWorkers[i->first].erase(worker);
      }
    }

  vector< Unit* > workerUnit;
  vector< Unit* > taskUnit;
  map<int,int> assignment;

  for(set<Unit*>::iterator i=freeWorkers.begin();i!=freeWorkers.end();i++)
    workerUnit.push_back(*i);

  // assign workers to resources that need more workers
  for(map<Unit*,int>::iterator i = desiredWorkerCount.begin(); i != desiredWorkerCount.end(); i++)
    if (i->second>(int)currentWorkers[i->first].size())
      for(int j=(int)currentWorkers[i->first].size();j<i->second;j++)
        taskUnit.push_back(i->first);

  //construct cost matrix
  //currently just uses euclidean distance, but walking distance would be more accurate
  RectangleArray<double> cost(workerUnit.size(),taskUnit.size());
  for(int w=0;w<(int)workerUnit.size();w++)
    for(int t=0;t<(int)taskUnit.size();t++)
      cost[w][t]=workerUnit[w]->getDistance(taskUnit[t]);

  //calculate assignment for workers and tasks (resources)
  assignment=computeAssignments(cost);

  //use assignment
  for(map<int,int>::iterator a=assignment.begin();a!=assignment.end();a++)
  {
    Unit* worker=workerUnit[a->first];
    Unit* resource=taskUnit[a->second];
    workers[worker].newResource = resource;
    currentWorkers[resource].insert(worker);
  }
}

bool mineralCompare (const pair<Unit*, int> i, const pair<Unit*, int> j) { return (i.second>j.second); }

void WorkerManager::rebalanceWorkers()
{
  mineralOrder.clear();
  desiredWorkerCount.clear();
  currentWorkers.clear();
  resourceBase.clear();
  int remainingWorkers = this->workers.size();
  
  for(map<Unit*,WorkerData >::iterator w = this->workers.begin(); w != this->workers.end(); w++)
  {
    currentWorkers[w->second.newResource].insert(w->first);
  }
  // iterate over all the resources of each active base
  for(set<Base*>::iterator b = this->basesCache.begin(); b != this->basesCache.end(); b++)
  {
    set<Unit*> baseMinerals = (*b)->getMinerals();
    vector< std::pair<Unit*,int> > baseMineralOrder;
    for(set<Unit*>::iterator m = baseMinerals.begin(); m != baseMinerals.end(); m++)
    {
      resourceBase[*m] = *b;
      desiredWorkerCount[*m] = 0;
      currentWorkers[*m].clear();
      baseMineralOrder.push_back(std::make_pair(*m,(*m)->getResources() - 2*(int)(*m)->getPosition().getDistance((*b)->getBaseLocation()->getPosition())));
    }
    sort(baseMineralOrder.begin(), baseMineralOrder.end(), mineralCompare);
    for(int i=0;i<(int)baseMineralOrder.size();i++)
    {
      Unit* mineral=baseMineralOrder[i].first;
      mineralOrder.push_back(make_pair(mineral, mineral->getResources() - 2*(int)mineral->getPosition().getDistance((*b)->getBaseLocation()->getPosition())-3000*i));
    }


    set<Unit*> baseGeysers = (*b)->getGeysers();
    for(set<Unit*>::iterator g = baseGeysers.begin(); g != baseGeysers.end(); g++)
    {
      resourceBase[*g] = *b;
      desiredWorkerCount[*g]=0;
      if ((*g)->getType().isRefinery() && (*g)->getPlayer()==Broodwar->self() && (*g)->isCompleted())
      {
        for(int w=0;w<this->WorkersPerGas && remainingWorkers>0;w++)
        {
          desiredWorkerCount[*g]++;
          remainingWorkers--;
        }
      }
      currentWorkers[*g].clear();
    }
{'fraction_non_alphanumeric': 0.1717579250720461, 'fraction_numerical': 0.0037052284890901604, 'mean_word_length': 3.885760257441673, 'pattern_counts': {'":': 0, '<': 70, '<?xml version=': 0, '>': 211, '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': '3132952', 'n_tokens_mistral': 3937, 'n_tokens_neox': 3867, 'n_words': 732}