src/pages/10_Set_Your_Lineup.py

from dataclasses import dataclass
import pandas as pd
import streamlit as st

from config import DEFAULT_ICON
from shared_page import common_page_config

from domain.constants import SEASON
from domain.playoffs import PLAYOFF_WEEK_TO_NAME, CURRENT_PLAYOFF_WEEK, ROSTER_WEEK_TO_PLAYOFF_WEEK
from domain.teams import SCHEDULE_NAME_TO_PFR_NAME_MAP, PLAYOFFS_TEAMS
from queries.nflverse.github_data import get_weekly_rosters
from queries.pfr.league_schedule import get_season_time_map
from login import check_password
from data_storage import update_selection, get_user_team


@dataclass
class PlayerOption:
    full_name: str
    gsis_id: str
    headshot_url: str
    position: str
    team: str
    gametime: pd.Timestamp | None
    week: int | None

    @classmethod
    def from_series(cls, input_series):
        return cls(
            full_name=input_series.full_name,
            gsis_id=input_series.gsis_id,
            headshot_url=input_series.headshot_url,
            position=input_series.position,
            team=input_series.team,
            gametime=input_series.gametime,
            week=int(input_series.week),
        )

    @classmethod
    def empty_player(cls, week: int | None = None):
        return cls(full_name="", gsis_id="", headshot_url="", position="", team="", gametime=None, week=week)

    def is_locked(self):
        if not self.gametime:
            return False
        else:
            return self.gametime < pd.Timestamp.now()


def initialize_empty_options_map() -> dict[str, dict[int, list[PlayerOption]]]:
    options_map: dict[str, dict[int, list[PlayerOption]]] = {}
    for pos in ["QB", "RB", "WR", "TE", "K", "DEF"]:
        options_map[pos] = {}
        for week in PLAYOFF_WEEK_TO_NAME.keys():
            options_map[pos][int(week)] = [PlayerOption.empty_player(week=week)]
    return options_map


def player_options_from_df(df_options) -> dict[str, dict[int, list[PlayerOption]]]:
    options_map = initialize_empty_options_map()
    for pos, pos_week_map in options_map.items():
        for week in pos_week_map.keys():
            df_pos_week = df_options[((df_options.week == week) & (df_options.position == pos))]
            if len(df_pos_week) > 0:
                player_options_list = df_pos_week.apply(PlayerOption.from_series, axis=1).tolist()
                options_map[pos][int(week)].extend(player_options_list)
    return options_map


@st.cache_data(ttl=60 * 60 * 24)
def load_options():
    df_rosters = get_weekly_rosters()

    # get game schedules
    week_game_times = get_season_time_map(SEASON)
    latest_game_time_defaults = {k: max(v.values()) for k, v in week_game_times.items() if v}

    # sort
    sort_by_cols = ["position", "week", "fantasy_points"]
    df_rosters.sort_values(sort_by_cols, ascending=False, inplace=True)

    # filter data from non-playoffs
    df_rosters = df_rosters[df_rosters.week.isin(ROSTER_WEEK_TO_PLAYOFF_WEEK.keys())]
    df_rosters["week"] = df_rosters["week"].map(ROSTER_WEEK_TO_PLAYOFF_WEEK)
    # set gametime
    if len(df_rosters) == 0:
        return initialize_empty_options_map()
    df_rosters["gametime"] = df_rosters.apply(
        lambda x: week_game_times.get(x.week, {}).get(
            SCHEDULE_NAME_TO_PFR_NAME_MAP[x.team], latest_game_time_defaults.get(x.week, None)
        ),
        axis=1,
    )

    df_rosters["in_playoffs"] = df_rosters.apply(lambda x: x.team in PLAYOFFS_TEAMS[x.week], axis=1)

    df_rosters = df_rosters[df_rosters.in_playoffs]
    player_options = player_options_from_df(df_rosters)
    return player_options


