!rm -rf /home/jovyan/.theano/

%matplotlib inline

import datetime
from itertools import product
import logging
import pickle
from warnings import filterwarnings

from matplotlib import pyplot as plt
from matplotlib.offsetbox import AnchoredText
from matplotlib.ticker import FuncFormatter, StrMethodFormatter
import numpy as np
import pandas as pd
import scipy as sp
import seaborn as sns
from sklearn.preprocessing import LabelEncoder
from theano import shared, tensor as tt

# keep theano from complaining about compile locks for small models
(logging.getLogger('theano.gof.compilelock')
        .setLevel(logging.CRITICAL))

# silence PyMC3 warnings (there aren't many)
filterwarnings(
    'ignore', ".*diverging samples after tuning.*",
    module='pymc3'
)
filterwarnings(
    'ignore', ".*acceptance probability.*",
    module='pymc3'
)

pct_formatter = StrMethodFormatter('{x:.1%}')

blue, green, *_ = sns.color_palette()

# configure pyplot for readability when rendered as a slideshow and projected
sns.set(color_codes=True)

plt.rc('figure', figsize=(8, 6))

LABELSIZE = 14
plt.rc('axes', labelsize=LABELSIZE)
plt.rc('axes', titlesize=LABELSIZE)
plt.rc('figure', titlesize=LABELSIZE)
plt.rc('legend', fontsize=LABELSIZE)
plt.rc('xtick', labelsize=LABELSIZE)
plt.rc('ytick', labelsize=LABELSIZE)

SEED = 207183 # from random.org, for reproducibility

%%bash
DATA_URI=https://raw.githubusercontent.com/polygraph-cool/last-two-minute-report/32f1c43dfa06c2e7652cc51ea65758007f2a1a01/output/all_games.csv
DATA_DEST=/tmp/all_games.csv

if [[ ! -e $DATA_DEST ]];
then
    wget -q -O $DATA_DEST $DATA_URI
fi

USECOLS = [
    'period',
    'seconds_left',
    'call_type',
    'committing_player',
    'disadvantaged_player',
    'review_decision',
    'play_id',
    'away',
    'home',
    'date',
    'score_away',
    'score_home',
    'disadvantaged_team',
    'committing_team'
]

orig_df = pd.read_csv(
    '/tmp/all_games.csv',
    usecols=USECOLS,
    index_col='play_id',
    parse_dates=['date']
)

foul_df = orig_df[
    orig_df.call_type
           .fillna("UNKNOWN")
           .str.startswith("Foul")
]

FOULS = [
    f"Foul: {foul_type}"
    for foul_type in [
        "Personal",
        "Shooting",
        "Offensive",
        "Loose Ball",
        "Away from Play"
    ]
]

TEAM_MAP = {
    "NKY": "NYK",
    "COS": "BOS",
    "SAT": "SAS",
    "CHi": "CHI",
    "LA)": "LAC",
    "AT)": "ATL",
    "ARL": "ATL"
}

def correct_team_name(col):
    def _correct_team_name(df):
        return df[col].apply(lambda team_name: TEAM_MAP.get(team_name, team_name))
    
    return _correct_team_name

def date_to_season(date):
    if date >= datetime.datetime(2017, 10, 17):
        return '2017-2018'
    elif date >= datetime.datetime(2016, 10, 25):
        return '2016-2017'
    elif date >= datetime.datetime(2015, 10, 27):
        return '2015-2016'
    else:
        return '2014-2015'

clean_df = (foul_df.where(lambda df: df.period == "Q4")
                   .where(lambda df: (df.date.between(datetime.datetime(2016, 10, 25),
                                                      datetime.datetime(2017, 4, 12))
                                     | df.date.between(datetime.datetime(2015, 10, 27),
                                                       datetime.datetime(2016, 5, 30)))
                         )
                   .assign(
                       review_decision=lambda df: df.review_decision.fillna("INC"),
                       committing_team=correct_team_name('committing_team'),
                       disadvantged_team=correct_team_name('disadvantaged_team'),
                       away=correct_team_name('away'),
                       home=correct_team_name('home'),
                       season=lambda df: df.date.apply(date_to_season)
                    )
                   .where(lambda df: df.call_type.isin(FOULS))
                   .dropna()
                   .drop('period', axis=1)
                   .assign(call_type=lambda df: (df.call_type
                                                   .str.split(': ', expand=True)  
                                                   .iloc[:, 1])))

