Mithy

--AIGlobals.lua
local Game = import('/lua/game.lua')
local Upgrades = import('/lua/common/CitadelUpgrades.lua').Upgrades


--Utility functions for expanding PriorityFunction versatility
function GetDeathMult()
    return Game.GameData.DeathPenaltyMultiplier[ScenarioInfo.Options.DeathPenalty or 'Normal']
end

function GetGoldMult()
    return Game.GameData.GoldIncomeMultiplier[ScenarioInfo.Options.GoldIncome or 'Normal']
end

function GetXPMult()
    return Game.GameData.ExperienceMultiplier[ScenarioInfo.Options.ExperienceRate or 'Normal']
end

function GetMinute()
    return math.floor(GetGameTimeSeconds() / 60)
end

function IsGeneral(unit)
    return EntityCategoryContains(categories.GENERAL, unit)
end

function IsAssassin(unit)
    return EntityCategoryContains(categories.ASSASSIN, unit)
end

function UpgradeCost(upgradeName)
    if Upgrades.Tree[upgradeName] then
        return Upgrades.Tree[upgradeName].Cost
    else
        WARN("AIGlobals.UpgradeCost -- Tried to get cost of non-existent upgrade '"..upgradeName.."'")
        return 0
    end
end

--Returns a fractional modifier from 0 to 1 based on the provided demigod's gold reserves, starting at the minimum specified
--amount of gold and ending at the min amount * rec multiplier.  If the rec multiplier is not specified, it defaults to 2.
function GoldThreshold(unit, min, rec)
    rec = min * (rec or 2)
    return math.min(1, math.max(0, unit:GetGold() - min) / (rec - min))
end

--Get our team's average hero level
function GetAverageLevel(unit)
    local levels, heroes = unit.Score.HeroLevel, 1
    for k, brain in unit:GetAIBrain():GetAlliedArmies() do
        if not brain.IsTeamArmy then
            heroes = heroes + 1
            levels = levels + brain.Score.HeroLevel
        end
    end
    return math.floor( (levels / heroes) + 0.5 )
end

--Get our team's total deaths
function GetTeamDeaths(unit)
    local deaths = unit.Score.HeroDeaths
    for k, brain in unit:GetAIBrain():GetAlliedArmies() do
        if not brain.IsTeamArmy then
            deaths = deaths + brain.Score.HeroDeaths
        end
    end
    return deaths
end

--Get our team's citadel health fraction
function GetCitadelHealth(unit)
    local cithealth = false
    local citadel = unit:GetAIBrain():GetStronghold()

    if citadel then
        cithealth = citadel:GetHealth() / citadel:GetMaxHealth()
    end

    return cithealth
end

function GetEnemyTeamBrain(unit)
    local unitbrain = unit:GetAIBrain()

    if not unitbrain.EnemyTeamBrain then
        for k, brain in ArmyBrains do
            if brain.TeamBrain and IsEnemy(unit:GetArmy(), brain:GetArmyIndex()) then
                unitbrain.EnemyTeamBrain = brain
            end
        end
    end

    return unitbrain.EnemyTeamBrain
end

--Used to get info on enemy citadel upgrades
function HasUpgradeByTeamBrain(brain, upgradeName)
    local upgradesTable = brain:GetStronghold().Sync.Upgrades

    if upgradesTable then
        for k,v in upgradesTable do
            if v == upgradeName then
                return true
            end
        end
    end

    return false
end

--Check our team's upgrades
function TeamHasUpgrade(unit, upgradeName)
    local hasupgrade = false
    local tBrain = unit:GetAIBrain():GetTeamArmy()

    if tBrain then
        hasupgrade = HasUpgradeByTeamBrain(tBrain, upgradeName)
    end

    return hasupgrade
end

--Check enemy team's upgrades
function EnemyHasUpgrade(unit, upgradeName)
    local hasupgrade = false
    local etBrain = GetEnemyTeamBrain(unit)

    if etBrain then
        local hasupgrade = HasUpgradeByTeamBrain(etBrain, upgradeName)
    end

    return hasupgrade
end



--Try to get the AI to respond to pings a bit more reliably
AIPingMaximumWeight = 90 # 37




