#include <fstream>

#include "OpponentModel.h"

#include "Bases.h"
#include "Random.h"
#include "The.h"

using namespace UAlbertaBot;

OpeningPlan OpponentModel::predictEnemyPlan() const
{
    // Don't bother to predict on island maps.
    // Data from other maps will only be misleading.
    if (Bases::Instance().isIslandStart())
    {
        return OpeningPlan::Unknown;
    }

    struct PlanInfoType
    {
        int wins;
        int games;
        double weight;
    };
    PlanInfoType planInfo[int(OpeningPlan::Size)];

    // 1. Initialize.
    for (int plan = int(OpeningPlan::Unknown); plan < int(OpeningPlan::Size); ++plan)
    {
        planInfo[plan].wins = 0;
        planInfo[plan].games = 0;
        planInfo[plan].weight = 0.0;
    }

    // 2. Gather info.
    double weight = 1.0;
    for (const GameRecord * record : _pastGameRecords)
    {
        if (sameMatchup(*record))
        {
            PlanInfoType & info = planInfo[int(record->getEnemyPlan())];
            info.games += 1;
            info.weight += weight;
            weight *= 1.25;        // more recent game records are more heavily weighted
        }
    }

    // 3. Decide.
    // For now, set the most heavily weighted plan other than Unknown as the expected plan. Ignore the other info.
    OpeningPlan bestPlan = OpeningPlan::Unknown;
    double bestWeight = 0.0;
    for (int plan = int(OpeningPlan::Unknown) + 1; plan < int(OpeningPlan::Size); ++plan)
    {
        if (planInfo[plan].weight > bestWeight)
        {
            bestPlan = OpeningPlan(plan);
            bestWeight = planInfo[plan].weight;
        }
    }

    return bestPlan;
}

// Does the opponent seem to play the same strategy every game?
// If we're pretty sure, set _singleStrategy to true.
// So far, we only check the plan. We have plenty of other data that could be helpful.
void OpponentModel::considerSingleStrategy()
{
    // Humans, we assume, are unpredictable.
    if (Config::Skills::HumanOpponent)
    {
        _singleStrategy = false;
        return;
    }

    // Gather info.
    int knownPlan = 0;
    int unknownPlan = 0;
    std::set<OpeningPlan> plansSeen;

    for (const GameRecord * record : _pastGameRecords)
    {
        if (sameMatchup(*record))
        {
            OpeningPlan plan = record->getEnemyPlan();
            if (plan == OpeningPlan::Unknown)
            {
                unknownPlan += 1;
            }
            else
            {
                knownPlan += 1;
                plansSeen.insert(plan);
            }
        }
    }

    // Decide.
    // If we don't recognize the majority of plans, we're not sure.
    if (knownPlan >= 2 && plansSeen.size() == 1 && unknownPlan <= knownPlan)
    {
        _singleStrategy = true;
    }
}