call_type_enc = LabelEncoder().fit(
    clean_df.call_type
)
n_call_type = call_type_enc.classes_.size

player_enc = LabelEncoder().fit(
    np.concatenate((
        clean_df.committing_player,
        clean_df.disadvantaged_player
    ))
)
n_player = player_enc.classes_.size

season_enc = LabelEncoder().fit(
    clean_df.season
)
n_season = season_enc.classes_.size

df = (clean_df[['seconds_left']]
              .round(0)
              .assign(
                call_type=call_type_enc.transform(clean_df.call_type),
                foul_called=1. * clean_df.review_decision.isin(['CC', 'INC']),
                player_committing=player_enc.transform(clean_df.committing_player),
                player_disadvantaged=player_enc.transform(clean_df.disadvantaged_player),
                score_committing=clean_df.score_home.where(
                    clean_df.committing_team == clean_df.home,
                    clean_df.score_away
                ),
                score_disadvantaged=clean_df.score_home.where(
                    clean_df.disadvantaged_team == clean_df.home,
                    clean_df.score_away
                ),
                season=season_enc.transform(clean_df.season)
              ))

def make_foul_rate_yaxis(ax, label="Observed foul call rate"):
    ax.yaxis.set_major_formatter(pct_formatter)
    ax.set_ylabel(label)
    
    return ax

import pymc3 as pm

with pm.Model() as base_model:
    β_season = pm.Normal('β_season', 0., 5., shape=n_season)
    p = pm.Deterministic('p', pm.math.sigmoid(β_season))

season = shared(df.season.values)

with base_model:
    y = pm.Bernoulli(
        'y', p[season],
        observed=df.foul_called.values
    )

NJOBS = 3

SAMPLE_KWARGS = {
    'draws': 1000,
    'njobs': NJOBS,
    'random_seed': [
        SEED + i for i in range(NJOBS)
    ]
}

bfmi = pm.bfmi(base_trace)
max_gr = max(np.max(gr_stats) for gr_stats in pm.gelman_rubin(base_trace).values())

CONVERGENCE_TITLE = lambda: f"BFMI = {bfmi:.2f}\nGelman-Rubin = {max_gr:.3f}"

resid_df = (df.assign(p_hat=base_trace['p'][:, df.season].mean(axis=0))
              .assign(resid=lambda df: df.foul_called - df.p_hat))

def make_time_axes(ax,
                   xlabel="Seconds remaining in game",
                   ylabel="Observed foul call rate"):
    ax.invert_xaxis()
    ax.set_xlabel(xlabel)
    
    return make_foul_rate_yaxis(ax, label=ylabel)

df['trailing_committing'] = (df.score_committing
                               .lt(df.score_disadvantaged)
                               .mul(1.)
                               .astype(np.int64))

with pm.Model() as poss_model:
    β_season = pm.Normal('β_season', 0., 5., shape=2)

def hierarchical_normal(name, shape, σ_shape=1):
    Δ = pm.Normal(
        f'Δ_{name}', 0., 1., shape=shape
    )
    σ = pm.HalfNormal(f'σ_{name}', 5., shape=σ_shape)
    
    return pm.Deterministic(name, Δ * σ)

with poss_model:
    β_call = hierarchical_normal('β_call', n_call_type)

df['score_diff'] = (df.score_disadvantaged
                      .sub(df.score_committing))

df['trailing_poss'] = (df.score_diff
                         .div(3)
                         .apply(np.ceil))

trailing_poss_enc = LabelEncoder().fit(df.trailing_poss)
trailing_poss = shared(
    trailing_poss_enc.transform(df.trailing_poss)
)
n_trailing_poss = trailing_poss_enc.classes_.size

df['remaining_poss'] = (df.seconds_left
                          .floordiv(25)
                          .add(1))

remaining_poss_enc = LabelEncoder().fit(df.remaining_poss)
remaining_poss = shared(
    remaining_poss_enc.transform(df.remaining_poss)
)
n_remaining_poss = remaining_poss_enc.classes_.size

def plot_foul_diff_heatmap(*_, data=None, **kwargs):
    ax = plt.gca()

    sns.heatmap(
        data.pivot_table(
            'diff',
            'trailing_poss',
            'remaining_poss'
        ),
        cmap='seismic', robust=True,
        cbar_kws={'format': pct_formatter}
    )
    
    ax.invert_yaxis()
    ax.set_title("Observed foul call rate")

