## UnitUtil
##
## Utilities for dealing with Civ4 Units and their related objects.
##
## Copyright (c) 2008 The BUG Mod.
##
## Author: EmperorFool

from CvPythonExtensions import *
import BugUtil
import PlayerUtil

# BUG - Mac Support - start
BugUtil.fixSets(globals())
# BUG - Mac Support - end

gc = CyGlobalContext()

NUM_UNITS = 0
NUM_CLASSES = 0
NUM_OR_BONUSES = 0
NUM_AND_TECHS = 0

TRAIN_CAN = 0
TRAIN_CANNOT = 1
TRAIN_LACK_BONUS = 2

# units that don't require any resources to build
unitsWithoutBonuses = set()
# units that require at least one resource to build
unitsWithBonuses = set()
# bonuses that are required to build at least one unit
strategicBonuses = set()
# unit ID -> set of tech IDs
unitTechs = dict()

# unit ID -> tuple ( set of units which can upgrade to it, 
#                    set of units to which it can be upgraded )
# e.g. Bowman -> ( (Archer),
#                  (Longbowman, Crossbowman, Rifleman, ...) )
genericUnits = set()
genericUnitIDs = dict()

upgradeUnits = dict()
genericUpgradeUnits = dict()

olderUnits = dict()
newerUnits = dict()

def init():
	"""
	Segregates units into two sets: those that require resources and those that don't.
	Creates a map of units from ID to the set of tech prerequisites.
	"""
	global NUM_UNITS, NUM_CLASSES
	NUM_UNITS = gc.getNumUnitInfos()
	NUM_CLASSES = gc.getNumUnitClassInfos()
	global NUM_OR_BONUSES, NUM_AND_TECHS
	NUM_OR_BONUSES = gc.getNUM_UNIT_PREREQ_OR_BONUSES()
	NUM_AND_TECHS = gc.getNUM_UNIT_AND_TECH_PREREQS()
	for eUnit in range(NUM_UNITS):
		unitInfo = gc.getUnitInfo(eUnit)
		BugUtil.debug("==== %s ====", unitInfo.getDescription())
		
		# generic unit
		classInfo = gc.getUnitClassInfo(unitInfo.getUnitClassType())
		eGenericUnit = classInfo.getDefaultUnitIndex()
		genericUnitIDs[eUnit] = eGenericUnit
		if eUnit == eGenericUnit:
			genericUnits.add(eUnit)
		else:
			BugUtil.debug("  unique of %s",
					gc.getUnitInfo(eGenericUnit).getDescription())
		
		# resource sets
		found = False
		eBonus = unitInfo.getPrereqAndBonus()
		if eBonus != -1:
			found = True
			strategicBonuses.add(eBonus)
			BugUtil.debug("  requires %s", gc.getBonusInfo(eBonus).getDescription())
		for i in range(NUM_OR_BONUSES):
			eBonus = unitInfo.getPrereqOrBonuses(i)
			if eBonus != -1:
				found = True
				strategicBonuses.add(eBonus)
				BugUtil.debug("  requires %s", gc.getBonusInfo(eBonus).getDescription())
		if found:
			unitsWithBonuses.add(eUnit)
		else:
			unitsWithoutBonuses.add(eUnit)
		
		# tech map
		techs = set()
		unitTechs[eUnit] = techs
		eTech = unitInfo.getPrereqAndTech()
		if eTech != -1:
			techs.add(eTech)
		for i in range(NUM_AND_TECHS):
			eTech = unitInfo.getPrereqAndTechs(i)
			if eTech != -1:
				techs.add(eTech)
		for eTech in techs:
			BugUtil.debug("  requires %s", gc.getTechInfo(eTech).getDescription())
		
	# upgrade maps
	for eUnit in range(NUM_UNITS):
		getOlderUnits(eUnit)
		getNewerUnits(eUnit)
	
	initOrders()

def unitInfos():
	"""Iterates through all CvUnitInfos."""
	for eUnit in range(NUM_UNITS):
		yield gc.getUnitInfo(eUnit)

def unitClassInfos():
	"""Iterates through all CvUnitClassInfos."""
	for eClass in range(NUM_CLASSES):
		yield gc.getUnitClassInfo(eClass)

