AIGlobals.SaveForUpgrade v2

--Mithy: New system for determining which upgrade an AI should be saving for, if any
--This is only run once per item or upgrade shopping check

--Map table for upgrades to save for, and from what rank / enemy rank
--Rank is the own War Rank to begin saving for this upgrade
--ERank an optional enemy War Rank at which to begin saving
--EHas is an optional bool that tells the AI to save if the enemy has this upgrade
-- all of these conditions are OR; any of them will trigger saving
-- item cost mod (e.g. coin purse) is taken into account when determining how much to save
SaveForUpgrades = {
    [1] = { Name = 'CGoldIncome01', Rank = 2 },
    [2] = { Name = 'CTroopNumber03', Rank = 6, ERank = 7, EHas = true },
    [3] = { Name = 'CTroopNumber05', Rank = 6, ERank = 7, EHas = true },
    [4] = { Name = 'CTroopNumber04', Rank = 7 },
    [5] = { Name = 'CTroopNumber06', Rank = 9 },
}
--Table for holding upcoming upgrades
local upgradeQueue = false

--Returns a number corresponding to this AI's current gold ranking among team AIs
function GetGoldRank(unit, brain)
    local goldTable = {}
    for k, abrain in ArmyBrains do
        if not abrain.TeamBrain and abrain.BrainController == 'AI' and IsAlly(brain:GetArmyIndex(), abrain:GetArmyIndex()) then
            table.insert(goldTable, {Army = abrain:GetArmyIndex(), Gold = abrain.mGold})
        end
    end
    if table.getn(goldTable) < 2 then
        return 1
    else
        table.sort(goldTable, sort_down_by 'Gold')
        return table.find(goldTable, brain:GetArmyIndex(), function(a, b) return a.Army == b end)
    end
    return false
end

--Main function for handling save logic; returns false, or upgrade name and upgrade cost
function SaveForUpgrade(unit, brain)
    LOG("SaveForUpgrade: "..repr(brain.Nickname).." / "..repr(unit:GetUnitId()).." @ "..repr(GetMinute()))
    local goldRank = GetGoldRank(unit, brain)
    local warRank = brain.Score.WarRank
    local enemyRank = GetEnemyTeamBrain(unit).Score.WarRank
    LOG("\tGold Rank: "..repr(goldRank)..", War Rank: "..repr(warRank)..", Enemy War Rank: "..repr(enemyRank))

    --Deferred init for the upgrade queue, to allow other mods to hook SaveForUpgrades
    if upgradeQueue == false then
        LOG("\tInitializing upgrade queue..")
        --Assign costs on startup
        for num, upgrade in SaveForUpgrades do
            upgrade.Cost = UpgradeCost(upgrade.Name)
        end
        --Then copy the map table to the queue
        upgradeQueue = table.deepcopy(SaveForUpgrades)
    end

    --Prune purchased upgrades
    while TeamHasUpgrade(unit, upgradeQueue[1].Name) do
        LOG("\tTeam has upgrade "..repr(upgradeQueue[1].Name)..", removing from queue..")
        table.remove(upgradeQueue, 1)
    end

    --Get a list of upgrades that we should be saving for now
    local currentUpgrades = {}
    LOG("\tUpgrades to save for:")
    for num, data in upgradeQueue do
        if (data.Rank and warRank >= data.Rank) or (data.ERank and enemyRank >= data.ERank) or (data.EHas and EnemyHasUpgrade(unit, data.Name)) then
            LOG("\t\t"..repr(data.Name).." - "..repr(data.Cost))
            table.insert(currentUpgrades, table.copy(data))
        end
    end
    if table.getn(currentUpgrades) < 1 then
        LOG("\t\t(none)")
    end

    --Return which one this AI should save for, if any
    local saveFor = currentUpgrades[goldRank]
    if saveFor then
        LOG("\tSaving for upgrade "..repr(saveFor.Name)..", cost "..repr(saveFor.Cost))
        return saveFor.Name, saveFor.Cost
    else
        LOG("\tNot saving.")
        return false
    end
end