from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import os
import xgboost as xgb
import pandas as pd
from concurrent.futures import ProcessPoolExecutor, as_completed

from events.get_events import get_shots
from xg.prepare_training_data import prepare_training_data


class ModelTraining:
    num_round = 29
    run_for_input_cols_combos = True
    params = {
              'max_depth': 4,
              'eta': 0.4,
              'objective': 'multi:softprob',
              'num_class': 2
              }

    columns_for_model_training = ['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']

    base_available_columns = ['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']

    columns_to_try = []

    removed_for_now = ["primary_position", "lag_2_is_same_team", "lag_3_is_same_player", "lag_3_is_same_team",
                       "lag_4_is_same_team", "lag_5_is_same_team", "lag_4_is_same_player", "lag_5_is_same_player",
                       "previous_event_type_lag4", "previous_event_type_lag5", "time_since_last_event_lag4",
                       "meters_per_sec_lag4", "lag_2_is_same_player", "time_since_last_event_lag2",
                       "time_since_last_event_lag3", "meters_per_sec_lag3", "previous_event_type_lag3"]

    training_data = None

    def train_model_for_params(self):
        raw_data = get_shots()
        training_data = prepare_training_data(raw_data, self.columns_for_model_training)
        train_xg_model(training_data, self.params, self.num_round)

    def train_models(self):
        raw_data = get_shots()
        if self.run_for_input_cols_combos:
            all_combinations = combs(self.columns_to_try)
            final_combinations = clean_combos(all_combinations, self.base_available_columns)
            for row_combo in final_combinations:
                self.training_data = prepare_training_data(raw_data, row_combo)
                self.run_for_num_round()
        else:
            self.run_for_num_round()

    def parallel_train_models(self, num_workers=4):
        raw_data = get_shots()
        if self.run_for_input_cols_combos:
            all_combinations = combs(self.columns_to_try)
            final_combinations = clean_combos(all_combinations, self.base_available_columns)

            with ProcessPoolExecutor(max_workers=num_workers) as executor:
                futures = [executor.submit(self.run_for_num_round, row_combo, raw_data) for row_combo in
                           final_combinations]

                for future in as_completed(futures):
                    print(f"Completed training for: {future.result()}")
        else:
            self.training_data = prepare_training_data(raw_data, self.columns_for_model_training)
            self.run_for_num_round()

    def run_for_num_round(self, row_combo, raw_data):
        training_data = prepare_training_data(raw_data, row_combo)
        for num_round in range(4, 30, 2):
            print("num round: " , num_round)
            self.run_for_max_depth(training_data, num_round)

    def run_for_max_depth(self, training_data, num_round):
        for max_depth in range(2, 16, 1):
            self.run_for_eta(training_data, num_round, max_depth)

    def run_for_eta(self, training_data, num_round, max_depth):

        for eta in range(2, 11, 1):
            params = {'max_depth': max_depth,
                      "eta": eta / 10,
                      'objective': self.params["objective"],
                      'num_class': self.params["num_class"]
                      }
            train_xg_model(training_data, params, num_round)


def train_xg_model(data, param, num_round):
    model_name = f"{param['max_depth']}_{param['eta']}_{num_round}"
    bst = xgb.train(param, data["training_data"], num_round)
    models_path = os.path.join(os.getcwd(), "xg", "model", f"{model_name}.json")
    bst.save_model(models_path)
    get_and_save_test_results(bst, data, param, num_round)


def get_and_save_test_results(model, data, params, num_round):
    preds = model.predict(data["testing_data"])
    xg = pd.DataFrame(preds,
                      columns=['col1', 'xg'])["xg"]
    mean_sq = mean_squared_error(data["testing_data_correct_result"], xg)
    mean_abs_err = mean_absolute_error(data["testing_data_correct_result"], xg)
    r2 = r2_score(data["testing_data_correct_result"], xg)

    save_test_results(params, num_round, data, mean_sq, mean_abs_err, r2)


