/* Be sure to to change all the protoss references if you aren't playing as protoss */

#include "BasicAIModule.h"
#include <set>
#include <list>
#include <utility>
using namespace BWAPI;
using namespace std;

// Starting base location
BWTA::BaseLocation* home;
// All workers assigned to gathering gas
UnitGroup gasWorkers;
// Set of all assimilators and the number workers assigned to them
std::set<std::pair<Unit*, int> > workersPerAssimilator;
// Closest minerals to start location
Unit* closestMinerals;

int WORKERS_PER_ASSIMILATOR = 3;

void BasicAIModule::onStart()
{
    this->showManagerAssignments = false;
    if (Broodwar->isReplay()) return;

    // Enable some cheat flags
    Broodwar->enableFlag(Flag::UserInput);

    BWTA::readMap();
    BWTA::analyze();

    // Let's speed things up a bit...
    Broodwar->setLocalSpeed(16);

    this->analyzed           = true;
    this->buildManager       = new BuildManager(&this->arbitrator);
    this->techManager        = new TechManager(&this->arbitrator);
    this->upgradeManager     = new UpgradeManager(&this->arbitrator);
    this->scoutManager       = new ScoutManager(&this->arbitrator);
    this->workerManager      = new WorkerManager(&this->arbitrator);
    this->buildOrderManager  = new BuildOrderManager(this->buildManager,this->techManager,this->upgradeManager,this->workerManager);
    this->baseManager        = new BaseManager();
    this->supplyManager      = new SupplyManager();
    this->defenseManager     = new DefenseManager(&this->arbitrator);
    this->informationManager = new InformationManager();
    this->unitGroupManager   = new UnitGroupManager();
    this->enhancedUI         = new EnhancedUI();

    this->supplyManager->setBuildManager(this->buildManager);
    this->supplyManager->setBuildOrderManager(this->buildOrderManager);
    this->techManager->setBuildingPlacer(this->buildManager->getBuildingPlacer());
    this->upgradeManager->setBuildingPlacer(this->buildManager->getBuildingPlacer());
    this->workerManager->setBaseManager(this->baseManager);
    this->workerManager->setBuildOrderManager(this->buildOrderManager);
    this->baseManager->setBuildOrderManager(this->buildOrderManager);

    BWAPI::Race race = Broodwar->self()->getRace();
    BWAPI::Race enemyRace = Broodwar->enemy()->getRace();
    BWAPI::UnitType workerType = *(race.getWorker());
    double minDist;
    BWTA::BaseLocation* natural = NULL;
    home = BWTA::getStartLocation(Broodwar->self());
    for each (BWTA::BaseLocation* b in BWTA::getBaseLocations())
    {
        if (b == home) continue;
        double dist = home->getGroundDistance(b);
        if (dist > 0)
        {
            if (natural == NULL || dist < minDist)
            {
                minDist = dist;
                natural = b;
            }
        }
    }

    // Find the closest minerals to the start location
    std::set<Unit*> minerals = home->getMinerals();
    Unit* nexus = *(SelectAll(UnitTypes::Protoss_Nexus).begin());
    double min;
    for each (Unit* m in minerals)
    {
        double curr = BWTA::getGroundDistance(nexus->getTilePosition(), m->getTilePosition());  
        
        if (curr < min)
        {
            min = curr;
            closestMinerals = m;
        }
    }

    this->buildOrderManager->enableDependencyResolver();
    
    /* Place your build order here */

    this->workerManager->enableAutoBuild();
    this->workerManager->setAutoBuildPriority(40);
}

