import os
import sys
import pandas as pd
from datetime import datetime
from settings import RAW_DATA_DIR, DataV3, DATA_V3_SUBVERSION
from src.features.helpers.processing import add_missing_timestamp_values
from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, \
    normalize_according_to_play_direction, check_group_event
from src.features.helpers.processing_v4 import home_has_possession, calculate_team_sitation

week_num = int(sys.argv[1])

data_v3 = DataV3(DATA_V3_SUBVERSION)
save_file_path = data_v3.get_step1_checkpoint_path(week_num)
try:
    clean_df = pd.read_csv(save_file_path)
    save_file_exists = True
except FileNotFoundError:
    save_file_exists = False

if not save_file_exists:
    print("Started loading data")
    play_df = pd.read_csv(os.path.join(RAW_DATA_DIR, 'plays.csv'))
    games_df = pd.read_csv(os.path.join(RAW_DATA_DIR, 'games.csv'))

    week_and_games = games_df[games_df.week == week_num]
    tracking_df = pd.read_csv(os.path.join(RAW_DATA_DIR, f'week{week_num}.csv'))
    print("Data loaded. Start processing timestamps")

    tracking_df = add_missing_timestamp_values(tracking_df)

    games_n_plays_df = play_df.merge(week_and_games, how='inner', on='gameId')
    m_grouped = games_n_plays_df.groupby(['gameId', 'playId'])

    df_t = tracking_df.merge(games_n_plays_df, how='left', on=['gameId', 'playId'])

    # Remove all events without 'pass_forward'
    df_t_grouped = df_t.groupby(['gameId', 'playId'])
    df_t_v3 = df_t.copy().sort_index()
    for name, group in df_t_grouped:
        game_id, play_id = name

        # if group does not contain pass forward, drop it
        if all(group.event != 'pass_forward'):
            df_t_v3 = df_t_v3[(df_t_v3.gameId != game_id) | (df_t_v3.playId != play_id)]

    df_t_v3_s = df_t_v3.sort_values(by=['gameId', 'playId', 'time', 'event'])
    df_t_v3_s = df_t_v3_s.reset_index(drop=True)

    df_t_grouped = df_t_v3_s.groupby(['gameId', 'playId'])

    # remove all values before 'pass_forward'
    print("Removing all values before pass forward event...")
    for name, group in df_t_grouped:
        game_id, play_id = name

        pass_forward_frame_id = group[group.event == 'pass_forward'].index.min() - 1
        remove_start = group.index.min()
        df_t_v3_s = df_t_v3_s.drop(df_t_v3_s.loc[remove_start:pass_forward_frame_id].index)

    pd.options.mode.chained_assignment = None
    gb = df_t_v3_s.groupby(['gameId', 'playId'])

    print('Getting closest players...')
    keep_indices = []
    for name, group in gb:
        game_id, play_id = name
        try:
            event_3rd = group.event.unique()[2]
        except IndexError:
            print('Number of events is < 3, skipping...')
            continue

        situation_df = group[group.event == event_3rd]

        # convert dataframe into series
        ball_row = situation_df[situation_df.team == 'football'].head(1)

        # remove ball
        player_situation_df = situation_df[situation_df.team != 'football']
        try:
            p1, p2 = get_closest_players(player_situation_df, ball_row.x.item(), ball_row.y.item())
        except ValueError:
            print('Value Error raised. This group will be skipped.')
            continue
        p_n_b_indices = get_players_and_ball_indices(group, p1, p2)
        if p_n_b_indices:
            keep_indices.extend(p_n_b_indices)

    clean_df = df_t_v3_s[df_t_v3_s.index.isin(keep_indices)]

    clean_df.to_csv(
        save_file_path,
        index=False
    )

print('Normalize...')
clean_df = normalize_according_to_play_direction(clean_df)

clean_df['homeHasPossession'] = clean_df.apply(
    lambda row: home_has_possession(row), axis=1
)

clean_df['teamSituation'] = clean_df.apply(
    lambda row: calculate_team_sitation(row), axis=1
)