call_name_df = df.assign(
    call_type=lambda df: call_type_enc.inverse_transform(
        df.call_type.values
    )
)

diff_df = (pd.merge(
                call_name_df,
                call_name_df.groupby('call_type')
                         .foul_called.mean()
                         .rename('avg_foul_called')
                         .reset_index()
             )
             .assign(diff=lambda df: df.foul_called - df.avg_foul_called))

with poss_model:
    β_poss = hierarchical_normal(
        'β_poss',
        (n_trailing_poss, n_remaining_poss, n_call_type),
        σ_shape=(1, 1, n_call_type)
    )

call_type = shared(df.call_type.values)

with poss_model:
    η_game = β_season[season] \
                + β_call[call_type] \
                + β_poss[
                    trailing_poss, remaining_poss, call_type
                ]

with poss_model:
    p = pm.Deterministic('p', pm.math.sigmoid(η_game))
    y = pm.Bernoulli('y', p, observed=df.foul_called)

bfmi = pm.bfmi(poss_trace)
max_gr = max(np.max(gr_stats) for gr_stats in pm.gelman_rubin(poss_trace).values())

resid_df = (df.assign(p_hat=poss_trace['p'].mean(axis=0))
              .assign(resid=lambda df: df.foul_called - df.p_hat))

N_BIN = 20

bin_ix, bins = pd.qcut(
    resid_df.p_hat, N_BIN,
    labels=np.arange(N_BIN),
    retbins=True
)

MODEL_NAME_MAP = {
    0: "Base",
    1: "Possession"
}

comp_df = (pm.compare(
                (base_trace, poss_trace),
                (base_model, poss_model)
             )
             .rename(index=MODEL_NAME_MAP)
             .loc[MODEL_NAME_MAP.values()])

with pm.Model() as irt_model:
    β_season = pm.Normal('β_season', 0., 5., shape=n_season)
    β_call = hierarchical_normal('β_call', n_call_type)
    β_poss = hierarchical_normal(
        'β_poss',
        (n_trailing_poss, n_remaining_poss, n_call_type),
        σ_shape=(1, 1, n_call_type)
    )
    
    η_game = β_season[season] \
                + β_call[call_type] \
                + β_poss[trailing_poss, remaining_poss, call_type]

player_committing = shared(df.player_committing.values)
player_disadvantaged = shared(df.player_disadvantaged.values)
n_player = player_enc.classes_.size

with irt_model:
    θ_player = hierarchical_normal(
        'θ_player', (n_player, n_season)
    )
    θ = θ_player[player_disadvantaged, season]

with irt_model:
    b_player = hierarchical_normal(
        'b_player', (n_player, n_season)
    )
    b = b_player[player_committing, season]

with irt_model:
    η_player = θ - b

with irt_model:
    η = η_game + η_player

with irt_model:
    p = pm.Deterministic('p', pm.math.sigmoid(η))
    y = pm.Bernoulli(
        'y', p,
        observed=df.foul_called
    )

bfmi = pm.bfmi(irt_trace)
max_gr = max(np.max(gr_stats) for gr_stats in pm.gelman_rubin(irt_trace).values())

resid_df = (df.assign(p_hat=irt_trace['p'].mean(axis=0))
              .assign(resid=lambda df: df.foul_called - df.p_hat))

N_BIN = 50

bin_ix, bins = pd.qcut(
    resid_df.p_hat, N_BIN,
    labels=np.arange(N_BIN),
    retbins=True
)

MODEL_NAME_MAP[2] = "IRT"

comp_df = (pm.compare(
                (base_trace, poss_trace, irt_trace),
                (base_model, poss_model, irt_model)
             )
             .rename(index=MODEL_NAME_MAP)
             .loc[MODEL_NAME_MAP.values()])

def varname_to_param(varname):
    return varname[0]

def varname_to_player(varname):
    return int(varname[3:-2])

def varname_to_season(varname):
    return int(varname[-1])