void BasicAIModule::onFrame()
{
    if (Broodwar->isReplay()) return;
    if (!this->analyzed) return;
    this->buildManager->update();
    this->buildOrderManager->update();
    this->baseManager->update();
    this->workerManager->update();
    this->techManager->update();
    this->upgradeManager->update();
    this->supplyManager->update();
    this->scoutManager->update();
    this->defenseManager->update();
    this->arbitrator.update();

    this->enhancedUI->update();

    if (Broodwar->getFrameCount() == 24*50)
        scoutManager->setScoutCount(1);

    // Only do these every second, for performance reasons
    if (Broodwar->getFrameCount() % 24 == 0)
    {
        // When an assimilator is constructed, assign workers to it
        for (std::set<std::pair<Unit*, int> >::iterator i = workersPerAssimilator.begin(); i != workersPerAssimilator.end(); i++)
        {
            // If the assimilator is still being constructed, or WORKERS_PER_ASSIMILATOR workers are
            // already asigned to it, ignore it for now
            if (!((*i).first->isCompleted()) || (*i).second >= WORKERS_PER_ASSIMILATOR) continue;

            // Assign a maximum of WORKERS_PER_ASSIMILATOR workers to each available assimilator
            int probeCount = 0;
            UnitGroup tempGasWorkers = SelectAll(Broodwar->self(), UnitTypes::Protoss_Probe)(FilterFlag::isIdle) + SelectAll(Broodwar->self(), UnitTypes::Protoss_Probe)(FilterFlag::isGatheringMinerals);
            tempGasWorkers = tempGasWorkers - gasWorkers;
            for (UnitGroup::iterator w = tempGasWorkers.begin(); w != tempGasWorkers.end() && probeCount < WORKERS_PER_ASSIMILATOR; w++, probeCount++)
            {
                (*w)->stop();
                (*w)->rightClick((*i).first);
                gasWorkers.insert((*w));
                (*i).second = (*i).second + 1;
            }
        }

        // Make sure all idle workers are always gathering minerals
        UnitGroup workers = SelectAll(Broodwar->self(), UnitTypes::Protoss_Probe)(FilterFlag::isIdle);
        workers = workers - gasWorkers;
        for each (Unit* w in workers)
        {
            w->rightClick(closestMinerals);
        }
    }

    std::set<Unit*> units = Broodwar->self()->getUnits();
    if (this->showManagerAssignments)
    {
        for(std::set<Unit*>::iterator i = units.begin(); i != units.end(); i++)
        {
            if (this->arbitrator.hasBid(*i))
            {
                int x = (*i)->getPosition().x();
                int y = (*i)->getPosition().y();
                std::list<std::pair<Arbitrator::Controller<BWAPI::Unit*,double>*, double> > bids = this->arbitrator.getAllBidders(*i);
                int y_off = 0;
                bool first = false;
                const char activeColor = '\x07', inactiveColor = '\x16';
                char color = activeColor;
                for(std::list<std::pair<Arbitrator::Controller<BWAPI::Unit*,double>*, double> >::iterator j = bids.begin(); j != bids.end(); j++)
                {
                    Broodwar->drawText(CoordinateType::Map, x, y+y_off, "%c%s: %d", color, j->first->getShortName().c_str(), (int)j->second);
                    y_off += 15;
                    color = inactiveColor;
                }
            }
        }
    }

    UnitGroup myPylonsAndGateways = SelectAll()(Pylon, Gateway)(HitPoints, "<=", 200);
    for each(Unit* u in myPylonsAndGateways)
    {
        Broodwar->drawCircleMap(u->getPosition().x(), u->getPosition().y(), 20, Colors::Red);
    }
}

void BasicAIModule::onUnitDestroy(BWAPI::Unit* unit)
{
    this->arbitrator.onRemoveObject(unit);
    this->buildManager->onRemoveUnit(unit);
    this->techManager->onRemoveUnit(unit);
    this->upgradeManager->onRemoveUnit(unit);
    this->workerManager->onRemoveUnit(unit);
    this->scoutManager->onRemoveUnit(unit);
    this->defenseManager->onRemoveUnit(unit);
    this->informationManager->onUnitDestroy(unit);
}

void BasicAIModule::onUnitShow(BWAPI::Unit* unit)
{
    this->informationManager->onUnitShow(unit);
    this->unitGroupManager->onUnitShow(unit);
}
void BasicAIModule::onUnitHide(BWAPI::Unit* unit)
{
    this->informationManager->onUnitHide(unit);
    this->unitGroupManager->onUnitHide(unit);
}

void BasicAIModule::onUnitMorph(BWAPI::Unit* unit)
{
    this->unitGroupManager->onUnitMorph(unit);

    // Once an assimilator has been constructed, add it to the set of assimilators
    if (unit->getPlayer() == Broodwar->self() && unit->getType() == UnitTypes::Protoss_Assimilator)
    {
        workersPerAssimilator.insert(std::pair<Unit*, int>(unit, 0));
    }
}
void BasicAIModule::onUnitRenegade(BWAPI::Unit* unit)
{
    this->unitGroupManager->onUnitRenegade(unit);
}

bool BasicAIModule::onSendText(std::string text)
{
    UnitType type = UnitTypes::getUnitType(text);
    if (text == "debug")
    {
        this->showManagerAssignments = true;
        this->buildOrderManager->enableDebugMode();
        this->scoutManager->enableDebugMode();

        return true;
    }
    if (text == "expand")
    {
        this->baseManager->expand(100);
    }
    if (type != UnitTypes::Unknown)
    {
        this->buildOrderManager->buildAdditional(1, type, 300);
    }
    else
    {
        TechType type = TechTypes::getTechType(text);
        if (type != TechTypes::Unknown)
        {
            this->techManager->research(type);
        }
        else
        {
            UpgradeType type = UpgradeTypes::getUpgradeType(text);
            if (type != UpgradeTypes::Unknown)
            {
                this->upgradeManager->upgrade(type);
            }
            else
            {
                Broodwar->printf("You typed '%s'!", text.c_str());
            }
        }
    }

    return true;
}