def format_player_option(player_opt: PlayerOption) -> str:
    return f"{player_opt.team} - {player_opt.full_name}"


def display_player(player_opt: PlayerOption | None):
    if player_opt:
        if player_opt.headshot_url:
            st.image(player_opt.headshot_url)
        if player_opt.full_name:
            st.write(player_opt.full_name)
            st.write(f"{player_opt.team} - {player_opt.gametime.strftime('%-m/%-d %-I:%M %p')}")


def position_cell(
    week: str, pos_str: str, pos_idx: int, options_map: dict[str, dict[int, list[PlayerOption]]], existing_selection_map
):
    pos_label = f"{week}-{pos_str}-{pos_idx}"
    selected_id = existing_selection_map.get(pos_label)
    options_list = options_map[pos_str][int(week)]
    disabled = False
    if isinstance(selected_id, str):
        try:
            selected_option_idx, selected_player = next(
                (i, v) for i, v in enumerate(options_list) if str(selected_id) == str(v.gsis_id)
            )
        except StopIteration:
            selected_player = PlayerOption.empty_player()
            selected_option_idx = 0
    else:
        selected_player = PlayerOption.empty_player()
        selected_option_idx = 0
    if int(week) > CURRENT_PLAYOFF_WEEK:
        options = []
        selected_option_idx = 0
        disabled = True
    elif int(week) < CURRENT_PLAYOFF_WEEK or selected_player.is_locked():
        options = [selected_player]
        selected_option_idx = 0
        disabled = True
    else:
        options = [x for x in options_list if not x.is_locked()]

    selected_player_from_box = st.selectbox(
        pos_str,
        options=options,
        format_func=format_player_option,
        index=selected_option_idx,
        key=pos_label,
        disabled=disabled,
    )
    if selected_player_from_box and int(week) == CURRENT_PLAYOFF_WEEK:
        if selected_player_from_box.gsis_id and selected_player_from_box.gsis_id != selected_id:
            if selected_player_from_box.is_locked():
                st.warning("Sorry player's game has already started", icon="🚨")
                display_player(selected_player)
                return
            elif selected_player_from_box.gsis_id in existing_selection_map.values():
                st.warning("Player already in lineup.  Please choose another.", icon="🚨")
                display_player(selected_player)
                return
            else:
                update_and_save_selection(pos_label, selected_player_from_box.gsis_id),
        display_player(selected_player_from_box)
    else:
        display_player(selected_player)


def update_and_save_selection(pos_label: str, selection_id: str):
    update_selection(st.session_state["logged_in_user"], pos_label, selection_id)


def get_page():
    page_title = "Select Your Team"
    st.set_page_config(page_title=page_title, page_icon=DEFAULT_ICON, layout="wide")
    common_page_config()
    if not check_password():
        st.write("Sorry, you must be logged in first to play")
        st.stop()

    st.title(page_title)

    if st.button("Refresh Data"):
        st.rerun()
    existing_selections = get_user_team(st.session_state["logged_in_user"])

    player_options = load_options()

    for week in range(1, 5):
        st.header(PLAYOFF_WEEK_TO_NAME[week])
        selection_cols = st.columns(8)

        with selection_cols[0]:
            position_cell(week, "QB", 1, player_options, existing_selections)
        with selection_cols[1]:
            position_cell(week, "RB", 1, player_options, existing_selections)
        with selection_cols[2]:
            position_cell(week, "RB", 2, player_options, existing_selections)
        with selection_cols[3]:
            position_cell(week, "WR", 1, player_options, existing_selections)
        with selection_cols[4]:
            position_cell(week, "WR", 2, player_options, existing_selections)
        with selection_cols[5]:
            position_cell(week, "TE", 1, player_options, existing_selections)
        with selection_cols[6]:
            position_cell(week, "K", 1, player_options, existing_selections)
        with selection_cols[7]:
            position_cell(week, "DEF", 1, player_options, existing_selections)


if __name__ == "__main__":
    get_page()