import gradio as gr
from gradio_calendar import Calendar
import polars as pl
import pandas as pd
import matplotlib.pyplot as plt
from plottable import Table, ColumnDefinition
from plottable.plots import circled_image

from data import df, game_df
from gradio_function import *
from css import css

import datetime

df = (
    df
    # .join(game_df, on='game_pk')
    # .with_columns(pl.col('game_date').str.to_datetime())
    .rename({
        'name': 'Name',
        'release_speed': 'KPH',
        'team': 'Team'
    })
)


def filter_pitcher_leaderboard_by_date(date, *args, **kwargs):
    day_df = df.filter(pl.col('game_date') == date)
    monday = date - datetime.timedelta(days=date.weekday())
    sunday = date + datetime.timedelta(days=6-date.weekday())
    week_df = df.filter((pl.col('game_date') >= monday) & (pl.col('game_date') <= sunday))
    daily_whiffs, daily_velos = compute_pitcher_leaderboards(day_df, *args, **kwargs)
    weekly_whiffs, weekly_velos = compute_pitcher_leaderboards(week_df, *args, **kwargs)
    return (
        f'<center><h1>Daily Leaderboard<h1><h2>{date.strftime("%B %d, %Y")}</h2><h3>{date.strftime("%A")}</h3></center>',
        f'<center><h1>Weekly Leaderboard<h1><h2>{monday.strftime("%B %d, %Y")} to {sunday.strftime("%B %d, %Y")}</h2><h3>{monday.strftime("%A")} to {sunday.strftime("%A")}</h3></center>',
        daily_whiffs.drop('Team'),
        daily_velos.drop('Team'),
        weekly_whiffs.drop('Team'),
        weekly_velos.drop('Team'),
        *plot_tables(daily_whiffs, daily_velos, 'Daily', f'{date.strftime("%B %d, %Y")}\n{date.strftime("%A")}'),
        *plot_tables(weekly_whiffs, weekly_velos, 'Weekly', f'{monday.strftime("%B %d, %Y")} to {sunday.strftime("%B %d, %Y")}\n{monday.strftime("%A")} to {sunday.strftime("%A")}'),
        gr.update(interactive=True),
        gr.update(interactive=True),
        gr.update(interactive=True),
        gr.update(interactive=True)
    )


def compute_pitcher_leaderboards(df, top_players, strict, ignore_zero_whiffs, show_rank, debug):
    # _df = df.filter(pl.col('game_date') == date)
    _df = df

    other_cols = ['Team', 'Name']
    
    if debug:
        other_cols = ['game_date'] + other_cols
        
    whiffs = (
        _df
        .group_by(['pitcher'])
        .agg(
            pl.col('whiff').sum().alias('Whiffs'),
            *[pl.col(col).first() for col in other_cols]
        )
        .select(*other_cols, 'Whiffs')
        .sort('Whiffs', descending=True)
    )
    if ignore_zero_whiffs:
        whiffs = whiffs.filter(pl.col('Whiffs') > 0)
    if len(whiffs) >top_players:
        whiffs = (
            whiffs
            .filter(pl.col('Whiffs') >= whiffs['Whiffs'][top_players])
        )
    if strict:
        whiffs = whiffs[:top_players]
    if show_rank:
        whiffs = (
            whiffs
            .with_row_index(offset=1)
            .rename({'index': 'Rank'})
        )
    
    velos = (
        _df
        .select(*other_cols, 'KPH')
        .with_columns((pl.col('KPH') / 1.609).round().cast(pl.Int16).alias('MPH'))
        .drop_nulls()
        .sort(['KPH', 'Name'], descending=[True, False])
    )
    if len(velos) > top_players:
        velos = velos.filter(pl.col('KPH') >= velos['KPH'][top_players])
    if strict:
        velos = velos[:top_players]
    if show_rank:
        velos = (
            velos
            .with_row_index(offset=1)
            .rename({'index': 'Rank'})
        )

    return whiffs, velos
    # return (
        # f'<center><h1>Daily Leaderboard<h1><h2>{date.strftime("%B %d, %Y")}</h2><h3>{date.strftime("%A")}</h3></center>',
        # whiffs,
        # velos,
        # gr.update(interactive=True),
        # gr.update(interactive=True)
    # )


def plot_tables(whiffs, velos, time_type, subheader):
    whiff_fig, whiff_ax = plt.subplots(figsize=(4, 6))

    whiffs = (
        whiffs
        .with_columns(
            pl.col('Team').map_elements(lambda team: f'assets/{team.lower()}.png', return_dtype=str)
        )
    ).to_pandas()
    if 'Rank' in whiffs.columns:
        whiffs = whiffs.set_index('Rank')
    else:
        whiffs.index = pd.Series(range(1, len(whiffs)+1), name='Rank')
    Table(
      (
        whiffs
      ),
      column_definitions=[
        ColumnDefinition(name="Rank", title="Rank", width=0.25),
        ColumnDefinition(name='Team', title='Team', width=0.25, plot_fn=circled_image, textprops={'ha': 'center'}),
        ColumnDefinition(name="Name", title="Player", textprops={'ha': 'left'}),
        ColumnDefinition(name="Whiffs", title="#", width=0.25)
      ],
      ax=whiff_ax
    )
    whiff_fig.suptitle(f'{time_type} Whiff Leaderboard\n{subheader}')

    velo_fig, velo_ax = plt.subplots(figsize=(4, 6))
    velos = (
        velos
        .with_columns(
            pl.col('Team').map_elements(lambda team: f'assets/{team.lower()}.png', return_dtype=str)
        )
    ).to_pandas()
    if 'Rank' in velos.columns:
        velos = velos.set_index('Rank')
    else:
        velos.index = pd.Series(range(1, len(velos)+1), name='Rank')
    Table(
      velos,
      column_definitions=[
        ColumnDefinition(name="Rank", title="Rank", width=0.25),
        ColumnDefinition(name='Team', title='Team', width=0.25, plot_fn=circled_image, textprops={'ha': 'center'}),
        ColumnDefinition(name="Name", title="Player", textprops={'ha': 'left'}),
        ColumnDefinition(name="KPH", title="KPH", width=0.25),
        ColumnDefinition(name='MPH', title='MPH', width=0.25)
      ],
      ax=velo_ax
    )
    velo_fig.suptitle(f'{time_type} Velocity Leaderboard\n{subheader}')
    
    return whiff_fig, velo_fig