print('Creating features...')
min_df = clean_df[[
    'time', 'x', 'y', 's', 'o', 'dir', 'event', 'team',
    'gameId', 'playId', 'frameId', 'isDefensivePI'
]]
gb_2 = clean_df.groupby(['gameId', 'playId', 'frameId'])

# ball direction and orientation are NaN
calc_df = pd.DataFrame(
    columns=[
        'time',
        'att_def_d', 'att_ball_d', 'def_ball_d',
        'att_s', 'def_s', 'ball_s',
        'att_o', 'def_o',
        'att_dir', 'def_dir',
        'event', 'gameId', 'playId', 'frameId', 'isDefensivePI'
    ]
)
GROUP_SIZE_MINIMUM = 3
for name, group in gb_2:
    game_id, play_id, frameId = name

    if len(group) < GROUP_SIZE_MINIMUM:
        continue

    ball = group[group.teamSituation == 'football'].head(1).squeeze()
    p_att = group[group.teamSituation == 'attacking'].head(1).squeeze()
    p_def = group[group.teamSituation == 'defending'].head(1).squeeze()

    group_row = group.head(1).squeeze()

    group_events = group.event.unique().tolist()

    dict_to_append = {
        'time': group_row.time,
        'att_def_d': calculate_distance(p_att.x, p_att.y, p_def.x, p_def.y),
        'att_ball_d': calculate_distance(p_att.x, p_att.y, ball.x, ball.y),
        'def_ball_d': calculate_distance(p_def.x, p_def.y, ball.x, ball.y),
        'att_s': p_att.s, 'def_s': p_def.s, 'ball_s': ball.s,
        'att_a': p_att.a, 'def_a': p_def.a, 'ball_a': ball.a,
        'att_o': p_att.o, 'def_o': p_def.o,
        'att_dir': p_att.dir, 'def_dir': p_def.dir,
        'event': group_row.event,
        'pass_arrived': check_group_event(group_events, 'pass_arrived'),
        'pass_outcome_caught': check_group_event(group_events, 'pass_outcome_caught'),
        'tackle': check_group_event(group_events, 'tackle'),
        'first_contact': check_group_event(group_events, 'first_contact'),
        'pass_outcome_incomplete': check_group_event(group_events, 'pass_outcome_incomplete'),
        'out_of_bounds': check_group_event(group_events, 'out_of_bounds'),
        'week': week_num,
        'gameId': group_row.gameId,
        'playId': group_row.playId,
        'frameId': group_row.frameId,
        'isDefensivePI': group_row.isDefensivePI
    }

    calc_df = calc_df.append(
        dict_to_append,
        ignore_index=True
    )

print("Saving data...")
calc_df.to_csv(
    data_v3.get_step1_end_path(week_num),
    index=False
)

print(f'End time: {datetime.now().strftime("%H:%M:%S")}')

