from plotnine import ggplot, annotate, aes, geom_text, geom_point, scale_size_manual, scale_shape_manual
from plotnine import scale_colour_manual, ylim, xlim, coord_fixed, theme_void
from plotnine.animation import PlotnineAnimation
import pandas as pd
import numpy as np
import warnings
from tqdm import tqdm


def visualize_play(play_data, path=None):
    # General field boundaries
    xmin = 0
    xmax = 160 / 3.
    # hash_right = 38.35
    hash_left = 12
    # hash_width = 3.3
    ymin = max(round(np.nanmin(play_data.x) - 10, -1), 0)
    ymax = min(round(np.nanmax(play_data.x) + 10, -1), 120)
    df_x_hash, df_y_hash = np.meshgrid([0, 23.36667, 29.96667, xmax], list(range(10, 110 + 1)))
    df_hash = pd.DataFrame({'x': df_x_hash.flatten(), 'y': df_y_hash.flatten()})
    df_hash['label'] = '_'

    df_hash = df_hash[(df_hash.y % 5 != 0) & (df_hash.y < ymax) & (df_hash.y > ymin)].reset_index(drop=True)

    df_hash_left = df_hash[df_hash['x'] < 55 / 2.]
    df_hash_right = df_hash[df_hash['x'] > 55 / 2.]

    # There are some odd warnings thrown, this takes care of them.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        play_data.is_copy = False
        pd.set_option('chained_assignment', None)
        plots = (plot(k, df_hash_left, df_hash_right, ymin, ymax, xmin, xmax, hash_left) for k in tqdm(play_data.groupby('frame.id')))

        ani = PlotnineAnimation(plots, interval=100, repeat_delay=500)

    # Save it out if a path is provided
    if path:
        ani.save(path)
    return ani


def plot(data, df_hash_left, df_hash_right, ymin, ymax, xmin, xmax, hash_left):
    df_frame = data[1]

    p = (ggplot() +
         geom_point(df_frame, aes(x='plot_x', y='plot_y', colour='team', group='nflId', pch='team', size='team'),
                    inherit_aes=False, na_rm=True) +
         geom_text(df_frame, aes(x='plot_x', y='plot_y', label='jerseyNumber'), colour="white", size=9,
                   inherit_aes=False, na_rm=True) +
         geom_text(df_hash_left, aes(label='label', x='x', y='y'), nudge_x=0.9, nudge_y=.8) +
         geom_text(df_hash_right, aes(label='label', x='x', y='y'), nudge_x=-1.1, nudge_y=.8) +
         annotate("segment", x=xmin, y=list(range(max(10, int(ymin)), min(int(ymax), 110) + 5, 5)),
                  xend=xmax, yend=list(range(max(10, int(ymin)), min(int(ymax), 110) + 5, 5))) +
         annotate("segment", x=[xmin, xmin, xmax, xmax], y=[ymin, ymax, ymax, ymin],
                  xend=[xmin, xmax, xmax, xmin], yend=[ymax, ymax, ymin, ymin]) +
         annotate("text", x=[hash_left] * 11, y=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110],
                  label=["G   ", 10, 20, 30, 40, 50, 40, 30, 20, 10, "   G"], angle=270, size=10) +
         annotate("text", x=[xmax - hash_left] * 11, y=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110],
                  label=["G   ", 10, 20, 30, 40, 50, 40, 30, 20, 10, "   G"], angle=270, size=10) +
         scale_size_manual(values=[6, 4, 6], guide=False) +
         scale_shape_manual(values=[19, 16, 19], guide=False) +
         scale_colour_manual(values=["#e31837", "#654321", "#002244"], guide=False) +
         ylim(ymin, ymax) +
         xlim(xmin, xmax) +
         coord_fixed() +
         theme_void())
    return p