def go_back_day(date):
    return date - datetime.timedelta(days=1)


def go_forward_day(date):
    return date + datetime.timedelta(days=1)


def go_back_week(date):
    return date - datetime.timedelta(days=7)


def go_forward_week(date):
    return date + datetime.timedelta(days=7)


def create_daily_pitcher_leaderboard():
    with gr.Blocks(
        css=css
    ) as demo:
        with gr.Row():
            # date_picker = gr.DateTime(
                # value=df['game_date'].max().strftime('%Y-%m-%d'),
                # include_time=False,
                # type='datetime',
                # label='Date',
                # scale=4
            # )
            date_picker = Calendar(
                value=df['game_date'].max().strftime('%Y-%m-%d'),
                type='datetime',
                label='Date',
                scale=2,
                min_width=50
            )
            top_players = gr.Number(10, label='# Top players', scale=1, min_width=100)
            strict = gr.Checkbox(False, label='Strict', info='Ignore ties and restrict to # top players', scale=2, min_width=100)
            ignore_zero_whiffs = gr.Checkbox(False, label='Ignore zero whiffs', info='Ignore zero whiff players if in top ranked', scale=2, min_width=100)
            show_rank = gr.Checkbox(False, label='Show rank', scale=1, min_width=100)
            debug = gr.Checkbox(False, label='Debug', info='Show dates', scale=1, min_width=100)
            search_btn = gr.Button('Search', scale=1, min_width=100)

        with gr.Row():
            prev_week_btn = gr.Button('Previous Week', interactive=False)
            prev_day_btn = gr.Button('Previous Day', interactive=False)
            next_day_btn = gr.Button('Next Day', interactive=False)
            next_week_btn = gr.Button('Next Week', interactive=False)

        with gr.Tab('Tables for viewing'):
            daily_header = gr.HTML('<center><h1>Daily Leaderboard<h1><h2 style="display: none;"></h2><h3 style="display: none;"></h3></center>')
            with gr.Row():
                daily_whiffs = gr.Dataframe(pl.DataFrame({'Name': [], 'Whiffs': []}), label='Whiffs', interactive=False, height=1000)
                daily_velos = gr.Dataframe(pl.DataFrame({'Name': [], 'KPH': [], 'MPH': []}), label='Velocity', interactive=False, height=1000)

            weekly_header = gr.HTML('<center><h1>Weekly Leaderboard<h1><h2 style="display: none;"></h2><h3 style="display: none;"></h3></center>')
            with gr.Row():
                weekly_whiffs = gr.Dataframe(pl.DataFrame({'Name': [], 'Whiffs': []}), label='Whiffs', interactive=False, height=1000)
                weekly_velos = gr.Dataframe(pl.DataFrame({'Name': [], 'KPH': [], 'MPH': []}), label='Velocity', interactive=False, height=1000)

        with gr.Tab('Tables for sharing'):
            gr.Markdown('''# Plotted leaderboards

            For easier sharing
            ''')
            with gr.Row():
                daily_whiffs_plot = gr.Plot(label='Whiffs')
                daily_velos_plot = gr.Plot(label='Velocity')

            with gr.Row():
                weekly_whiffs_plot = gr.Plot(label='Whiffs')
                weekly_velos_plot = gr.Plot(label='Velocity')
            
        search_kwargs = dict(
            fn=filter_pitcher_leaderboard_by_date,
            inputs=[date_picker, top_players, strict, ignore_zero_whiffs, show_rank, debug],
            outputs=[daily_header, weekly_header, daily_whiffs, daily_velos, weekly_whiffs, weekly_velos, daily_whiffs_plot, daily_velos_plot, weekly_whiffs_plot, weekly_velos_plot, prev_day_btn, next_day_btn, prev_week_btn, next_week_btn]
        )
        search_btn.click(**search_kwargs)
        for btn, fn in (
            (prev_day_btn, go_back_day),
            (next_day_btn, go_forward_day),
            (prev_week_btn, go_back_week),
            (next_week_btn, go_forward_week)
        ):
            (
                btn
                .click(fn, date_picker, date_picker)
                .then(**search_kwargs)
            )
        
        
        
    return demo

demo = create_daily_pitcher_leaderboard()

if __name__ == '__main__':
    # demo = create_daily_pitcher_leaderboard()
    demo.launch()