irt_df = (pm.trace_to_dataframe(
                irt_trace, varnames=['θ_player', 'b_player']
            )
            .rename(columns=lambda col: col.replace('_player', ''))
            .T
            .apply(
               lambda s: pd.Series.describe(
                   s, percentiles=[0.055, 0.945]
               ),
               axis=1
            )
            [['mean', '5.5%', '94.5%']]
            .rename(columns={
                 '5.5%': 'low',
                 '94.5%': 'high'
            })
            .rename_axis('varname')
            .reset_index()
            .assign(
                param=lambda df: df.varname.apply(varname_to_param),
                player=lambda df: df.varname.apply(varname_to_player),
                season=lambda df: df.varname.apply(varname_to_season)
            )
            .drop('varname', axis=1))

player_irt_df = irt_df.pivot_table(
    index='player',
    columns=['param', 'season'],
    values='mean'
)

def plot_latent_params(df):
    fig, ax = plt.subplots()
    
    n, _ = df.shape
    y = np.arange(n)

    ax.errorbar(
        df['mean'], y,
        xerr=(df[['high', 'low']]
                .sub(df['mean'], axis=0)
                .abs()
                .values.T),
        fmt='o'
    )

    ax.set_yticks(y)
    ax.set_yticklabels(
        player_enc.inverse_transform(df.player)
    )
    ax.set_ylabel("Player")
    
    return fig, ax

top_bot_irt_df = (irt_df.groupby('param')
                    .apply(
                        lambda df: pd.concat((
                                df.nlargest(10, 'mean'),
                                df.nsmallest(10, 'mean')
                            ),
                            axis=0, ignore_index=True
                        )
                    )
                    .reset_index(drop=True))

def p_val_to_asterisks(p_val):
    if p_val < 0.0001:
        return "****"
    elif p_val < 0.001:
        return "***"
    elif p_val < 0.01:
        return "**"
    elif p_val < 0.05:
        return "*"
    else:
        return ""

def plot_corr(x, y, **kwargs):
    corrcoeff, p_val = sp.stats.pearsonr(x, y)
    asterisks = p_val_to_asterisks(p_val)
    
    artist = AnchoredText(
        f'{corrcoeff:.2f}{asterisks}',
        loc=10, frameon=False,
        prop=dict(size=LABELSIZE)
    )
    plt.gca().add_artist(artist)
    plt.grid(b=False)

PARAM_MAP = {
    'θ': r"$\hat{\theta}$",
    'b': r"$\hat{b}$"
}

def replace_label(label):
    param, season = eval(label)

    return "{param}\n({season})".format(
        param=PARAM_MAP[param],
        season=season_enc.inverse_transform(season)
    )
    
def style_grid(grid):
    for ax in grid.axes.flat:
        ax.grid(False)
        ax.set_xticklabels([]);
        ax.set_yticklabels([]);
        
        if ax.get_xlabel():
            ax.set_xlabel(replace_label(ax.get_xlabel()))

        if ax.get_ylabel():
            ax.set_ylabel(replace_label(ax.get_ylabel()))
            
    return grid

player_all_season = set(df.groupby('player_disadvantaged')
                          .filter(lambda df: df.season.nunique() == n_season)
                          .player_committing) \
                        & set(df.groupby('player_committing')
                                .filter(lambda df: df.season.nunique() == n_season)
                                .player_committing)

MIN = 10

player_has_min = set(df.groupby('player_disadvantaged')
                       .filter(lambda df: df.season.value_counts().gt(MIN).all())
                       .player_committing) \
                    & set(df.groupby('player_committing')
                            .filter(lambda df: df.season.value_counts().gt(MIN).all())
                            .player_committing)

with open('./irt.pickle', mode='wb') as dest:
    pickle.dump(
        {
            'call_type_enc': call_type_enc,
            'player_enc': player_enc,
            'season_enc': season_enc,
            'trailing_poss_enc': trailing_poss_enc,
            'remaining_poss_enc': remaining_poss_enc,
            'df': df,
            'irt_trace': irt_trace
        },
        dest
    )

import dash
import dash_html_components as html

import dash_core_components as dcc
import plotly.graph_objs as go

CALL_TYPE = 0
PLAYER_COMMITTING = 0
PLAYER_DISADVANTAGED = 1
SEASON = 1

def get_pp_data_df(call_type=CALL_TYPE,
                   player_committing=PLAYER_COMMITTING,
                   player_disadvantaged=PLAYER_DISADVANTAGED,
                   season=SEASON):
    return (pd.DataFrame(
                list(product(
                    np.arange(n_trailing_poss),
                    np.arange(n_remaining_poss)
                )),
                columns=[
                    'trailing_poss',
                    'remaining_poss'
                ]
              )
              .assign(
                call_type=call_type,
                player_committing=player_committing,
                player_disadvantaged=1,
                season=SEASON
              ))