--Citadel Upgrade priorities
CitadelUpgradeWeights = {
    # Fortified Structure - prioritize on enemy catapults and citadel health
    CBuildingHealth01 = {
        PriorityFunction = function(unit)
            if unit:GetAIBrain().Score.WarScore >= 300 then
                return 250
            else
                return 0
            end
        end,
    },
    CBuildingHealth02 = {
        PriorityFunction = function(unit)
            local priority = 0
            if GetCitadelHealth(unit) < 0.98 then
                priority = 600
            elseif EnemyHasUpgrade(unit, 'CTroopNumber04') then
                priority = 100
            elseif unit:GetAIBrain().Score.WarRank >= 7 or GetEnemyTeamBrain(unit).Score.WarRank >= 7 then
                priority = 10
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingHealth02'))
        end,
    },
    CBuildingHealth03 = {
        PriorityFunction = function(unit)
            local priority = 0
            if GetCitadelHealth(unit) < 0.9 then
                priority = 600
            elseif EnemyHasUpgrade(unit, 'CTroopNumber04') then
                priority = 15
            elseif unit:GetAIBrain().Score.WarRank >= 8 or GetEnemyTeamBrain(unit).Score.WarRank >= 8 then
                priority = 5
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingHealth03'))
        end,
    },
    CBuildingHealth04 = {
        PriorityFunction = function(unit)
            local priority = 0
            if GetCitadelHealth(unit) < 0.8 then
                priority = 600
            elseif EnemyHasUpgrade(unit, 'CTroopNumber04') then
                priority = 15
            elseif unit:GetAIBrain().Score.WarRank >= 10 or GetEnemyTeamBrain(unit).Score.WarRank >= 10 then
                priority = 5
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingHealth04'))
        end,
    },

    # Building Firepower - prioritize by enemy giants, then by structure upgrades, then war rank
    CBuildingStrength01 = {
        PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CTroopNumber06') then
                priority = 250
            elseif TeamHasUpgrade(unit, 'CUpgradeStructure01') then
                priority = 10
            elseif unit:GetAIBrain().Score.WarRank >= 7 or GetEnemyTeamBrain(unit).Score.WarRank >= 7 then
                priority = 15
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingStrength01'))
        end,
    },
    CBuildingStrength02 = {
        PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CTroopNumber06') then
                priority = 225
            elseif TeamHasUpgrade(unit, 'CUpgradeStructure01') then
                priority = 10
            elseif unit:GetAIBrain().Score.WarRank >= 8 or GetEnemyTeamBrain(unit).Score.WarRank >= 8 then
                priority = 10
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingStrength02'))
        end,
    },
    CBuildingStrength03 = {
        PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CTroopNumber06') then
                priority = 200
            elseif TeamHasUpgrade(unit, 'CUpgradeStructure01') then
                priority = 10
            elseif unit:GetAIBrain().Score.WarRank >= 9 or GetEnemyTeamBrain(unit).Score.WarRank >= 9 then
                priority = 5
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingStrength03'))
        end,
    },
    CBuildingStrength04 = {
        PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CTroopNumber06') then
                priority = 175
            elseif TeamHasUpgrade(unit, 'CUpgradeStructure01') then
                priority = 10
            elseif unit:GetAIBrain().Score.WarRank >= 10 or GetEnemyTeamBrain(unit).Score.WarRank >= 10 then
                priority = 5
            end
            return priority * GoldThreshold(unit, UpgradeCost('CBuildingStrength04'))
        end,
    },

    # Currency - timed priorities
    CGoldIncome01 = {
        Priority = 200,
    },
    CGoldIncome02 = { --12 / 18 minute cutoff
        PriorityFunction = function(unit)
            local mult = math.floor(GetMinute() * 8)
            local base = 100 * GetGoldMult()
            local priority = math.max(0, base - mult)
            --LOG("mithy: CitadelUpgradeWeights: CGoldIncome02: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": "..repr(base).." - "..repr(mult).." = "..repr(priority))
            return priority
        end,
    },
    CGoldIncome03 = { --10 / 15 minute cutoff
        PriorityFunction = function(unit)
            local mult = math.floor(GetMinute() * 6)
            local base = 60 * GetGoldMult()
            local priority = math.max(0, base - mult)
            --LOG("mithy: CitadelUpgradeWeights: CGoldIncome03: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": "..repr(base).." - "..repr(mult).." = "..repr(priority))
            return priority
        end,
    },

    # Blacksmith - match enemy upgrades, ultra-prioritize if either team has giants
    CTroopStrength01 = {
       PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 450
            elseif TeamHasUpgrade(unit, 'CTroopNumber03') or EnemyHasUpgrade(unit, 'CTroopStrength01') then
                return 25
            elseif unit:GetAIBrain().Score.WarRank >= 7 then
                return 10
            else
                return 0
            end
        end,
    },
    CTroopStrength02 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 425
            elseif TeamHasUpgrade(unit, 'CTroopNumber05') or EnemyHasUpgrade(unit, 'CTroopStrength02') then
                return 15
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 5
            else
                return 0
            end
        end,
    },
    CTroopStrength03 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 400
            elseif TeamHasUpgrade(unit, 'CTroopNumber04') or EnemyHasUpgrade(unit, 'CTroopStrength03') then
                return 15
            elseif unit:GetAIBrain().Score.WarRank >= 9 then
                return 5
            else
                return 0
            end
        end,
    },
    CTroopStrength04 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 375
            elseif EnemyHasUpgrade(unit, 'CTroopStrength04') then
                return 10
            elseif unit:GetAIBrain().Score.WarRank >= 10 then
                return 5
            else
                return 0
            end
        end,
    },

    # Armory - match enemy upgrades, ultra-prioritize if either team has giants
    CTroopArmor01 = {
       PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 475
            elseif TeamHasUpgrade(unit, 'CTroopNumber03') or EnemyHasUpgrade(unit, 'CTroopArmor01') then
                return 25
            elseif unit:GetAIBrain().Score.WarRank >= 7 then
                return 15
            else
                return 0
            end
        end,
    },
    CTroopArmor02 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 450
            elseif TeamHasUpgrade(unit, 'CTroopNumber05') or EnemyHasUpgrade(unit, 'CTroopArmor02') then
                return 20
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 5
            else
                return 0
            end
        end,
    },
    CTroopArmor03 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 425
            elseif TeamHasUpgrade(unit, 'CTroopNumber04') or EnemyHasUpgrade(unit, 'CTroopArmor03') then
                return 15
            elseif unit:GetAIBrain().Score.WarRank >= 9 then
                return 5
            else
                return 0
            end
        end,
    },
    CTroopArmor04 = {
        PriorityFunction = function(unit)
            if TeamHasUpgrade(unit, 'CTroopNumber06') or EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 400
            elseif EnemyHasUpgrade(unit, 'CTroopArmor04') then
                return 10
            elseif unit:GetAIBrain().Score.WarRank >= 10 then
                return 5
            else
                return 0
            end
        end,
    },

    # Experience - always get I, priorities for II-IV based on average team level, xp game setting, and time elapsed.
    CPortalFrequency01 = {
        Priority = 200,
    },
    CPortalFrequency02 = {
        PriorityFunction = function(unit)
            local mult = math.floor( (GetMinute() + GetAverageLevel(unit)) ) * 3
            local base = math.floor( 90 / GetXPMult() )
            local priority = math.max(0, base - mult)
            --LOG("mithy: CitadelUpgradeWeights: CPortalFrequency02: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": "..repr(base).." - "..repr(mult).." = "..repr(priority))
            return priority
        end,
    },
    CPortalFrequency03 = {
        PriorityFunction = function(unit)
            local mult = math.floor( (GetMinute() + GetAverageLevel(unit)) ) * 3
            local base = math.floor( 75 / GetXPMult() )
            local priority = math.max(0, base - mult)
            --LOG("mithy: CitadelUpgradeWeights: CPortalFrequency03: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": "..repr(base).." - "..repr(mult).." = "..repr(priority))
            return priority
        end,
    },
    CPortalFrequency04 = {
        PriorityFunction = function(unit)
            local mult = math.floor( (GetMinute() + GetAverageLevel(unit)) ) * 3
            local base = math.floor( 60 / GetXPMult() )
            local priority = math.max(0, base - mult)
            --LOG("mithy: CitadelUpgradeWeights: CPortalFrequency04: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": "..repr(base).." - "..repr(mult).." = "..repr(priority))
            return priority
        end,
    },

    # Graveyard - priority based on death penalty, number of team deaths, and average team level.
    CDeathPenalty01 = {
        PriorityFunction = function(unit)
            local mult = math.floor( (2*GetTeamDeaths(unit) + GetAverageLevel(unit)) ) * 4
            local minimum = math.floor( 80 / GetDeathMult() )
            local priority = math.max(0, mult - minimum)
            --LOG("mithy: CitadelUpgradeWeights: CDeathPenalty01: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": " ..repr(mult).." - "..repr(minimum).." = "..repr(priority))
            return priority
        end,
    },
    CDeathPenalty02 = {
        PriorityFunction = function(unit)
            local mult = math.floor( (2*GetTeamDeaths(unit) + GetAverageLevel(unit)) ) * 4
            local minimum = math.floor( 160 / GetDeathMult() )
            local priority = math.max(0, mult - minimum)
            --LOG("mithy: CitadelUpgradeWeights: CDeathPenalty02: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": " ..repr(mult).." - "..repr(minimum).." = "..repr(priority))
            return priority
        end,
    },
    CDeathPenalty03 = {
        PriorityFunction = function(unit)
            if GetDeathMult() >= 1 then
                local mult = math.floor( (2*GetTeamDeaths(unit) + GetAverageLevel(unit)) ) * 4
                local minimum = math.floor( 240 / GetDeathMult() )
                local priority = math.max(0, mult - minimum)
                --LOG("mithy: CitadelUpgradeWeights: CDeathPenalty03: "..repr(unit:GetAIBrain():GetTeamArmy().Name)..": " ..repr(mult).." - "..repr(minimum).." = "..repr(priority))
                return priority
            else
                return 0
            end
        end,
    },

    # Additional troops - AI will now attempt to match troop upgrades with their opponent even before WR 8
    # This keeps AI games from becoming too one-sided once one team gets priests+
    # Priests
    CTroopNumber03 = {
         PriorityFunction = function(unit)
            if EnemyHasUpgrade(unit, 'CTroopNumber03') then
                return 500
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 90
            else
                return 0
            end
        end,
    },
    # Angels
    CTroopNumber05 = {
         PriorityFunction = function(unit)
            if EnemyHasUpgrade(unit, 'CTroopNumber05') then
                return 500
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 80
            else
                return 0
            end
        end,
    },
    # Catapults
    CTroopNumber04 = {
         PriorityFunction = function(unit)
            if EnemyHasUpgrade(unit, 'CTroopNumber04') then
                return 500
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 70
            else
                return 0
            end
        end,
    },
    # Giants
    CTroopNumber06 = {
         PriorityFunction = function(unit)
            if EnemyHasUpgrade(unit, 'CTroopNumber06') then
                return 500
            elseif unit:GetAIBrain().Score.WarRank >= 8 then
                return 50
            else
                return 0
            end
        end,
    },
    # Trebuchets - match, prioritize vs catapults+
    CUpgradeStructure01 = {
         PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CUpgradeStructure01') or EnemyHasUpgrade(unit, 'CTroopNumber04') then
                priority = 100
            elseif GetEnemyTeamBrain(unit).Score.WarRank >= 8 then
                priority = 10
            end
            return priority * GoldThreshold(unit, UpgradeCost('CUpgradeStructure01') * 2)
        end,
    },
    # Finger of God - match, prioritize vs catapults+
    CUpgradeStructure02 = {
         PriorityFunction = function(unit)
            local priority = 0
            if EnemyHasUpgrade(unit, 'CUpgradeStructure02') or EnemyHasUpgrade(unit, 'CTroopNumber04') then
                priority = 100
            elseif GetEnemyTeamBrain(unit).Score.WarRank >= 9 then
                priority = 10
            end
            return priority * GoldThreshold(unit, UpgradeCost('CUpgradeStructure02') * 2)
        end,
    },
}