[((13, 10, 13, 36), 'settings.DataV3', 'DataV3', ({(13, 17, 13, 35): 'DATA_V3_SUBVERSION'}, {}), '(DATA_V3_SUBVERSION)', False, 'from settings import RAW_DATA_DIR, DataV3, DATA_V3_SUBVERSION\n'), ((98, 11, 98, 58), 'src.features.helpers.processing_v3.normalize_according_to_play_direction', 'normalize_according_to_play_direction', ({(98, 49, 98, 57): 'clean_df'}, {}), '(clean_df)', False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((116, 10, 125, 1), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((16, 15, 16, 42), 'pandas.read_csv', 'pd.read_csv', ({(16, 27, 16, 41): 'save_file_path'}, {}), '(save_file_path)', True, 'import pandas as pd\n'), ((30, 18, 30, 59), 'src.features.helpers.processing.add_missing_timestamp_values', 'add_missing_timestamp_values', ({(30, 47, 30, 58): 'tracking_df'}, {}), '(tracking_df)', False, 'from src.features.helpers.processing import add_missing_timestamp_values\n'), ((23, 26, 23, 65), 'os.path.join', 'os.path.join', ({(23, 39, 23, 51): 'RAW_DATA_DIR', (23, 53, 23, 64): '"""plays.csv"""'}, {}), "(RAW_DATA_DIR, 'plays.csv')", False, 'import os\n'), ((24, 27, 24, 66), 'os.path.join', 'os.path.join', ({(24, 40, 24, 52): 'RAW_DATA_DIR', (24, 54, 24, 65): '"""games.csv"""'}, {}), "(RAW_DATA_DIR, 'games.csv')", False, 'import os\n'), ((27, 30, 27, 79), 'os.path.join', 'os.path.join', ({(27, 43, 27, 55): 'RAW_DATA_DIR', (27, 57, 27, 78): 'f"""week{week_num}.csv"""'}, {}), "(RAW_DATA_DIR, f'week{week_num}.csv')", False, 'import os\n'), ((86, 24, 86, 67), 'src.features.helpers.processing_v3.get_players_and_ball_indices', 'get_players_and_ball_indices', ({(86, 53, 86, 58): 'group', (86, 60, 86, 62): 'p1', (86, 64, 86, 66): 'p2'}, {}), '(group, p1, p2)', False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((101, 16, 101, 40), 'src.features.helpers.processing_v4.home_has_possession', 'home_has_possession', ({(101, 36, 101, 39): 'row'}, {}), '(row)', False, 'from src.features.helpers.processing_v4 import home_has_possession, calculate_team_sitation\n'), ((105, 16, 105, 44), 'src.features.helpers.processing_v4.calculate_team_sitation', 'calculate_team_sitation', ({(105, 40, 105, 43): 'row'}, {}), '(row)', False, 'from src.features.helpers.processing_v4 import home_has_possession, calculate_team_sitation\n'), ((143, 21, 143, 75), 'src.features.helpers.processing_v3.calculate_distance', 'calculate_distance', ({(143, 40, 143, 47): 'p_att.x', (143, 49, 143, 56): 'p_att.y', (143, 58, 143, 65): 'p_def.x', (143, 67, 143, 74): 'p_def.y'}, {}), '(p_att.x, p_att.y, p_def.x, p_def.y)', False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((144, 22, 144, 74), 'src.features.helpers.processing_v3.calculate_distance', 'calculate_distance', ({(144, 41, 144, 48): 'p_att.x', (144, 50, 144, 57): 'p_att.y', (144, 59, 144, 65): 'ball.x', (144, 67, 144, 73): 'ball.y'}, {}), '(p_att.x, p_att.y, ball.x, ball.y)', False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((145, 22, 145, 74), 'src.features.helpers.processing_v3.calculate_distance', 'calculate_distance', ({(145, 41, 145, 48): 'p_def.x', (145, 50, 145, 57): 'p_def.y', (145, 59, 145, 65): 'ball.x', (145, 67, 145, 73): 'ball.y'}, {}), '(p_def.x, p_def.y, ball.x, ball.y)', False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((151, 24, 151, 71), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(151, 42, 151, 54): 'group_events', (151, 56, 151, 70): '"""pass_arrived"""'}, {}), "(group_events, 'pass_arrived')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((152, 31, 152, 85), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(152, 49, 152, 61): 'group_events', (152, 63, 152, 84): '"""pass_outcome_caught"""'}, {}), "(group_events, 'pass_outcome_caught')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((153, 18, 153, 59), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(153, 36, 153, 48): 'group_events', (153, 50, 153, 58): '"""tackle"""'}, {}), "(group_events, 'tackle')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((154, 25, 154, 73), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(154, 43, 154, 55): 'group_events', (154, 57, 154, 72): '"""first_contact"""'}, {}), "(group_events, 'first_contact')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((155, 35, 155, 93), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(155, 53, 155, 65): 'group_events', (155, 67, 155, 92): '"""pass_outcome_incomplete"""'}, {}), "(group_events, 'pass_outcome_incomplete')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((156, 25, 156, 73), 'src.features.helpers.processing_v3.check_group_event', 'check_group_event', ({(156, 43, 156, 55): 'group_events', (156, 57, 156, 72): '"""out_of_bounds"""'}, {}), "(group_events, 'out_of_bounds')", False, 'from src.features.helpers.processing_v3 import get_closest_players, get_players_and_ball_indices, calculate_distance, normalize_according_to_play_direction, check_group_event\n'), ((175, 19, 175, 33), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n')]