## UnitGrouper
##
## Builds groups of units for use in reporting or screens.
##
## Copyright (c) 2008 The BUG Mod.
##
## Author: EmperorFool

from CvPythonExtensions import *
import BugUtil
import UnitUtil

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

# globals
gc = CyGlobalContext()

# Base grouping classes

class Grouper:
	"""
	Holds all Grouping definitions.
	"""
	def __init__(self):
		self.groupings = []
		self.groupingsByKey = {}
	
	def _addGrouping(self, grouping):
		grouping.index = len(self.groupings)
		self.groupings.append(grouping)
		self.groupingsByKey[grouping.key] = grouping
	
	def getGrouping(self, key):
		if key in self.groupingsByKey:
			return self.groupingsByKey[key]
		else:
			return None
	
	def __getitem__(self, key):
		if isinstance(key, int):
			return self.groupings[key]
		else:
			return self.groupingsByKey(key)
	
	def __iter__(self):
		return self.groupings.__iter__()
	
	def iterkeys(self):
		return self.groupingsByKey.iterkeys()
	
	def itervalues(self):
		return self.groupingsByKey.itervalues()
	
	def iteritems(self):
		return self.groupingsByKey.iteritems()

class Grouping:
	"""
	Applies a formula to place units into groups.
	
	key: used for sorting groupings; must be in the range [0, 999] inclusive
	title: used to display the group
	"""
	def __init__(self, key, title):
		self.index = None
		self.key = key
		if title.startswith("TXT_KEY_"):
			self.title = BugUtil.getPlainText(title)
		else:
			self.title = title
		self.groups = {}
	
	def _addGroup(self, group):
		self.groups[group.key] = group
	
	def calcGroupKeys(self, unit, player, team):
		return None

class Group:
	"""
	Represents a single group value within a grouping.
	
	key: used for sorting groups; must be in the range [0, 999] inclusive
	title: used to display the group
	"""
	def __init__(self, grouping, key, title):
		self.grouping = grouping
		self.key = key
		if title.startswith("TXT_KEY_"):
			self.title = BugUtil.getPlainText(title)
		else:
			self.title = title
	
	def getTitle(self):
		return self.title


# Grouping definitions

class UnitTypeGrouping(Grouping):
	"""
	Groups units by their unit type.
	Ex: Warrior, Maceman, Panzer
	"""
	def __init__(self):
		Grouping.__init__(self, "type", "TXT_KEY_UNIT_GROUPER_TYPE_GROUPING")
		
		for i in range(gc.getNumUnitInfos()):
			info = gc.getUnitInfo(i)
			if info:
				self._addGroup(Group(self, i, info.getDescription()))
	
	def calcGroupKeys(self, unit, player, team):
		return (unit.getUnitType(),)

class UnitCombatGrouping(Grouping):
	"""
	Groups units by their combat type.
	Ex: None, Melee, Gunpowder, Naval
	"""
	def __init__(self):
		Grouping.__init__(self, "combat", "TXT_KEY_UNIT_GROUPER_COMBAT_GROUPING")
		self.NONE = 0
		
		self._addGroup(Group(self, self.NONE, "TXT_KEY_UNIT_GROUPER_COMBAT_GROUP_NONE"))
		for i in range(gc.getNumUnitCombatInfos()):
			info = gc.getUnitCombatInfo(i)
			if info:
				self._addGroup(Group(self, i + 1, info.getDescription()))
	
	def calcGroupKeys(self, unit, player, team):
		return (gc.getUnitInfo(unit.getUnitType()).getUnitCombatType() + 1,)

class LevelGrouping(Grouping):
	"""
	Groups units by their level, 1 to MAX_LEVEL (50).
	Units over level MAX_LEVEL are put into the MAX_LEVEL group.
	"""
	def __init__(self):
		Grouping.__init__(self, "level", "TXT_KEY_UNIT_GROUPER_LEVEL_GROUPING")
		
		self.MAX_LEVEL = 50
		for i in range(self.MAX_LEVEL):
			self._addGroup(Group(self, i, BugUtil.getText("TXT_KEY_UNIT_GROUPER_LEVEL_GROUP", (str(i),))))
		self._addGroup(Group(self, self.MAX_LEVEL, BugUtil.getText("TXT_KEY_UNIT_GROUPER_LEVEL_GROUP", ("%d+" % self.MAX_LEVEL,))))
	
	def calcGroupKeys(self, unit, player, team):
		return (max(0, min(unit.getLevel(), self.MAX_LEVEL)),)