def isGeneric(eUnit):
	"""Returns True if the given unit is generic and valid."""
	return eUnit in genericUnits

def isUnique(eUnit):
	"""Returns True if the given unit is unique and valid."""
	return eUnit not in genericUnits and eUnit != -1

def getGeneric(eUnit):
	"""
	Returns the generic unit counterpart to the given unique unit or the same
	unit if it's not unique.
	"""
	return genericUnitIDs[eUnit]

def getGenerics(units):
	"""
	Returns a set of the generic units which are counterparts to the unique units
	in the given set. Generic units in the set are not returned.
	"""
	generics = set()
	for eUnit in units:
		eGenericUnit = genericUnitIDs[eUnit]
		if eGenericUnit != eUnit:
			generics.add(eGenericUnit)
	return generics

def getGenericUpgrades(eUnit):
	"""Returns the set of all generic units to which eUnit can upgrade."""
	return genericUpgradeUnits[eUnit]

def getUpgrades(eUnit):
	"""Returns the set of all units to which eUnit can upgrade."""
	return upgradeUnits[eUnit]

def getOlderUnits(eUnit):
	"""Returns the set of all units from which eUnit can upgrade."""
	if eUnit in olderUnits:
		return olderUnits[eUnit]
	unitInfo = gc.getUnitInfo(eUnit)
	eClass = unitInfo.getUnitClassType()
	units = set()
	olderUnits[eUnit] = units
	for eOldUnit in range(NUM_UNITS):
		oldUnitInfo = gc.getUnitInfo(eOldUnit)
		if oldUnitInfo.getUpgradeUnitClass(eClass):
			#BugUtil.debug("%s -> %s", oldUnitInfo.getDescription(), unitInfo.getDescription())
			units.add(eOldUnit)
			units |= getOlderUnits(eOldUnit)
	return units

def getNewerUnits(eUnit):
	"""Returns the set of all units to which eUnit can upgrade."""
	if eUnit in newerUnits:
		return newerUnits[eUnit]
	unitInfo = gc.getUnitInfo(eUnit)
	upgrades = set()
	upgradeUnits[eUnit] = upgrades
	genericUpgrades = set()
	genericUpgradeUnits[eUnit] = genericUpgrades
	newer = set()
	newerUnits[eUnit] = newer
	for eNewUnit in range(NUM_UNITS):
		newUnitInfo = gc.getUnitInfo(eNewUnit)
		if unitInfo.getUpgradeUnitClass(newUnitInfo.getUnitClassType()):
			#BugUtil.debug("%s -> %s", unitInfo.getDescription(), newUnitInfo.getDescription())
			upgrades.add(eNewUnit)
			if isGeneric(eNewUnit):
				genericUpgrades.add(eNewUnit)
			newer.add(eNewUnit)
			newer |= getNewerUnits(eNewUnit)
	return newer

def isUnitOrUpgradeInSet(eUnit, units):
	"""
	Returns True if eUnit is in the given set of units or can be upgraded
	to at least one unit in it.
	"""
	return eUnit in units or len(getNewerUnits(eUnit) & units) > 0

def areUpgradesInSet(eUnit, units):
	"""
	Returns True if every immediate upgrade of eUnit is in the given set.
	
	This ignores transitive upgrades (Warrior doesn't check for Macemen).
	
	Need to take UUs into consideration.
	"""
	upgrades = getUpgrades(eUnit)
	return upgrades <= units

def replaceUniqueUnits(units):
	"""
	Replaces unique units with their generic counterparts in the given set.
	"""
	uniques = set()
	generics = set()
	for eUnit in units:
		if isUnique(eUnit):
			uniques.add(eUnit)
			generics.add(getGeneric(eUnit))
	units -= uniques
	units += generics

def findObsoleteUnits(units):
	"""
	Returns a set containing the units whose immediate upgrades are all in the set,
	taking unique units into consideration.
	
	For example, if the set contains Maceman and Redcoat, neither is returned,
	but if it contains Grenadier as well, Maceman is returned.
	"""
	#result = units.copy()
	generics = getGenerics(units)
	obsoletes = set()
	for eUnit in units:
#		unitInfo = gc.getUnitInfo(eUnit)
		upgrades = getGenericUpgrades(eUnit)
		if upgrades:
			for eUpgradeUnit in upgrades:
				if eUpgradeUnit not in units and eUpgradeUnit not in generics:
#					BugUtil.debug("findObsoleteUnits - %s, cannot build %s", 
#							unitInfo.getDescription(),
#							gc.getUnitInfo(eUpgradeUnit).getDescription())
					break
			else:
				obsoletes.add(eUnit)
	return obsoletes

def getKnownUnits(playerOrID):
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	units = set()
	for eUnit in range(NUM_UNITS):
		for eTech in unitTechs[eUnit]:
			if team.isHasTech(eTech):
				units.add(eUnit)
	return units

def getKnowableUnits(playerOrID):
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	units = set()
	for eUnit in range(NUM_UNITS):
		for eTech in unitTechs[eUnit]:
			if not (team.isHasTech(eTech) or player.canResearch(eTech, False)):
				break
		else:
			units.add(eUnit)
	return units

def getTrainableUnits(playerOrID, knowableUnits, checkCities=True, military=None):
	"""
	Returns the set of all units the player can train, including obsolete ones.
	"""
	game = CyGame()
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	civInfo = gc.getCivilizationInfo(player.getCivilizationType())
	if checkCities:
		cities = PlayerUtil.getPlayerCities(player)
	else:
		cities = None
	units = set()
	BugUtil.debug("%s =========", player.getName())
	for eClass in range(NUM_CLASSES):
		eUnit = civInfo.getCivilizationUnits(eClass)
		if eUnit == -1 or eUnit not in knowableUnits:
			#BugUtil.debug("  %s -> unknowable", gc.getUnitClassInfo(eClass).getDescription())
			continue
		unitInfo = gc.getUnitInfo(eUnit)
		# military
		if military is not None:
			combat = (unitInfo.getUnitCombatType() > 0 or unitInfo.getNukeRange() != -1
					or unitInfo.getAirCombat() > 0)
			if military != combat:
				#BugUtil.debug("  %s -> combat is %s", unitInfo.getDescription(), combat)
				continue
		# OCC and Settlers
		if game.isOption(GameOptionTypes.GAMEOPTION_ONE_CITY_CHALLENGE) and unitInfo.isFound():
			BugUtil.debug("  %s -> no founding units in OCC", unitInfo.getDescription())
			continue
		# techs
		for eTech in unitTechs[eUnit]:
			if not team.isHasTech(eTech):
				BugUtil.debug("  %s -> doesn't know %s", unitInfo.getDescription(), 
						gc.getTechInfo(eTech).getDescription())
				missing = True
				break
		else:
			missing = False
		if missing:
			continue
		# state religion
		eReligion = unitInfo.getStateReligion()
		if eReligion != -1 and player.getStateReligion() != eReligion:
			BugUtil.debug("  %s -> wrong state religion", unitInfo.getDescription())
			continue
		# nukes
		if (game.isNoNukes() or not game.isNukesValid()) and unitInfo.getNukeRange() != -1:
			BugUtil.debug("  %s -> no nukes", unitInfo.getDescription())
			continue
		# getSpecialUnitType, game.isSpecialUnitValid
		eSpecialType = unitInfo.getSpecialUnitType()
		if eSpecialType != -1 and not game.isSpecialUnitValid(eSpecialType):
			BugUtil.debug("  %s -> special unit type %s invalid", unitInfo.getDescription(),
					gc.getSpecialUnitInfo(eSpecialType).getDescription())
			continue
		# cities
		if cities and not canAnyCityBuildUnit(eUnit, cities, -1, True):
			BugUtil.debug("  %s -> no city can train unit", unitInfo.getDescription())
			continue
		BugUtil.debug("  %s", unitInfo.getDescription())
		units.add(eUnit)
	return units