// If the opponent model has collected useful information,
// set _recommendedOpening, the opening to play (or instructions for choosing it).
// Leaving _recommendedOpening blank continues as if the opponent model were turned off.
// Also fill in the opponent summary _summary, but unneeded fields are not set.
// This runs once before play starts, when all we know is the opponent
// and whatever the game records tell us about the opponent.
void OpponentModel::considerOpenings()
{
    // Gather basic information from the game records.
    for (const GameRecord * record : _pastGameRecords)
    {
        if (sameMatchup(*record))
        {
            ++_summary.totalGames;
            if (record->getWin())
            {
                ++_summary.totalWins;
            }
            OpeningInfoType & info = _summary.openingInfo[record->getOpeningName()];
            if (record->getMapName() == BWAPI::Broodwar->mapFileName())
            {
                info.sameGames += 1;
                if (record->getWin())
                {
                    info.sameWins += 1;
                }
            }
            else
            {
                info.otherGames += 1;
                if (record->getWin())
                {
                    info.otherWins += 1;
                }
            }
            if (record->getExpectedEnemyPlan() == record->getEnemyPlan())
            {
                // The plan was recorded as correctly predicted in that game.
                _summary.planInfo.sameGames += 1;
                if (record->getWin())
                {
                    _summary.planInfo.sameWins += 1;
                }
            }
            else
            {
                // The plan was not correctly predicted.
                _summary.planInfo.otherGames += 1;
                if (record->getWin())
                {
                    _summary.planInfo.otherWins += 1;
                }
            }
        }
    }

    UAB_ASSERT(_summary.totalWins == _summary.planInfo.sameWins + _summary.planInfo.otherWins, "bad total");
    UAB_ASSERT(_summary.totalGames == _summary.planInfo.sameGames + _summary.planInfo.otherGames, "bad total");

    // For the first games, stick to the counter openings based on the predicted plan.
    if (_summary.totalGames <= 5)
    {
        _recommendedOpening = getOpeningForEnemyPlan(_expectedEnemyPlan);
        return;										// with or without expected play
    }

    UAB_ASSERT(_summary.totalGames > 0 && _summary.totalWins >= 0, "bad total");
    UAB_ASSERT(_summary.openingInfo.size() > 0 && int(_summary.openingInfo.size()) <= _summary.totalGames, "bad total");

    // If we keep winning, stick to the winning track.
    if (_summary.totalWins == _summary.totalGames ||
        _singleStrategy && _summary.planInfo.sameWins > 0 &&
        _summary.planInfo.sameWins == _summary.planInfo.sameGames)   // Unknown plan is OK
    {
        // BWAPI::Broodwar->printf("winning track");
        _recommendedOpening = getOpeningForEnemyPlan(_expectedEnemyPlan);
        return;										// with or without expected play
    }

    // Randomly choose any opening that always wins, or always wins on this map.
    // This bypasses the map weighting below.
    // The algorithm is reservoir sampling with reservoir size = 1.
    // It gives equal probabilities without remembering all the elements.
    std::string alwaysWins;
    double nAlwaysWins = 0.0;
    std::string alwaysWinsOnThisMap;
    double nAlwaysWinsOnThisMap = 0.0;
    for (const auto & item : _summary.openingInfo)
    {
        const OpeningInfoType & info = item.second;
        if (info.sameWins + info.otherWins > 0 && info.sameWins + info.otherWins == info.sameGames + info.otherGames)
        {
            nAlwaysWins += 1.0;
            if (Random::Instance().flag(1.0 / nAlwaysWins))
            {
                alwaysWins = item.first;
            }
        }
        if (info.sameWins > 0 && info.sameWins == info.sameGames)
        {
            nAlwaysWinsOnThisMap += 1.0;
            if (Random::Instance().flag(1.0 / nAlwaysWinsOnThisMap))
            {
                alwaysWinsOnThisMap = item.first;
            }
        }
    }
    if (!alwaysWins.empty())
    {
        //BWAPI::Broodwar->printf("always wins");
        _recommendedOpening = alwaysWins;
        return;
    }
    if (!alwaysWinsOnThisMap.empty())
    {
        //BWAPI::Broodwar->printf("always wins on this map");
        _recommendedOpening = alwaysWinsOnThisMap;
        return;
    }

    // If we haven't decided yet:

    // Compute "weighted" win rates which combine map win rates and overall win rates, as an
    // estimate of the true win rate on this map. The estimate is ad hoc, using an assumption
    // that is sure to be wrong.
    for (auto it = _summary.openingInfo.begin(); it != _summary.openingInfo.end(); ++it)
    {
        OpeningInfoType & info = it->second;

        // Evidence provided by game results is proportional to the square root of the number of games.
        // So let's pretend that a game played on the same map provides mapPower times as much evidence
        // as a game played on another map.
        double mapPower = info.sameGames ? (info.sameGames + info.otherGames) / sqrt(info.sameGames) : 0.0;

        info.weightedWins = mapPower * info.sameWins + info.otherWins;
        info.weightedGames = mapPower * info.sameGames + info.otherGames;
    }

    if (_singleStrategy)
    {
        singleStrategyEnemyOpenings();
    }
    else
    {
        multipleStrategyEnemyOpenings();
    }
}