class PromotionGrouping(Grouping):
	"""
	Groups units by their promotions.
	Ex: Combat 1, Cover, Tactics
	"""
	def __init__(self):
		Grouping.__init__(self, "promo", "TXT_KEY_UNIT_GROUPER_PROMOTION_GROUPING")
		
		self.NONE = 0
		self.NO_PROMOS = (0,)
		self._addGroup(Group(self, self.NONE, "TXT_KEY_UNIT_GROUPER_PROMOTION_GROUP_NONE"))
		for i in range(gc.getNumPromotionInfos()):
			info = gc.getPromotionInfo(i)
			if info:
				self._addGroup(Group(self, i + 1, '<img=%s size=16></img> %s' % 
												  (info.getButton(), info.getDescription())))
	
	def calcGroupKeys(self, unit, player, team):
		promos = []
		for iPromo in range(gc.getNumPromotionInfos()):
			if unit.isHasPromotion(iPromo):
				promos.append(iPromo + 1)
		if not promos:
			promos = self.NO_PROMOS
		return promos

class LocationGrouping(Grouping):
	"""
	Groups units by their location on the map.
	Ex: Domestic City, Friendly City, Enemy Territory
	"""
	def __init__(self):
		Grouping.__init__(self, "loc", "TXT_KEY_UNIT_GROUPER_LOCATION_GROUPING")
		(
			self.DOMESTIC_CITY,
			self.DOMESTIC_TERRITORY,
			self.TEAM_CITY,
			self.TEAM_TERRITORY,
			self.FRIENDLY_CITY,
			self.FRIENDLY_TERRITORY,
			self.NEUTRAL_TERRITORY,
			self.ENEMY_TERRITORY,
			self.BARBARIAN_TERRITORY
		) = range(9)
		
		self._addGroup(Group(self, self.DOMESTIC_CITY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_DOMESTIC_CITY"))
		self._addGroup(Group(self, self.DOMESTIC_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_DOMESTIC_TERRITORY"))
		self._addGroup(Group(self, self.TEAM_CITY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_TEAM_CITY"))
		self._addGroup(Group(self, self.TEAM_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_TEAM_TERRITORY"))
		self._addGroup(Group(self, self.FRIENDLY_CITY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_FRIENDLY_CITY"))
		self._addGroup(Group(self, self.FRIENDLY_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_FRIENDLY_TERRITORY"))
		self._addGroup(Group(self, self.NEUTRAL_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_NEUTRAL_TERRITORY"))
		self._addGroup(Group(self, self.ENEMY_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_ENEMY_TERRITORY"))
		self._addGroup(Group(self, self.BARBARIAN_TERRITORY, "TXT_KEY_UNIT_GROUPER_LOCATION_GROUP_BARBARIAN_TERRITORY"))
	
	def calcGroupKeys(self, unit, player, team):
		plot = unit.plot()
		if not plot or plot.isNone():
			return None
		if plot.isBarbarian():
			return (self.BARBARIAN_TERRITORY,)
		teamId = team.getID()
		ownerId = plot.getRevealedOwner(teamId, False)
		if ownerId == -1:
			return (self.NEUTRAL_TERRITORY,)
		elif ownerId == player.getID():
			if plot.isCity():
				return (self.DOMESTIC_CITY,)
			else:
				return (self.DOMESTIC_TERRITORY,)
		else:
			owner = gc.getPlayer(ownerId)
			ownerTeamId = owner.getTeam()
			if ownerTeamId == teamId:
				if plot.isCity():
					return (self.TEAM_CITY,)
				else:
					return (self.TEAM_TERRITORY,)
			elif team.isAtWar(ownerTeamId):
				return (self.ENEMY_TERRITORY,)
			else:
				if plot.isCity():
					return (self.FRIENDLY_CITY,)
				else:
					return (self.FRIENDLY_TERRITORY,)

class OrderGrouping(Grouping):
	"""
	Groups units by their current order/action.
	Ex: Fortify, Go To, Blockade
	"""
	def __init__(self):
		Grouping.__init__(self, "order", "TXT_KEY_UNIT_GROUPER_ORDER_GROUPING")
		(
			self.ORDER_NONE,
			self.ORDER_SKIP,
			self.ORDER_SLEEP,
			self.ORDER_FORTIFY,
			self.ORDER_HEAL,
			self.ORDER_SENTRY,
			self.ORDER_INTERCEPT,
			self.ORDER_PATROL,
			self.ORDER_PLUNDER,
			self.ORDER_BUILD,
			self.ORDER_CONSTRUCT,
			self.ORDER_GOTO,
			self.ORDER_EXPLORE,
			self.ORDER_AUTO_BUILD,
			self.ORDER_AUTO_NETWORK,
			self.ORDER_AUTO_CITY,
			self.ORDER_AUTO_RELIGION,
			self.ORDER_OTHER,
		) = range(18)
		
		self._addGroup(Group(self, self.ORDER_NONE, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_NONE"))
		self._addGroup(Group(self, self.ORDER_SKIP, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_SKIP"))
		self._addGroup(Group(self, self.ORDER_SLEEP, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_SLEEP"))
		self._addGroup(Group(self, self.ORDER_FORTIFY, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_FORTIFY"))
		self._addGroup(Group(self, self.ORDER_HEAL, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_HEAL"))
		self._addGroup(Group(self, self.ORDER_SENTRY, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_SENTRY"))
		self._addGroup(Group(self, self.ORDER_INTERCEPT, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_INTERCEPT"))
		self._addGroup(Group(self, self.ORDER_PATROL, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_PATROL"))
		self._addGroup(Group(self, self.ORDER_PLUNDER, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_PLUNDER"))
		self._addGroup(Group(self, self.ORDER_BUILD, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_BUILD"))
		self._addGroup(Group(self, self.ORDER_CONSTRUCT, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_CONSTRUCT"))
		self._addGroup(Group(self, self.ORDER_GOTO, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_GOTO"))
		self._addGroup(Group(self, self.ORDER_EXPLORE, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_EXPLORE"))
		self._addGroup(Group(self, self.ORDER_AUTO_BUILD, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_AUTO_BUILD"))
		self._addGroup(Group(self, self.ORDER_AUTO_NETWORK, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_AUTO_NETWORK"))
		self._addGroup(Group(self, self.ORDER_AUTO_CITY, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_AUTO_CITY"))
		self._addGroup(Group(self, self.ORDER_AUTO_RELIGION, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_AUTO_RELIGION"))
		self._addGroup(Group(self, self.ORDER_OTHER, "TXT_KEY_UNIT_GROUPER_ORDER_GROUP_OTHER"))
	
	def calcGroupKeys(self, unit, player, team):
		eOrder = UnitUtil.getOrder(unit)
		if eOrder >= self.ORDER_OTHER:
			return (self.ORDER_OTHER,)
		else:
			return (eOrder,)

class StandardGrouper(Grouper):
	def __init__(self):
		Grouper.__init__(self)
		
		self._addGrouping(UnitTypeGrouping())
		self._addGrouping(UnitCombatGrouping())
		self._addGrouping(LevelGrouping())
		self._addGrouping(PromotionGrouping())
		self._addGrouping(LocationGrouping())
		self._addGrouping(OrderGrouping())


# Classes for tracking stats about groups and units

class GrouperStats:
	"""
	Holds stats for a set of groupings.
	"""
	def __init__(self, grouper):
		self.grouper = grouper
		self.groupings = {}

		for grouping in self.grouper.groupings:
			self._addGrouping(GroupingStats(grouping))
	
	def _addGrouping(self, grouping):
		self.groupings[grouping.grouping.key] = grouping
	
	def processUnit(self, player, team, unit):
		stats = UnitStats(unit.getOwner(), unit.getID(), unit)
		for grouping in self.groupings.itervalues():
			grouping._processUnit(player, team, stats)
		return stats
	
	def getGrouping(self, key):
		if key in self.groupings:
			return self.groupings[key]
		else:
			return None
	
	def itergroupings(self):
		return self.groupings.itervalues()

class GroupingStats:
	"""
	Holds stats for a grouping.
	"""
	def __init__(self, grouping):
		self.grouping = grouping
		self.groups = {}
		
		for group in self.grouping.groups.itervalues():
			self._addGroup(GroupStats(group))
	
	def _addGroup(self, group):
		self.groups[group.group.key] = group
	
	def _processUnit(self, player, team, unitStats):
		keys = self.grouping.calcGroupKeys(unitStats.unit, player, team)
		for key in keys:
			self.groups[key]._addUnit(unitStats)
	
	def itergroups(self):
		return self.groups.itervalues()

class GroupStats:
	"""
	Holds stats for a group of units.
	"""
	def __init__(self, group):
		self.group = group
		self.units = set()
	
	def _addUnit(self, unitStats):
		self.units.add(unitStats)
	
	def title(self):
		return self.group.title
	
	def size(self):
		return len(self.units)
	
	def isEmpty(self):
		return self.size() == 0

class UnitStats:
	"""
	Holds stats about a single unit.
	"""
	def __init__(self, playerId, unitId, unit):
		self.key = (playerId, unitId)
		self.unit = unit

	def __hash__(self):
		return hash(self.key)

	def __eq__(self, other):
		return self.key == other.key
