implement mix tracklist balancing

app/utils/mixes.py

@@ -0,0 +1,95 @@
+from app.models.track import Track
+from typing import List, Dict, Tuple
+from collections import Counter
+
+
+def violates_gap_rule(balanced_mix: Dict[int, Track], position: int, track: Track, gap: int = 3) -> bool:
+    """
+    Check if placing the track at the given position violates the gap rule.
+    """
+    track_artists = set(artist['artisthash'] for artist in track.artists)
+    
+    for i in range(max(0, position - gap), position):
+        if i in balanced_mix:
+            existing_artists = set(artist['artisthash'] for artist in balanced_mix[i].artists)
+            if track_artists.intersection(existing_artists):
+                return True
+    return False
+
+def find_next_position(balanced_mix: Dict[int, Track], start: int, track: Track, total_tracks: int) -> int:
+    """
+    Find the next available position for the track, starting from 'start' and wrapping around.
+    """
+    for i in range(start, total_tracks):
+        if i not in balanced_mix and not violates_gap_rule(balanced_mix, i, track):
+            return i
+    for i in range(start):
+        if i not in balanced_mix and not violates_gap_rule(balanced_mix, i, track):
+            return i
+    return start  # If no better position is found, return the original position
+
+def is_tracklist_balanced(tracks: List[Track], gap: int = 3) -> Tuple[bool, bool]:
+    """
+    Checks if a tracklist is balanced or can be balanced.
+    
+    Args:
+    - tracks: List of Track objects
+    - gap: Minimum number of tracks between songs by the same artist (default 3)
+    
+    Returns:
+    - A tuple (can_be_balanced, is_currently_balanced)
+    """
+    total_tracks = len(tracks)
+    
+    # Count tracks per artist (considering only the first artist)
+    artist_counts = Counter(track.artists[0]['artisthash'] for track in tracks)
+    
+    # Calculate the maximum number of tracks an artist can have in a balanced list
+    max_tracks_per_artist = (total_tracks + gap) // (gap + 1)
+    
+    # Check if it's mathematically possible to balance the tracklist
+    can_be_balanced = all(count <= max_tracks_per_artist for count in artist_counts.values())
+
+    if not can_be_balanced:
+        return False, False
+    
+    # Check if the current arrangement is balanced
+    is_currently_balanced = True
+    artist_last_positions = {}
+    
+    for i, track in enumerate(tracks):
+        artist = track.artists[0]['artisthash']
+        if artist in artist_last_positions:
+            if i - artist_last_positions[artist] <= gap:
+                is_currently_balanced = False
+                break
+        artist_last_positions[artist] = i
+    
+    return can_be_balanced, is_currently_balanced
+
+def balance_mix(tracks: List[Track]) -> List[Track]:
+    """
+    Balances the mix by ensuring that the tracks in a mix are distributed evenly.
+    Preserves the overall rating order of tracks while minimizing disruption.
+    """
+    can_be_balanced, is_balanced = is_tracklist_balanced(tracks)
+    
+    if not can_be_balanced:
+        print("Warning: This tracklist cannot be perfectly balanced.")
+        # Proceed with best-effort balancing
+    
+    if is_balanced:
+        return tracks  # Already balanced, no need to modify
+    
+    balanced_mix: Dict[int, Track] = {}
+    total_tracks = len(tracks)
+
+    for i, track in enumerate(tracks):
+        if i in balanced_mix or not violates_gap_rule(balanced_mix, i, track):
+            balanced_mix[i] = track
+        else:
+            new_position = find_next_position(balanced_mix, i, track, total_tracks)
+            balanced_mix[new_position] = track
+
+    # Convert the dictionary back to a list, preserving the new order
+    return [balanced_mix[i] for i in sorted(balanced_mix.keys())]