import numpy as np
import pandas as pd
import streamlit as st
import pymongo
import os

st.set_page_config(layout="wide")

@st.cache_resource
def init_conn():
    uri = os.getenv('mongo_uri')
    client = pymongo.MongoClient(uri, retryWrites=True, serverSelectionTimeoutMS=500000)
    db = client["MLB_Database"]

    return db

db = init_conn()

@st.cache_resource(ttl = 60)
def init_baselines():
    collection = db["HR_Table"] 
    cursor = collection.find()
    raw_display = pd.DataFrame(cursor)
    raw_display.rename(columns={"Names": "Player"}, inplace = True)
    raw_frame = raw_display.drop(columns=['_id']).drop_duplicates(subset='Player')
    raw_frame = raw_frame[raw_frame['Binom_xHR'] != "#N/A"]

    return raw_frame

hr_frame = init_baselines()
st.title("HR Finder Table")

with st.container():
    col1, col2, col3 = st.columns(3)
    with col1:
        disp_options = st.radio("Display options:", options = ['Basics', 'Exclude DFS Info', 'Include DFS Info'], key='display_options')
    with col2:
        team_options = st.multiselect("Parse Teams:", options = hr_frame.Team.unique(), key = 'team_options')
    with col3:
        pos_options = st.multiselect("Parse Positions:", options = ['C', '1B', '2B', '3B', 'SS', 'OF'], key = 'pos_options')

if len(team_options) > 0:
    hr_frame = hr_frame[hr_frame['Team'].isin(team_options)]

if len(pos_options) > 0:
    position_mask = hr_frame['Position'].apply(lambda x: any(pos in x for pos in pos_options))
    hr_frame = hr_frame[position_mask]

if disp_options == 'Basics':
    st.session_state['disp_frame'] = hr_frame[['Player', 'xHR/PA', 'Opp_xHR/PA', 'BP_Binom_xHR', 'Binom_xHR']]
    st.session_state['disp_frame'] = st.session_state['disp_frame'].rename(columns={'xHR/PA': 'Hitter', 'Opp_xHR/PA': 'SP', 'BP_Binom_xHR': 'BP', 'Binom_xHR': 'xHRs'})
    st.session_state['disp_frame'] = st.session_state['disp_frame'].sort_values(by='xHRs', ascending=False)
elif disp_options == 'Exclude DFS Info':
    st.session_state['disp_frame'] = hr_frame.drop(columns=['Salary', 'Position', 'FD_Position', 'Order', 'NQ_Factor'])
    st.session_state['disp_frame'] = st.session_state['disp_frame'].sort_values(by='Binom_xHR', ascending=False)
else:
    st.session_state['disp_frame'] = hr_frame.copy()
    st.session_state['disp_frame'] = st.session_state['disp_frame'].sort_values(by='Binom_xHR', ascending=False)

st.session_state['disp_frame'] = st.session_state['disp_frame'].set_index('Player', drop = True)

if 'disp_frame' in st.session_state:
    st.dataframe(st.session_state['disp_frame'].style.background_gradient(axis=0).background_gradient(cmap = 'RdYlGn').format(precision=2), height=1200, use_container_width = True)