import polars as pl
import numpy as np
import joblib
from shiny import App, reactive, render, ui
import matplotlib.pyplot as plt
import matplotlib.ticker as tkr
import seaborn as sns
import adjustText
sns.set_style('whitegrid')


import matplotlib
cmap_sum = matplotlib.colors.LinearSegmentedColormap.from_list("", ['#FFFFFF','#FFB000','#FE6100'])

xwoba_model = joblib.load('joblib_model/xwoba_model.joblib')

x = np.arange(-30,90.5,.5)
y = np.arange(0,120.5,0.1)

xx, yy = np.meshgrid(x, y)

df = pl.DataFrame({'launch_angle': xx.ravel(), 'launch_speed': yy.ravel()})

df = df.with_columns(
    pl.Series('xwoba', xwoba_model.predict_proba(df.select(['launch_angle','launch_speed'])) @ [0, 0.883, 1.244, 1.569, 2.004])
)

df = df.with_columns(
    pl.Series('xslg', xwoba_model.predict_proba(df.select(['launch_angle','launch_speed'])) @ [0, 1, 2, 3, 4])
)

app_ui = ui.page_sidebar(
    ui.sidebar(
        ui.markdown("""
        ### How to use this app
        
        1. Click anywhere on the plot to select a point, or manually enter coordinates
        2. The selected point's coordinates will update automatically
        3. The xwOBA value will be calculated based on these coordinates
        """),
        ui.hr(),
        ui.input_numeric("x_select", "Launch Speed (mph)", value=110),
        ui.input_numeric("y_select", "Launch Angle (°)", value=30),
        ui.input_switch("flip_stat", "xwOBA", value=False),


    ),
    ui.output_plot("plot",width='900px',height='900px', click=True)
)


def server(input, output, session):
    # Store the coordinates in reactive values
    x_coord = reactive.value(110)
    y_coord = reactive.value(30)
    
    @reactive.effect
    @reactive.event(input.plot_click)
    def _():
        # Update reactive values when plot is clicked
        click_data = input.plot_click()
        if click_data is not None:
            x_coord.set(click_data["x"])
            y_coord.set(click_data["y"])
            # Update the numeric inputs
            ui.update_numeric("x_select", value=round(click_data["x"],1))
            ui.update_numeric("y_select", value=round(click_data["y"],1))
    
    @reactive.effect
    @reactive.event(input.x_select, input.y_select)
    def _():
        # Update reactive values when numeric inputs change
        x_coord.set(round(input.x_select(),1))
        y_coord.set(round(input.y_select(),1))


    @render.plot
    def plot():
        switch = input.flip_stat()
        fig, ax = plt.subplots(1, 1, figsize=(9, 9))


        if switch:
            h = ax.hexbin(df['launch_speed'], 
                        df['launch_angle'],
                        C=df['xwoba'], 
                        gridsize=(40,25), 
                        cmap=cmap_sum,
                        vmin=0.0,
                        vmax=2.0,)
            bounds=[0.0,0.4,0.8,1.2,1.6,2.0]
            fig.colorbar(h, ax=ax, label='xwOBA',format=tkr.FormatStrFormatter('%.3f'),shrink=0.5,
                        ticks=bounds)
            
        else:
            h = ax.hexbin(df['launch_speed'], 
                        df['launch_angle'],
                        C=df['xslg'], 
                        gridsize=(40,25), 
                        cmap=cmap_sum,
                        vmin=0.0,
                        vmax=4.0,)
            bounds=[0.0,0.5,1,1.5,2,2.5,3,3.5,4]
            fig.colorbar(h, ax=ax, label='xSLG',format=tkr.FormatStrFormatter('%.3f'),shrink=0.5,
                        ticks=bounds)



        ax.set_xlabel('Launch Speed')
        ax.set_ylabel('Launch Angle')
        if switch:
            ax.set_title('Exit Velocity vs Launch Angle\nExpected Weighted On Base Average (xwOBA)\nBy: @TJStats, Data:MLB')
        else:
            ax.set_title('Exit Velocity vs Launch Angle\nExpected Total Bases (xSLG)\nBy: @TJStats, Data:MLB')
            
        ax.grid(False)
        ax.axis('square')
        ax.set_xlim(0, 120)
        ax.set_ylim(-30, 90)

        x_select = input.x_select()
        y_select = input.y_select()


        sns.scatterplot(x=[x_select],y=[y_select],color='#648FFF',s=50,ax=ax,edgecolor='k',zorder=100)

        
        if switch:
            xwoba_value = (xwoba_model.predict_proba([[y_select,x_select]]) @ [0, 0.883, 1.244, 1.569, 2.004])[0]
            texts = [ax.text(x_select+3, y_select+3, f'xwOBA: {xwoba_value:.3f}', color='black', fontsize=12, weight='bold',
                            zorder=1000, bbox=dict(facecolor='white', alpha=0.5, edgecolor='black'))]
    
        else:
            xwoba_value = (xwoba_model.predict_proba([[y_select,x_select]]) @ [0, 1, 2, 3, 4])[0]
            texts = [ax.text(x_select+3, y_select+3, f'xSLG: {xwoba_value:.3f}', color='black', fontsize=12, weight='bold',
                            zorder=1000, bbox=dict(facecolor='white', alpha=0.5, edgecolor='black'))]
    



        adjustText.adjust_text(texts,

                                arrowprops=dict(arrowstyle='->', color='#DC267F'),avoid_self=True,
                                min_arrow_len =5)
        # xwoba_value = 

        ax.axhline(y=y_select, color='k', linestyle='--',linewidth=1,alpha=0.5)
        ax.axvline(x=x_select, color='k', linestyle='--',linewidth=1,alpha=0.5)

        # ax.axis('square')


app = App(app_ui, server)