From 8c24d1686b33c8373e8a874b81145e8f148c91db Mon Sep 17 00:00:00 2001
From: Lyberta <lyberta@lyberta.net>
Date: Mon, 30 Jul 2018 20:40:33 +0300
Subject: [PATCH] Teamplay: Third pass at autobalance.

---
 qcsrc/server/teamplay.qc | 82 +++++++++++++++++++++++++++++-----------
 qcsrc/server/teamplay.qh |  7 ++++
 2 files changed, 66 insertions(+), 23 deletions(-)

diff --git a/qcsrc/server/teamplay.qc b/qcsrc/server/teamplay.qc
index cfd15f248..4a4305bac 100644
--- a/qcsrc/server/teamplay.qc
+++ b/qcsrc/server/teamplay.qc
@@ -838,8 +838,6 @@ void TeamBalance_AutoBalanceBots()
 	TeamBalance_GetTeamCounts(balance, NULL);
 	int smallest_team_index = 0;
 	int smallest_team_player_count = 0;
-	int largest_team_index = 0;
-	int largest_team_player_count = 0;
 	for (int i = 1; i <= NUM_TEAMS; ++i)
 	{
 		entity team_ = TeamBalance_GetTeamFromIndex(balance, i);
@@ -858,35 +856,43 @@ void TeamBalance_AutoBalanceBots()
 			smallest_team_index = i;
 			smallest_team_player_count = player_count;
 		}
-		if (largest_team_index == 0)
+	}
+	//PrintToChatAll(sprintf("Smallest team: %f", smallest_team_index));
+	//PrintToChatAll(sprintf("Smallest team players: %f", smallest_team_player_count));
+	entity lowest_bot = NULL;
+	int teams = BITS(NUM_TEAMS);
+	while (teams != 0)
+	{
+		int largest_team_index = TeamBalance_GetLargestTeamIndex(balance,
+			teams);
+		if (smallest_team_index == largest_team_index)
 		{
-			largest_team_index = i;
-			largest_team_player_count = player_count;
+			TeamBalance_Destroy(balance);
+			return;
 		}
-		else if (player_count > largest_team_player_count)
+		entity largest_team = TeamBalance_GetTeamFromIndex(balance,
+			largest_team_index);
+		int largest_team_player_count = TeamBalanceTeam_GetNumberOfPlayers(
+			largest_team);
+		if (largest_team_player_count - smallest_team_player_count < 2)
 		{
-			largest_team_index = i;
-			largest_team_player_count = player_count;
+			TeamBalance_Destroy(balance);
+			return;
 		}
+		//PrintToChatAll(sprintf("Largest team: %f", largest_team_index));
+		//PrintToChatAll(sprintf("Largest team players: %f", largest_team_player_count));
+		lowest_bot = TeamBalance_GetPlayerForTeamSwitch(largest_team_index,
+			smallest_team_index, true);
+		if (lowest_bot != NULL)
+		{
+			break;
+		}
+		teams &= ~Team_IndexToBit(largest_team_index);
 	}
 	TeamBalance_Destroy(balance);
-	//PrintToChatAll(sprintf("Smallest team: %f", smallest_team_index));
-	//PrintToChatAll(sprintf("Largest team: %f", largest_team_index));
-	//PrintToChatAll(sprintf("Smallest team players: %f", smallest_team_player_count));
-	//PrintToChatAll(sprintf("Largest team players: %f", largest_team_player_count));
-	if (smallest_team_index == largest_team_index)
-	{
-		return;
-	}
-	if (largest_team_player_count - smallest_team_player_count < 2)
-	{
-		return;
-	}
-	entity lowest_bot = TeamBalance_GetPlayerForTeamSwitch(largest_team_index,
-		smallest_team_index, true);
 	if (lowest_bot == NULL)
 	{
-		//PrintToChatAll("No bot found");
+		//PrintToChatAll("No bot found after searching through all the teams");
 		return;
 	}
 	if (!Player_SetTeamIndex(lowest_bot, smallest_team_index))
@@ -896,6 +902,36 @@ void TeamBalance_AutoBalanceBots()
 	KillPlayerForTeamChange(lowest_bot);
 }
 
+int TeamBalance_GetLargestTeamIndex(entity balance, int teams)
+{
+	int largest_team_index = 0;
+	int largest_team_player_count = 0;
+	for (int i = 1; i <= NUM_TEAMS; ++i)
+	{
+		if (!(Team_IndexToBit(i) & teams))
+		{
+			continue;
+		}
+		entity team_ = TeamBalance_GetTeamFromIndex(balance, i);
+		if (!TeamBalanceTeam_IsAllowed(team_))
+		{
+			continue;
+		}
+		int player_count = TeamBalanceTeam_GetNumberOfPlayers(team_);
+		if (largest_team_index == 0)
+		{
+			largest_team_index = i;
+			largest_team_player_count = player_count;
+		}
+		else if (player_count > largest_team_player_count)
+		{
+			largest_team_index = i;
+			largest_team_player_count = player_count;
+		}
+	}
+	return largest_team_index;
+}
+
 entity TeamBalance_GetPlayerForTeamSwitch(int source_team_index,
 	int destination_team_index, bool is_bot)
 {
diff --git a/qcsrc/server/teamplay.qh b/qcsrc/server/teamplay.qh
index 404d728a2..e265df49f 100644
--- a/qcsrc/server/teamplay.qh
+++ b/qcsrc/server/teamplay.qh
@@ -226,6 +226,13 @@ int TeamBalance_CompareTeams(entity balance, int team_index_a, int team_index_b,
 /// \brief Switches a bot from one team to another if teams are not balanced.
 void TeamBalance_AutoBalanceBots();
 
+/// \brief Returns the index of the team with most players that is contained in
+/// the given bitmask of teams.
+/// \param[in] balance Team balance entity.
+/// \param[in] teams Bitmask of teams to search in.
+/// \return Index of the team with most players.
+int TeamBalance_GetLargestTeamIndex(entity balance, int teams);
+
 /// \brief Returns the player who is the most suitable for switching between
 /// the given teams.
 /// \param[in] source_team_index Index of the team to search in.
-- 
2.39.2