def get_irt_df(player_committing_=PLAYER_COMMITTING,
               player_disadvantaged_=PLAYER_DISADVANTAGED,
               season_=SEASON):
    return pd.DataFrame({
        'θ': irt_trace['θ_player'][:, player_disadvantaged_, season_],
        'b': irt_trace['b_player'][:, player_committing_, season_]
    })

def get_pp_df(call_type_=CALL_TYPE,
              player_committing_=PLAYER_COMMITTING,
              player_disadvantaged_=PLAYER_DISADVANTAGED,
              season_=SEASON):
    pp_data_df = get_pp_data_df(
        call_type=call_type_,
        player_committing=player_committing_,
        player_disadvantaged=player_disadvantaged_,
        season=season_
    )
    
    call_type.set_value(pp_data_df.call_type.values)
    player_committing.set_value(pp_data_df.player_committing.values)
    player_disadvantaged.set_value(pp_data_df.player_disadvantaged.values)
    season.set_value(pp_data_df.season.values)
    trailing_poss.set_value(pp_data_df.trailing_poss.values)
    remaining_poss.set_value(pp_data_df.remaining_poss.values)
    
    with irt_model:
        pp_trace = pm.sample_ppc(irt_trace)
        
    return pp_data_df.assign(
        p_hat=pp_trace['y'].mean(axis=0)
    )

def get_irt_scatter(irt_df):
    return {
        'data': [
            go.Scatter(
                x=irt_df.b,
                y=irt_df.θ,
                mode='markers',
                marker={'opacity': 0.5},
                hoverinfo='none'
            )
        ],
        'layout': go.Layout(
            title="Latent parameters",
            xaxis={'title': "b (committing player)"},
            yaxis={'title': "θ (disadvantaged player)"}
        )
    }

def get_irt_hist(irt_df):
    return {
        'data': [
            go.Histogram(
                x=irt_df.θ - irt_df.b,
                histnorm='probability',
                hoverinfo='none'
            )
        ],
        'layout': go.Layout(
            title="Latent parameter difference",
            xaxis={'title': "η_game"}
        )
    }

def get_foul_figure(pp_df):
    table = pp_df.pivot_table(
        'p_hat',
        'trailing_poss',
        'remaining_poss'
    )
    
    return {
        'data': [
            go.Heatmap(
                x=remaining_poss_enc.inverse_transform(
                    table.columns
                ),
                y=trailing_poss_enc.inverse_transform(
                    table.index
                ),
                z=table.values.tolist(),
                zmin=0., zmax=0.75,
                hoverinfo='z',
                colorbar={'tickformat': '.1%'}
            )
        ],
        'layout': go.Layout(
            title="Predicted foul call probability",
            xaxis={'title': "Remaining possessions"},
            yaxis={'title': "Trailing possessions"}
        )
    }

PLAYER_OPTIONS = [
    {
        'label': name,
        'value': i
    }
    for i, name in enumerate(player_enc.classes_)
]

INPUT_TABLE = html.Table(
    [
        html.Tr([
            html.Td(
                [
                    html.Div("Season"),
                    dcc.Dropdown(
                        id='season-dropdown',
                        options=[
                            {
                                'label': season,
                                'value': i
                            }
                            for i, season in enumerate(season_enc.classes_)
                        ],
                        value=0
                    )
                ],
                style={'width': '50%'}
            ),
            html.Td([
                html.Div("Call type"),
                dcc.Dropdown(
                    id='call-type-dropdown',
                    options=[
                        {
                            'label': call_type.lstrip("Foul: "),
                            'value': i
                        }
                        for i, call_type in enumerate(call_type_enc.classes_)
                    ],
                    value=0
                )
            ])
        ]),
        html.Tr([
            html.Td([
                    html.Div("Disadvantaged player"),
                    dcc.Dropdown(
                        id='player-disadvantaged-dropdown',
                        options=PLAYER_OPTIONS,
                        value=1
                    )
            ]),
            html.Td([
                    html.Div("Committing player"),
                    dcc.Dropdown(
                        id='player-committing-dropdown',
                        options=PLAYER_OPTIONS,
                        value=0
                    )
            ])
        ])
    ],
    style={'width': '100%'}
)