def getTrainableAndUntrainableUnits(playerOrID, knowableUnits, military=None):
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	cities = PlayerUtil.getPlayerCities(player)
	# separate units into two groups: yes and no
	units = getTrainableUnits(playerOrID, knowableUnits, False, military)
	yesUnits = set()
	noUnits = set()
	BugUtil.debug("-----------------------")
	for eUnit in units:
		if canAnyCityBuildUnit(eUnit, cities, -1, True):
			BugUtil.debug("  yes %s", gc.getUnitInfo(eUnit).getDescription())
			yesUnits.add(eUnit)
		else:
			BugUtil.debug("  no  %s", gc.getUnitInfo(eUnit).getDescription())
			noUnits.add(eUnit)
	return yesUnits, noUnits

def getKnownTrainableUnits(playerOrID, askingPlayerOrID, knowableUnits, bonuses, military=None):
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	askingPlayer = PlayerUtil.getPlayer(askingPlayerOrID)
	eAskingTeam, askingTeam = PlayerUtil.getPlayerTeamAndID(askingPlayer)
	#trade = player.canTradeNetworkWith(askingPlayer.getID())
	cities = PlayerUtil.getPlayerCities(player, 
			lambda city: city.isRevealed(eAskingTeam, False))
	# separate units into two groups: yes and maybe
	units = getTrainableUnits(playerOrID, knowableUnits, False, military)
	yesUnits = set()
	maybeUnits = set()
	BugUtil.debug("-----------------------")
	for eUnit in units:
		if not canAnyCityBuildUnit(eUnit, cities, eAskingTeam, False):
			BugUtil.debug("  no    %s", gc.getUnitInfo(eUnit).getDescription())
		elif hasBonusesForUnit(eUnit, bonuses):
			BugUtil.debug("  yes   %s", gc.getUnitInfo(eUnit).getDescription())
			yesUnits.add(eUnit)
		elif bonuses is None:
			BugUtil.debug("  maybe %s", gc.getUnitInfo(eUnit).getDescription())
			maybeUnits.add(eUnit)
	return yesUnits, maybeUnits

def hasBonusesForUnit(eUnit, bonuses):
	if eUnit not in unitsWithBonuses:
		return True
	if not bonuses:
		return False
	unitInfo = gc.getUnitInfo(eUnit)
	eBonus = unitInfo.getPrereqAndBonus()
	if eBonus != -1 and eBonus not in bonuses:
		return False
	requiresBonus = False
	for i in range(NUM_OR_BONUSES):
		eBonus = unitInfo.getPrereqOrBonuses(i)
		if eBonus != -1:
			requiresBonus = True
			if eBonus in bonuses:
				break
	else:
		if requiresBonus:
			return False
	return True

def canAnyCityBuildUnit(eUnit, cities=None, askingTeamOrID=-1, checkBonuses=True):
	eAskingTeam = PlayerUtil.getTeamID(askingTeamOrID)
	unitInfo = gc.getUnitInfo(eUnit)
	if cities:
		for city in cities:
			if canCityBuildUnit(unitInfo, city, eAskingTeam, checkBonuses):
				return True
		return False
	else:
		return canCityBuildUnit(unitInfo, None, eAskingTeam, checkBonuses)

def canCityBuildUnit(unitInfo, city, eAskingTeam, checkBonuses=True):
	# religion
	if unitInfo.isPrereqReligion():
		if not city or city.getReligionCount() > 0:
			# EF: Seems odd to enforce NO religions in the city, 
			#     but this is how CvPlot.canTrain() does it.
			#     The function should actually be called isPrereqNoReligion().
			return False
	eReligion = unitInfo.getPrereqReligion()
	if eReligion != -1 and not (city and city.isHasReligion(eReligion)):
		return False
	# corporation
	eCorp = unitInfo.getPrereqCorporation()
	if eCorp != -1 and not (city and city.isActiveCorporation(eCorp)):
		return False
	# skipping isPrereqBonuses as the land part looks broken
	# and we don't want to limit work boats if all resources are covered
	# domain
	if unitInfo.getDomainType() == DomainTypes.DOMAIN_SEA:
		if not (city and PlayerUtil.isSaltWaterPort(city, eAskingTeam)):
			return False
		# EF: this is how CyPlot does it
		#plot = city.plot()
		#if not plot.isWater() or not plot.isCoastalLand(unitInfo.getMinAreaSize()):
		#	continue
	else:
		minArea = unitInfo.getMinAreaSize()
		if minArea != -1:
			if eAskingTeam != -1 or not city or city.plot().area().getNumTiles() < minArea:
				return False
	# holy city
	eReligion = unitInfo.getHolyCity()
	if eReligion != -1 and not (city and city.isHolyCity(eReligion)):
		return False
	# building
	eBuilding = unitInfo.getPrereqBuilding()
	if eBuilding != -1:
		if eAskingTeam != -1 or not city:
			return False
		if city.getNumBuilding(eBuilding) == 0:
			eSpecialBuilding = gc.getBuildingInfo(eBuilding).getSpecialBuildingType()
			if eSpecialBuilding == -1 or not gc.getPlayer(city.getOwner()).isSpecialBuildingNotRequired(eSpecialBuilding):
				return False
	# resources
	if checkBonuses and not cityHasBonusesForUnit(unitInfo, city):
		return False
	# passes all tests for this city
	return True