// The enemy always plays the same plan against us (we think).
// Seek the single opening that best counters it.
void OpponentModel::singleStrategyEnemyOpenings()
{
    const OpponentSummary & summary = getSummary();

    // Explore different actions this proportion of the time.
    // The number varies depending on the overall win rate: Explore less if we're usually winning.
    const double overallWinRate = double(summary.totalWins) / summary.totalGames;
    UAB_ASSERT(overallWinRate >= 0.0 && overallWinRate <= 1.0, "bad total");
    const double explorationRate = 0.05 + (1.0 - overallWinRate) * 0.30;

    // Decide whether to explore.
    if (summary.totalWins == 0 || Random::Instance().flag(explorationRate))
    {
        //BWAPI::Broodwar->printf("single strategy - explore");
        _recommendedOpening = getExploreOpening(summary);
        return;
    }

    // BWAPI::Broodwar->printf("single strategy - exploit");

    // We're not exploring. Choose an opening with the best weighted win rate.
    // This is a variation on the epsilon-greedy method (where epsilon is not a constant).
    double bestScore = -1.0;		// every opening will have a win rate >= 0
    double nBest = 1.0;
    for (auto it = summary.openingInfo.begin(); it != summary.openingInfo.end(); ++it)
    {
        const OpeningInfoType & info = it->second;

        double score = weightedWinRate(info.weightedWins, info.weightedGames);

        if (score > bestScore)
        {
            _recommendedOpening = it->first;
            bestScore = score;
            nBest = 1.0;
        }
        else if (bestScore - score < 0.001)	// we know score <= bestScore
        {
            // We choose randomly among openings with essentially equal score, using reservoir sampling.
            nBest += 1.0;
            if (Random::Instance().flag(1.0 / nBest))
            {
                _recommendedOpening = it->first;
            }
        }
    }
}

// The enemy plays more than one plan against us.
// Seek a mix of openings that will keep the enemy on its toes.
void OpponentModel::multipleStrategyEnemyOpenings()
{
    const OpponentSummary & summary = getSummary();

    std::vector< std::pair<std::string, double> > openings;

    // The exploration option. Its win rate is the mean weighted win rate of all openings tried.
    double totalWeightedWins = 0.0;
    double totalWeightedGames = 0.0;
    for (auto it = summary.openingInfo.begin(); it != summary.openingInfo.end(); ++it)
    {
        const OpeningInfoType & info = it->second;

        totalWeightedWins += info.weightedWins;
        totalWeightedGames += info.weightedGames;
    }

    double meanWeightedWinRate = weightedWinRate(totalWeightedWins, totalWeightedGames);
    double nextTotal = std::max(0.1, meanWeightedWinRate * meanWeightedWinRate);
    openings.push_back(std::pair<std::string, double>("explore", nextTotal));

    // The specific openings already tried.
    for (auto it = summary.openingInfo.begin(); it != summary.openingInfo.end(); ++it)
    {
        const OpeningInfoType & info = it->second;

        if (info.weightedWins > 0.0)
        {
            double rate = weightedWinRate(info.weightedWins, info.weightedGames);
            nextTotal += rate * rate;
            openings.push_back(std::pair<std::string, double>(it->first, nextTotal));
        }
    }

    // Choose randomly by weight.

    double w = Random::Instance().range(nextTotal);
    for (size_t i = 0; i < openings.size(); ++i)
    {
        if (w < openings[i].second)
        {
            _recommendedOpening = openings[i].first;
            break;
        }
    }

    //BWAPI::Broodwar->printf("multiple strategy choice %s", _recommendedOpening.c_str());

    if (_recommendedOpening == "explore")
    {
        _recommendedOpening = getExploreOpening(summary);
    }
}