def save_test_results(param, num_round, data, mean_sq, mean_abs_err, r2):
    path = os.path.join(os.getcwd(), f"test_results.txt")
    file1 = open(path, "a")  # append mode
    file1.write(f"{data['clist']};{param};{num_round};{mean_sq};{mean_abs_err};{r2}\n")
    file1.close()


def combs(a):
    if not a:
        return [[]]
    sub_combs = combs(a[1:])
    return sub_combs + [comb + [a[0]] for comb in sub_combs]



def clean_combos(combos, base_available_columns):
    final = []
    for combo in combos:
        combo = [*combo, *base_available_columns]
        event_lag2 = False
        event_lag3 = False
        event_lag4 = False
        event_lag5 = False

        same_team_lag1 = False
        same_team_lag2 = False
        same_team_lag3 = False
        same_team_lag4 = False
        same_team_lag5 = False

        same_player_lag1 = False
        same_player_lag2 = False
        same_player_lag3 = False
        same_player_lag4 = False
        same_player_lag5 = False

        mps_lag0 = False
        mps_lag1 = False
        mps_lag2 = False
        mps_lag3 = False
        mps_lag4 = False

        if len(combo) == 0:
            continue

        if len(combo) == 1 and "primary_position" in combo:
            continue
        if len(combo) == 1 and "primary_position_detailed" in combo:
            continue
        if not ("primary_position" in combo) and not ("primary_position_detailed" in combo):
            continue
        if "primary_position" in combo and "primary_position_detailed" in combo:
            continue
        if "previous_event_type_lag2" in combo:
            event_lag2 = True
        if "previous_event_type_lag3" in combo:
            event_lag3 = True
        if "previous_event_type_lag4" in combo:
            event_lag4 = True
        if "previous_event_type_lag5" in combo:
            event_lag5 = True

        if "lag_1_is_same_player" in combo:
            same_player_lag1 = True
        if "lag_2_is_same_player" in combo:
            same_player_lag2 = True
        if "lag_3_is_same_player" in combo:
            same_player_lag3 = True
        if "lag_4_is_same_player" in combo:
            same_player_lag4 = True
        if "lag_5_is_same_player" in combo:
            same_player_lag5 = True

        if "lag_1_is_same_team" in combo:
            same_team_lag1 = True
        if "lag_2_is_same_team" in combo:
            same_team_lag2 = True
        if "lag_3_is_same_team" in combo:
            same_team_lag3 = True
        if "lag_4_is_same_team" in combo:
            same_team_lag4 = True
        if "lag_5_is_same_team" in combo:
            same_team_lag5 = True

        if "meters_per_sec" in combo:
            mps_lag0 = True
        if "meters_per_sec_lag1" in combo:
            mps_lag1 = True
        if "meters_per_sec_lag2" in combo:
            mps_lag2 = True
        if "meters_per_sec_lag3" in combo:
            mps_lag3 = True
        if "meters_per_sec_lag4" in combo:
            mps_lag4 = True

        if event_lag5 and not event_lag4:
            continue
        if event_lag4 and not event_lag3:
            continue
        if event_lag3 and not event_lag2:
            continue

        if same_player_lag5 and not same_player_lag4:
            continue
        if same_player_lag4 and not same_player_lag3:
            continue
        if same_player_lag3 and not same_player_lag2:
            continue
        if same_player_lag2 and not same_player_lag1:
            continue

        if same_team_lag5 and not same_team_lag4:
            continue
        if same_team_lag4 and not same_team_lag3:
            continue
        if same_team_lag3 and not same_team_lag2:
            continue
        if same_team_lag2 and not same_team_lag1:
            continue

        if mps_lag4 and not mps_lag3:
            continue
        if mps_lag3 and not mps_lag2:
            continue
        if mps_lag2 and not mps_lag1:
            continue
        if mps_lag1 and not mps_lag0:
            continue

        final.append(combo)
    return final