def cityHasBonusesForUnit(unitInfo, city):
	if not city:
		return False
	eBonus = unitInfo.getPrereqAndBonus()
	if eBonus != -1 and not city.hasBonus(eBonus):
		return False
	requiresBonus = False
	for i in range(NUM_OR_BONUSES):
		eBonus = unitInfo.getPrereqOrBonuses(i)
		if eBonus != -1:
			requiresBonus = True
			if city.hasBonus(eBonus):
				break
	else:
		if requiresBonus:
			return False
	return True

def getCanTrainUnits(playerOrID, askingPlayerOrID=None, military=None):
	"""
	Returns the set of all units the player can train.
	
	Searches all of the player's cities to find which units can be trained.
	
	If askingPlayerOrID is given, only cities they have seen are checked, and
	only units whose prerequisite techs they know or can research are returned.
	Also, if the two players' trade networks are not connected, units that
	require resources to train are returned in a second set.
	
	If military is provided, only military or civilian units are checked
	depending on its value, True or False, respectively.
	
	*** OBSOLETE ***
	
	"""
	player, team = PlayerUtil.getPlayerAndTeam(playerOrID)
	askingPlayer = PlayerUtil.getPlayer(askingPlayerOrID)
	if askingPlayer:
		eAskingTeam, askingTeam = PlayerUtil.getPlayerTeamAndID(askingPlayer)
		trade = player.canTradeNetworkWith(askingPlayer.getID())
	civInfo = gc.getCivilizationInfo(player.getCivilizationType())
	units = set()
	maybeUnits = set()
	for eClass in range(NUM_CLASSES):
		eUnit = civInfo.getCivilizationUnits(eClass)
		if eUnit == -1:
			classInfo = gc.getUnitClassInfo(eClass)
			BugUtil.debug("%s doesn't have %s", civInfo.getDescription(), classInfo.getDescription())
			eUnit = classInfo.getDefaultUnitIndex()
		unitInfo = gc.getUnitInfo(eUnit)
		if unitInfo:
			if ((military == True and unitInfo.getUnitCombatType() <= 0)
			or (military == False and unitInfo.getUnitCombatType() > 0)):
				BugUtil.debug("skipping (non-)military %s", unitInfo.getDescription())
				continue
			if askingPlayer:
				for eTech in unitTechs[eUnit]:
					if not (askingTeam.isHasTech(eTech) or askingPlayer.canResearch(eTech, False)):
						BugUtil.debug("%s doesn't comprehend %s", askingPlayer.getName(), gc.getTechInfo(eTech).getDescription())
						skip = True
						break
				else:
					skip = False
				if skip:
					BugUtil.debug("skipping unknowable %s", unitInfo.getDescription())
					continue
			for city in PlayerUtil.playerCities(player):
				if askingPlayer:
					if not city.isRevealed(eAskingTeam, False):
						continue
					if city.canTrain(eUnit, False, not trade):
						if eUnit in unitsWithBonuses:
							maybeUnits.add(eUnit)
						else:
							units.add(eUnit)
						break
				else:
					if city.canTrain(eUnit, False, False):
						units.add(eUnit)
						break
	BugUtil.debug("%s can train:", player.getName())
	for eUnit in units:
		unitInfo = gc.getUnitInfo(eUnit)
		BugUtil.debug("  %s", unitInfo.getDescription())
	if askingPlayer:
		BugUtil.debug("%s can maybe train:", player.getName())
		for eUnit in maybeUnits:
			unitInfo = gc.getUnitInfo(eUnit)
			BugUtil.debug("  %s", unitInfo.getDescription())
		return units, maybeUnits
	else:
		return units