// Return 0.0 if there are no games.
double OpponentModel::weightedWinRate(double weightedWins, double weightedGames) const
{
    return weightedGames < 0.1 ? 0.0 : weightedWins / weightedGames;
}

// Possibly update the expected enemy plan.
// This runs later in the game, when we may have more information.
void OpponentModel::reconsiderEnemyPlan()
{
    if (_planRecognizer.getPlan() != OpeningPlan::Unknown)
    {
        // We already know the actual plan. No need to form an expectation.
        return;
    }

    if (!_gameRecord.getEnemyIsRandom() || the.enemy()->getRace() == BWAPI::Races::Unknown)
    {
        // For now, we only update the expected plan if the enemy went random
        // and we have now learned its race. 
        // The new information should narrow down the possibilities.
        return;
    }

    // Don't reconsider too often.
    if (BWAPI::Broodwar->getFrameCount() % 12 != 8)
    {
        return;
    }

    // We set the new expected plan even if it is Unknown. Better to know that we don't know.
    _expectedEnemyPlan = predictEnemyPlan();
}

// Find the past game record which best matches the current game and remember it.
void OpponentModel::setBestMatch()
{
    int bestScore = -1;
    GameRecord * bestRecord = nullptr;

    for (GameRecord * record : _pastGameRecords)
    {
        int score = _gameRecord.distance(*record);
        if (score != -1 && (!bestRecord || score < bestScore))
        {
            bestScore = score;
            bestRecord = record;
        }
    }

    _bestMatch = bestRecord;
}

// We have decided to explore openings. Return an appropriate opening name.
// "random" means choose randomly among all openings.
// "matchup" means choose among openings for this matchup.
// getOpeningForEnemyPlan(_expectedEnemyPlan) means choose an opening to counter the predicted enemy plan.
std::string OpponentModel::getExploreOpening(const OpponentSummary & opponentSummary)
{
    UAB_ASSERT(opponentSummary.totalGames > 0, "no records");

    const double wrongPlanRate = double(opponentSummary.planInfo.otherGames) / opponentSummary.totalGames;
    // Is the predicted enemy plan likely to be right?
    if (opponentSummary.totalGames > 30 && Random::Instance().flag(0.75))
    {
        return "random";
    }
    else if (Random::Instance().flag(0.8 * wrongPlanRate * double(std::min(opponentSummary.totalGames, 20)) / 20.0))
    {
        return "matchup";
    }
    else
    {
        return getOpeningForEnemyPlan(_expectedEnemyPlan);
    }
}

// We expect the enemy to follow the given opening plan.
// Recommend an opening to counter that plan.
// The counters are configured; all we have to do is name the strategy mix.
// The empty opening "" means play the regular openings, no plan recognized.
std::string OpponentModel::getOpeningForEnemyPlan(OpeningPlan enemyPlan)
{
    if (enemyPlan == OpeningPlan::Unknown)
    {
        return "";
    }
    return "Counter " + OpeningPlanString(enemyPlan);
}

// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --

OpponentModel::OpponentModel()
    : _bestMatch(nullptr)
    , _singleStrategy(false)
    , _initialExpectedEnemyPlan(OpeningPlan::Unknown)
    , _expectedEnemyPlan(OpeningPlan::Unknown)
{
    std::string name = the.enemy()->getName();

    // Replace characters that the filesystem may not like with '_'.
    // TODO Obviously not a thorough job.
    std::replace(name.begin(), name.end(), ' ', '_');

    _filename = "om_" + name + ".txt";
}

