import os

import pandas as pd
from matplotlib import pyplot as plt
import xgboost as xgb
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error

from visualizations.pitch import draw_pitch
from visualizations.shots import draw_shots


def run_example():

    X, chosen_shots_dmatrix, chosen_shots = load_and_prepare_data()

    predictions = evaluate_shots(chosen_shots_dmatrix)

    shots, chosen_shots = prepare_shots(predictions, X, chosen_shots)
    print_shots_stats(shots, chosen_shots)
    shots = shots.to_dict("records")
    draw_shot_map(shots)



def load_and_prepare_data():
    path = r'./chosen_shots'

    chosen_shots = pd.read_csv(os.path.join(path, "haaland_shots.csv"), sep=",")
    X = chosen_shots[['py', 'px', 'previous_event_type', 'shot_body_part_type', 'shot_after_corner',
                                  'time_since_last_event', 'meters_per_sec', 'primary_position_detailed',
                                  'shot_after_free_kick', 'previous_event_type_lag2', 'meters_per_sec_lag2',
                                  'meters_per_sec_lag1', 'lag_1_is_same_player', 'lag_1_is_same_team',
                                  'time_since_last_event_lag1', 'meters_per_sec_lag3', 'time_since_last_event_lag3',
                      "goal"]]
    chosen_shots_dmatrix = xgb.DMatrix(X[['py', 'px', 'previous_event_type', 'shot_body_part_type', 'shot_after_corner',
                                  'time_since_last_event', 'meters_per_sec', 'primary_position_detailed',
                                  'shot_after_free_kick', 'previous_event_type_lag2', 'meters_per_sec_lag2',
                                  'meters_per_sec_lag1', 'lag_1_is_same_player', 'lag_1_is_same_team',
                                  'time_since_last_event_lag1', 'meters_per_sec_lag3', 'time_since_last_event_lag3']])
    return X, chosen_shots_dmatrix, chosen_shots


def evaluate_shots(shots):
    model = xgb.Booster()
    model.load_model("./chosen_model/4_0.4_29.json")

    predictions = model.predict(shots)

    return predictions


def prepare_shots(predictions, X, chosen_shots):

    evaluated_shots_pd = pd.DataFrame(predictions,
                                      columns=['col1',
                                               'xg'])

    """extract the evaluation and save as pandas DF"""
    shots_with_xg_pd = X.assign(xg=evaluated_shots_pd["xg"])[
        ["px", "py", "xg", "goal"]]

    return shots_with_xg_pd, chosen_shots


def print_shots_stats(shots, chosen_shots):
    xgsum = 0.0
    goals = 0
    wyscout_xg = 0.0


    wy_xg = chosen_shots.shot_xg
    wy_goal = chosen_shots.goal
    wy_mean_sq = mean_squared_error(wy_goal, wy_xg)
    wy_r2 = r2_score(wy_goal, wy_xg)
    wy_mean_abs = mean_absolute_error(wy_goal, wy_xg)

    my_xg = shots.xg
    my_mean_sq = mean_squared_error(wy_goal, my_xg)
    my_r2 = r2_score(wy_goal, my_xg)
    my_mean_abs = mean_absolute_error(wy_goal, my_xg)

    print("wy_mean_sq: ", wy_mean_sq)
    print("wy_r2: ", wy_r2)
    print("wy_mean_abs: ", wy_mean_abs)

    print("my_mean_sq: ", my_mean_sq)
    print("my_r2: ", my_r2)
    print("my_mean_abs: ", my_mean_abs)

    shots = shots.to_dict("records")
    chosen_shots = chosen_shots.to_dict("records")

    for shot in shots:
        xgsum += shot["xg"]
    for chosen_shot in chosen_shots:
        goals += chosen_shot["goal"]
        wyscout_xg += chosen_shot["shot_xg"]

    print(f"total_xg: {xgsum}")
    print(f"wyscout_xg: {wyscout_xg}")
    print(f"goals: {goals}")


def draw_shot_map(shots):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    fig, ax = draw_pitch(ax, fig)
    draw_shots(shots, ax)

    fig.savefig('my_plot.png')
