Enhance player swapping logic to maintain team diversity by implementing constraints that ensure at least three teams are represented in the lineup, while also considering the current player's team during replacements.

src/sports/nhl_functions.py

@@ -429,7 +429,6 @@ def init_team_results(model_source: DataFrame, position_reqs: dict, salary_cap:
                     # Add the position to our swapped tracking list
                     swapped_positions.append(pos_to_swap)
 
-                    # ... existing code for finding replacement player ...
                     # Find the assigned position (e.g., 'OF', 'SP', etc.)
                     match = re.match(r'[A-Z]+', pos_to_swap.strip())
                     if match:
@@ -437,6 +436,20 @@ def init_team_results(model_source: DataFrame, position_reqs: dict, salary_cap:
                     else:
                         continue
                     
+                    # Get current player being swapped and their team
+                    current_player = new_row[pos_to_swap]
+                    current_player_team = player_team_map.get(current_player)
+                    
+                    # Count teams in current lineup (before swap)
+                    current_lineup_players = [new_row[pos] for pos in pos_labels]
+                    current_teams = set(player_team_map.get(p) for p in current_lineup_players if player_team_map.get(p))
+                    
+                    # Count how many players from the current player's team are in the lineup
+                    players_from_same_team = sum(1 for p in current_lineup_players if player_team_map.get(p) == current_player_team)
+                    
+                    # If this player is the only one from their team, we need to be careful
+                    must_maintain_team = players_from_same_team == 1 and len(current_teams) <= 3
+                    
                     # Handle special positions
                     if assigned_pos == 'FLEX' or assigned_pos == 'FLEX1' or assigned_pos == 'FLEX2':
                         util_positions = ['C', 'W', 'D']
@@ -461,6 +474,30 @@ def init_team_results(model_source: DataFrame, position_reqs: dict, salary_cap:
                     
                     if eligible.empty:
                         continue
+                    
+                    # Filter eligible players to maintain 3-team constraint
+                    if must_maintain_team:
+                        # Must pick from same team OR ensure 3 teams remain
+                        # Get teams that would remain after removing current player
+                        remaining_teams = set(player_team_map.get(p) for p in current_lineup_players if p != current_player and player_team_map.get(p))
+                        
+                        # Filter to players whose team either:
+                        # 1. Is the same as current player (maintains that team)
+                        # 2. Is already in remaining_teams AND remaining_teams has at least 2 other teams
+                        def maintains_3_teams(player_team):
+                            if player_team == current_player_team:
+                                return True  # Same team, maintains the team
+                            if player_team in remaining_teams and len(remaining_teams) >= 3:
+                                return True  # Team already represented and we have enough teams
+                            if player_team not in remaining_teams and len(remaining_teams) >= 2:
+                                return True  # New team, but we still have 2 other teams + this one = 3
+                            return False
+                        
+                        eligible = eligible[eligible['Team'].apply(maintains_3_teams)]
+                    
+                    if eligible.empty:
+                        continue
+                        
                     # Randomly select a replacement
                     replacement = eligible.sample(1).iloc[0]
                     new_row[pos_to_swap] = replacement['Name']
@@ -479,9 +516,6 @@ def init_team_results(model_source: DataFrame, position_reqs: dict, salary_cap:
             # Get all position columns (exclude totals)
             pos_labels = [col for col in final_df.columns if col not in ['Total_Salary', 'Total_Median', 'Target_Team', 'Own']]
 
-            # Create a mapping from player name to team for fast lookup
-            player_team_map = filtered_df.set_index('Name')['Team'].to_dict()
-
             # Count target team players in each row
             def count_target_team(row):
                 return sum(1 for pos in pos_labels if player_team_map.get(row[pos], None) == target_team)