// Read past game records from the opponent model file, and do initial analysis.
void OpponentModel::read()
{
    // This creates the skills, which is necessary for reading the data files.
    the.skillkit.initialize();

    if (Config::IO::ReadOpponentModel)
    {
        std::ifstream inFile;
        
        inFile.open(Config::IO::ReadDir + _filename);

        // There may not be a file to read. Check for a prepared file in the AI directory.
        if (!inFile.good())
        {
            inFile.clear();
            inFile.open(Config::IO::PreparedDataDir + _filename);
            if (!inFile.good())
            {
                // No prepared file either. That's OK.
                return;
            }
        }
        // At this point, we have a file to read.

        while (inFile.good())
        {
            // NOTE We allocate records here and never free them if valid.
            //      Their lifetime is the whole game.
            GameRecord * record = new GameRecord(inFile);
            if (record->isValid())
            {
                _pastGameRecords.push_back(record);
            }
            else
            {
                delete record;
            }
        }

        inFile.close();
    }

    // Make immediate decisions that may take into account the game records.
    // The initial expected enemy plan is set only here. That's the idea.
    // The current expected enemy plan may be reset later.
    _expectedEnemyPlan = _initialExpectedEnemyPlan = predictEnemyPlan();
    considerSingleStrategy();
    considerOpenings();
}

// Write the game records to the opponent model file.
void OpponentModel::write()
{
    if (Config::IO::WriteOpponentModel)
    {
        std::ofstream outFile(Config::IO::WriteDir + _filename, std::ios::trunc);

        // If it fails, there's not much we can do about it.
        if (outFile.bad())
        {
            return;
        }

        // The number of initial game records to skip over without rewriting.
        // In normal operation, nToSkip is 0 or 1.
        int nToSkip = 0;
        if (int(_pastGameRecords.size()) >= Config::IO::MaxGameRecords)
        {
            nToSkip = _pastGameRecords.size() - Config::IO::MaxGameRecords + 1;
        }

        // Rewrite any old records that were read in.
        // Not needed for local testing or for SSCAIT, necessary for other competitions.
        for (GameRecord * record : _pastGameRecords)
        {
            if (nToSkip > 0)
            {
                --nToSkip;
            }
            else
            {
                record->write(outFile);
            }
        }

        // And write the record of this game.
        _gameRecord.write(outFile);

        outFile.close();
    }
}

void OpponentModel::update()
{
    _planRecognizer.update();
    reconsiderEnemyPlan();

    if (Config::IO::ReadOpponentModel || Config::IO::WriteOpponentModel)
    {
        _gameRecord.update();

        // TODO the rest is turned off for now, not currently useful
        return;

        if (BWAPI::Broodwar->getFrameCount() % 32 == 31)
        {
            setBestMatch();
        }

        if (_bestMatch)
        {
            //_bestMatch->debugLog();
            //BWAPI::Broodwar->drawTextScreen(200, 10, "%cmatch %s %s", white, _bestMatch->mapName, _bestMatch->openingName);
            BWAPI::Broodwar->drawTextScreen(220, 6, "%cmatch", white);
        }
        else
        {
            BWAPI::Broodwar->drawTextScreen(220, 6, "%cno best match", white);
        }
    }
}

// Fill in the snapshot with a prediction of what the opponent may have at a given time.
// Not currently usable.
void OpponentModel::predictEnemy(int lookaheadFrames, PlayerSnapshot & snap) const
{
    const int t = BWAPI::Broodwar->getFrameCount() + lookaheadFrames;

    // Use the best-match past game record if possible.
    // Otherwise, take a current snapshot and call it the prediction.
    if (_bestMatch && _bestMatch->findClosestSnapshot(t, snap))
    {
        // All done.
    }
    else
    {
        snap.takeEnemy();
    }
}

// Is this game record for the same matchup as the current game?
// We usually ignore records which aren't.
bool OpponentModel::sameMatchup(const GameRecord & record) const
{
    return _gameRecord.sameMatchup(record);
}