(
	ORDER_NONE,
	
	ORDER_SKIP,
	ORDER_SLEEP,
	ORDER_FORTIFY,
	ORDER_HEAL,
	
	ORDER_SENTRY,
	ORDER_INTERCEPT,
	ORDER_PATROL,
	ORDER_PLUNDER,
	
	ORDER_BUILD,  # improvement
	ORDER_CONSTRUCT,  # building
	ORDER_GOTO,
	ORDER_EXPLORE,
	
	ORDER_AUTO_BUILD,
	ORDER_AUTO_NETWORK,
	ORDER_AUTO_CITY,
	ORDER_AUTO_RELIGION,
) = range(17)

ORDERS_BY_ACTIVITY = {
	ActivityTypes.ACTIVITY_AWAKE: ORDER_NONE,
	ActivityTypes.ACTIVITY_INTERCEPT: ORDER_INTERCEPT,
	ActivityTypes.ACTIVITY_PATROL: ORDER_PATROL,
	ActivityTypes.ACTIVITY_PLUNDER: ORDER_PLUNDER,
	ActivityTypes.ACTIVITY_HEAL: ORDER_HEAL,
	ActivityTypes.ACTIVITY_SENTRY: ORDER_SENTRY,
	ActivityTypes.ACTIVITY_HOLD: ORDER_SKIP,
}
ORDERS_BY_AUTOMATION = {
	AutomateTypes.AUTOMATE_EXPLORE: ORDER_EXPLORE,
	AutomateTypes.AUTOMATE_BUILD: ORDER_AUTO_BUILD,
	AutomateTypes.AUTOMATE_NETWORK: ORDER_AUTO_NETWORK,
	AutomateTypes.AUTOMATE_CITY: ORDER_AUTO_CITY,
	AutomateTypes.AUTOMATE_RELIGION: ORDER_AUTO_RELIGION,
}
MOVE_TO_MISSIONS = [
	MissionTypes.MISSION_MOVE_TO, 
	MissionTypes.MISSION_MOVE_TO_UNIT,
]

def getOrder(unit):
	group = unit.getGroup()
	eActivityType = group.getActivityType()
	if eActivityType in ORDERS_BY_ACTIVITY:
		return ORDERS_BY_ACTIVITY[eActivityType]
	eAutomationType = group.getAutomateType()
	if eAutomationType in ORDERS_BY_AUTOMATION:
		return ORDERS_BY_AUTOMATION[eAutomationType]
	if (group.getLengthMissionQueue() > 0):
		# TODO: loop to find the first non-goto and check in ORDERS_BY_MISSION
		eMissionType = group.getMissionType(0)
		if eMissionType == MissionTypes.MISSION_BUILD:
			return ORDER_BUILD
		elif eMissionType in MOVE_TO_MISSIONS:
			return ORDER_GOTO
	elif (unit.isWaiting()):
		if (unit.isFortifyable()):
			return ORDER_FORTIFY
		else:
			return ORDER_SLEEP
	return ORDER_NONE

def initOrders():
	"""
	Adds orders added by BULL.
	"""
	try:
		ORDERS_BY_ACTIVITY[ActivityTypes.ACTIVITY_SENTRY_WHILE_HEAL] = ORDER_HEAL
		ORDERS_BY_ACTIVITY[ActivityTypes.ACTIVITY_SENTRY_LAND_UNITS] = ORDER_SENTRY
		ORDERS_BY_ACTIVITY[ActivityTypes.ACTIVITY_SENTRY_NAVAL_UNITS] = ORDER_SENTRY
		MOVE_TO_MISSIONS.append(MissionTypes.MISSION_MOVE_TO_SENTRY)
	except:
		pass
