playoff.py

from __future__ import annotations

import time
from typing import Any

import pandas as pd
from pulp import PULP_CBC_CMD, LpMaximize, LpProblem, LpVariable, lpSum

POSITION: str = 'Pos'
PROJECTION: str = 'Total Points'
PLAYER: str = 'Player'

MAX_LINEUPS: int = 10

lineup_configs: dict[str, dict[str, int]] = {'playoff-league': {'QB': 2, 'RB': 4, 'WR': 4, 'TE': 2, 'K': 1, 'DST': 1}}


def calculate_lineups(
    lineup_type: dict[str, int],
    output_file: str,
    csv_file: str,
    must_include_players: list[str] | None = None,
    exclude_players: list[str] | None = None,
) -> None:
    must_include_players = must_include_players or []
    exclude_players = exclude_players or []

    players = pd.read_csv(csv_file, usecols=[PLAYER, POSITION, PROJECTION])
    # Trim whitespace from columns
    players = players.apply(lambda x: x.str.strip() if x.dtype == 'object' else x)

    # Validate and warn about missing players
    all_player_names = set(players[PLAYER].unique())
    missing_must_include = sorted(set(must_include_players) - all_player_names)
    if missing_must_include:
        print(f'WARNING: Must-include players not found in CSV: {", ".join(missing_must_include)}')

    missing_exclude = sorted(set(exclude_players) - all_player_names)
    if missing_exclude:
        print(f'WARNING: Exclude players not found in CSV: {", ".join(missing_exclude)}')

    # Filter out excluded players
    players = players[~players[PLAYER].isin(exclude_players)]

    player_data: dict[str, dict[str, float]] = {}
    for _, row in players.iterrows():
        pos: str = row[POSITION]
        if pos not in player_data:
            player_data[pos] = {}
        player_data[pos][row[PLAYER]] = row[PROJECTION]

    lineup_results: list[dict[str, Any]] = []
    previous_lineups: list[list[tuple[str, str]]] = []

    for lineup_num in range(1, MAX_LINEUPS + 1):
        prob: LpProblem = LpProblem('Fantasy', LpMaximize)

        player_vars: dict[str, dict[str, LpVariable]] = {}
        for pos, players in player_data.items():
            player_vars[pos] = LpVariable.dicts(f'{pos}_players', players.keys(), cat='Binary')

        prob += (
            lpSum(
                [
                    player_data[pos][player] * player_vars[pos][player]
                    for pos in player_vars
                    for player in player_vars[pos]
                ]
            ),
            'Total_Points',
        )

        # Enforce lineup constraints (how many players from each position)
        for pos, count in lineup_type.items():
            prob += (
                lpSum([player_vars[pos][player] for player in player_vars[pos]]) == count,
                f'{pos}_constraint',
            )

        # Enforce must-include players
        for must_include in must_include_players:
            for pos in player_vars:
                if must_include in player_vars[pos]:
                    prob += player_vars[pos][must_include] == 1, f'must_include_{must_include}_{lineup_num}'
                    break

        # Add each unique lineup only once
        for counter, prev_lineup in enumerate(previous_lineups):
            prob += (
                lpSum([player_vars[pos][player] for pos, player in prev_lineup]) <= len(prev_lineup) - 1,
                f'unique_lineup_{lineup_num}_{counter}',
            )

        prob.solve(PULP_CBC_CMD(msg=0))  # Suppress noisy output

        current_lineup_players: list[tuple[str, str]] = [
            (pos, player) for pos in player_vars for player, var in player_vars[pos].items() if var.varValue == 1
        ]

        # Only find unique lineups up to MAX_LINEUPS
        if not current_lineup_players or len(previous_lineups) >= MAX_LINEUPS:
            break

        # Add the current lineup's players to the list of previous lineups
        previous_lineups.append(current_lineup_players)

        lineup: dict[str, Any] = {'Lineup #': lineup_num}
        total_score: float = 0.0

        for column_count, (pos, player) in enumerate(current_lineup_players, start=1):
            proj: float = player_data[pos][player]
            lineup[f'Player {column_count} Position'] = pos
            lineup[f'Player {column_count} Name'] = player
            lineup[f'Player {column_count} Projected Points'] = proj
            total_score += proj

        lineup['Total Score'] = f'{total_score:.1f}'
        lineup_results.append(lineup)

    pd.DataFrame(lineup_results).to_csv(f'{output_file}.csv', index=False, header=True)


def generate_lineup_files(
    csv_file: str,
    must_include_players: list[str] | None = None,
    exclude_players: list[str] | None = None,
) -> None:
    must_include_players = must_include_players or []
    exclude_players = exclude_players or []

    for name, config in lineup_configs.items():
        calculate_lineups(config, name, csv_file, must_include_players, exclude_players)

    print('Lineup files created')

    csv_files = [f'{name}.csv' for name in lineup_configs]
    combined_df = pd.concat([pd.read_csv(file) for file in csv_files])
    combined_df['Total Score'] = pd.to_numeric(combined_df['Total Score'])
    combined_df = combined_df.sort_values(by='Total Score', ascending=False)
    combined_df.to_csv('combined_lineups.csv', index=False, header=False)

    # Truncate player names to first 12 chars so they fit better in the terminal
    for col in combined_df.columns:
        if combined_df[col].dtype == 'object':
            combined_df[col] = combined_df[col].str[:12]

    print(combined_df.to_string(index=False, header=False))


if __name__ == '__main__':
    start_time: float = time.time()
    file_name: str = 'playoff.csv'

    must_include = []
    exclude_list = []

    generate_lineup_files(file_name, must_include_players=must_include, exclude_players=exclude_list)
    end_time: float = time.time()

    print(f'Total execution time: {end_time - start_time:.2f} seconds')