// The recognized enemy opening plan.
OpeningPlan OpponentModel::getEnemyPlan() const
{
    return _planRecognizer.getPlan();
}

// String for displaying the recognized enemy opening plan in the UI.
std::string OpponentModel::getEnemyPlanString() const
{
    return OpeningPlanString(_planRecognizer.getPlan());
}

// String for displaying the expected enemy opening plan in the UI.
std::string OpponentModel::getExpectedEnemyPlanString() const
{
    return OpeningPlanString(_expectedEnemyPlan);
}

// The recognized enemy plan, or the current expected enemy plan if none.
OpeningPlan OpponentModel::getBestGuessEnemyPlan() const
{
    if (_planRecognizer.getPlan() != OpeningPlan::Unknown)
    {
        return _planRecognizer.getPlan();
    }
    return _expectedEnemyPlan;
}

// One of:
// 1 the recognized plan
// 2 the expected plan if the enemy is single-strategy
// 3 the expected plan if the enemy went random and we have learned their race
OpeningPlan OpponentModel::getDarnLikelyEnemyPlan() const
{
    if (_planRecognizer.getPlan() != OpeningPlan::Unknown)
    {
        return _planRecognizer.getPlan();
    }

    if (_singleStrategy)
    {
        return _expectedEnemyPlan;
    }

    if (_gameRecord.getEnemyIsRandom() &&
        the.enemy()->getRace() != BWAPI::Races::Unknown)
    {
        return _expectedEnemyPlan;
    }

    return OpeningPlan::Unknown;
}

// Based on past games, when may the enemy have air units?
// 0 means no prediction, otherwise it's a frame number (past or future doesn't matter).
// The answer may vary over time depending on what we think the enemy opening is.
// The caller will call repeatedly as long as it is interested in the answer.
// NOTE This is really "does the enemy have air tech?" Even a robo facility counts.
int OpponentModel::predictEnemyAirUnitTiming() const
{
	static OpeningPlan lastEnemyPlan = OpeningPlan::None;
	static int lastReturnValue = 0;

	OpeningPlan enemyPlanNow = getBestGuessEnemyPlan();		// may be Unknown, can never be None

	// Cache the prediction until and unless the enemy plan changes.
	if (enemyPlanNow == lastEnemyPlan)
	{
		return lastReturnValue;
	}
	lastEnemyPlan = enemyPlanNow;

	// Timings when the enemy follows any plan, or the same plan as this game.
	// 0 means never.
	int anyPlan = 0;			// earliest timing ever seen, for when there are few games
	std::vector<int> samePlan;	// all timings seen; we'll take one that's early-ish but not the earliest
	size_t countZeros = 0;		// how often the enemy never got air tech, when following the same plan

	samePlan.reserve(_pastGameRecords.size());

	for (const GameRecord * record : _pastGameRecords)
	{
		int timing = record->getFrameEnemyGetsAirUnits();
		if (timing)
		{
			anyPlan = (anyPlan == 0) ? timing : std::min(anyPlan, timing);
		}

		if (enemyPlanNow == record->getEnemyPlan())
		{
			if (timing)
			{
				samePlan.push_back(timing);
			}
			else
			{
				countZeros += 1;
			}
		}
	}

	// Figure out the return value.
	// Don't return it yet; we have to save it for the cache.
	if (samePlan.size() >= 4)
	{
		if (countZeros >= samePlan.size())
		{
			// Usually no air tech for the enemy. Predict that.
			lastReturnValue = 0;
		}
		else
		{
			std::sort(samePlan.begin(), samePlan.end());
			lastReturnValue = samePlan.at(samePlan.size() / 3 - 1);		// -1 because it is zero-based
		}
	}
	else
	{
		lastReturnValue = anyPlan;
	}

	return lastReturnValue;
}

OpponentModel & OpponentModel::Instance()
{
    static OpponentModel instance;
    return instance;
}
