from mesa import Model

# TODO activate player with ball
from mesa.time import RandomActivation
from mesa.space import SingleGrid
from mesa.datacollection import DataCollector

from g1_utils.utils import (
    calculate_distance_two_points,
    train_linear_two_points,
    apply_linear_model,
    check_if_folder_exists,
)
from g2_player.FootballPlayer import FootballPlayer

from pathlib import Path
import os
import datetime
import re
import logging
import random
import numpy as np
import yaml
import matplotlib.pyplot as plt
from matplotlib import colors

logger = logging.getLogger("FootballModel")


class FootballModel(Model):
    def __init__(self, players_per_team, width, height, game_length):
        # load parameters
        self.parameters = yaml.full_load(open("general_parameters.yaml"))

        # generate game id
        self.game_id = self.__generate_game_id()

        self.players_per_team = players_per_team

        self.result = {"A": 0, "B": 0}
        self.game_length = game_length

        self.grid = SingleGrid(width, height, False)
        self.distance_model = self.__train_distance_model()
        self.goal_coordinates = self.__determine_goal_midpoints()
        self.penalty_box_measures = self.__determine_penalty_box()
        self.scoring_probabilities = self.__determine_scoring_probabilities()

        # TODO: change so it only activates player with ball
        self.schedule = RandomActivation(self)

        ## Create Agents
        # generate_ids
        self.ids_team_a = ["A_" + str(id) for id in range(self.players_per_team)]
        self.ids_team_b = ["B_" + str(id) for id in range(self.players_per_team)]
        self.ids = self.ids_team_a + self.ids_team_b

        self.id_dict = {}
        for i in range(len(self.ids)):
            # Create dict linking real ids to used ids
            self.id_dict[self.ids[i]] = i

            player = FootballPlayer(
                unique_id=self.ids[i], team=self.ids[i][0], model=self
            )

            self.schedule.add(player)

            ## Place agent on grid
            # TODO this should be later be improved to create a formation of some sorts
            empty_place = self.grid.find_empty()
            self.grid.place_agent(player, empty_place)
            logger.info("Agent " + str(self.ids[i]) + " placed in " + str(empty_place))

        # Give the ball to a player
        ball_options = self.determine_furthest_back_per_team()
        kickoff_team = random.choice(list(ball_options.keys()))
        chosen_player = ball_options[kickoff_team]

        logger.info("Team " + str(kickoff_team) + " wins the coin toss")
        self.schedule.agents[self.id_dict[chosen_player]].has_ball = True
        logger.info("Ball given to player " + str(chosen_player))

        self.plot_grid(minute=0, subtitle="Game kickoff")

    def __generate_game_id(self):
        """Generate game id depending on current time
        
        Returns:
            str: Game id
        """
        time_str = str(datetime.datetime.now())

        replacements = {replace: "" for replace in ["-", " ", ":", "."]}
        replacements = dict((re.escape(k), v) for k, v in replacements.items())

        pattern = re.compile("|".join(replacements.keys()))
        game_id = pattern.sub(lambda x: replacements[re.escape(x.group(0))], time_str)

        return game_id

    def __train_distance_model(self):
        """Calculates the maximum distance in the pitch and trains a linear model
            to determine passing accuracy-distance ratio

        Returns:
            dict: dict with slope and intercept
        """

        max_distance = calculate_distance_two_points(
            (0, 0), (self.grid.width, self.grid.height)
        )

        output_dict = train_linear_two_points(
            (1, self.parameters["min_distance_passing_perc"]),
            (max_distance, self.parameters["max_distance_passing_perc"]),
        )

        return output_dict

    def __determine_penalty_box(self):
        """Calculate measures needed to draw penalty boxes
        
        Returns:
            dict: Dict with four required lengths
            #TODO explain this somehow
        """
        penalty_box_height = (16.5 / 105) * self.grid.height
        penalty_box_width = (40 / 68) * self.grid.width
        penalty_box_buffer = (self.grid.width - penalty_box_width) / 2

        measures_penalty_boxes = {
            "length_a": round(penalty_box_buffer, 2),
            "length_b": round(penalty_box_buffer + penalty_box_width, 2),
            "length_c": round(penalty_box_height, 2),
            "length_d": round(self.grid.height - penalty_box_height, 2),
        }

        return measures_penalty_boxes

    def __determine_goal_midpoints(self):
        goal_width = (self.grid.width - 1) / 2

        return {"A": (goal_width, -0.5), "B": (goal_width, self.grid.height - 0.5)}

    def __determine_scoring_probabilities(self):
        """Determine scoring probabilities from every point on the pitch
        
        Returns:
            dict: dict with scoring probabilities for each position
        """
        # calculate distance to opposing goal for each team
        dist_A = {
            (x, y): calculate_distance_two_points(self.goal_coordinates["B"], (x, y))
            for a, x, y in self.grid.coord_iter()
        }

        dist_B = {
            (x, y): calculate_distance_two_points(self.goal_coordinates["A"], (x, y))
            for a, x, y in self.grid.coord_iter()
        }

        # train probabilities model
        min_dist = dist_A[min(dist_A, key=dist_A.get)]
        max_dist = dist_A[max(dist_A, key=dist_A.get)]

        # TODO this should be adjustable in yaml
        point_min = (min_dist, self.parameters["min_distance_shooting_perc"])
        point_max = (max_dist, self.parameters["max_distance_shooting_perc"])

        dist_model = train_linear_two_points(point_min, point_max)

        prob_A = {
            pos: apply_linear_model(dist_A[pos], dist_model) for pos in dist_A.keys()
        }

        prob_B = {
            pos: apply_linear_model(dist_B[pos], dist_model) for pos in dist_B.keys()
        }

        output_dict = {
            "A": prob_A,
            "B": prob_B,
        }

        return output_dict

    def update_result(self):
        """Count number of goals scored per team and update result
        """
        team_A_goals = sum(
            [
                player.goals_scored
                for player in self.schedule.agents
                if player.team == "A"
            ]
        )
        team_B_goals = sum(
            [
                player.goals_scored
                for player in self.schedule.agents
                if player.team == "B"
            ]
        )

        self.result = {"A": team_A_goals, "B": team_B_goals}
        logger.info(
            "Score stands: A - "
            + str(team_A_goals)
            + " : "
            + str(team_B_goals)
            + " - B"
        )

    def determine_furthest_back_per_team(self):
        """Determine player that is furthest back for each team
        
        Returns:
            dict: Dict with teams as keys and player id as value
        """
        pos_dict = {agent.unique_id: agent.pos[1] for agent in self.schedule.agents}

        team_A = {key: pos_dict[key] for key in pos_dict.keys() if key[0] == "A"}
        team_B = {key: pos_dict[key] for key in pos_dict.keys() if key[0] == "B"}

        # Calcualte min value for positions in each team
        min_A = team_A[min(team_A, key=team_A.get)]
        min_B = team_B[max(team_B, key=team_B.get)]

        # # Get players located in such value
        players_A = [key for key in team_A.keys() if pos_dict[key] == min_A]
        players_B = [key for key in team_B.keys() if pos_dict[key] == min_B]

        # Return random player from list
        output_dict = {
            "A": random.choice(players_A),
            "B": random.choice(players_B),
        }

        return output_dict

    def who_has_ball(self):
        """Determine which player has the ball
        
        Returns:
            int: Index of player who has ball
        """
        has_ball = [agent.has_ball for agent in self.schedule.agents]
        index_has_ball = [i for i, x in enumerate(has_ball) if x]
        if len(index_has_ball) == 1:
            id_has_ball = [
                key
                for key in self.id_dict.keys()
                if self.id_dict[key] == index_has_ball[0]
            ]
            return {
                "index": index_has_ball[0],
                "id": id_has_ball[0],
                "player": self.schedule.agents[index_has_ball[0]],
            }

        elif len(index_has_ball) == 0:
            logger.error("No one has ball")

        else:
            logger.error("More than one player has the ball")

    def plot_grid(self, minute=1, subtitle="No subtitle specified", save_fig=False):
        """Plot game grid
        
        Args:
            subtitle (str, optional): Subtitle for plot. 
                Defaults to "No subtitle specified".
            save_fig (bool, optional): If True, save figure in determined folder
                structure. Defaults to False.
        """
        # gather ball position
        ball_coord = self.who_has_ball()["player"].pos
        ball_coord_x = ball_coord[0]
        ball_coord_y = ball_coord[1]
        # gather goal_posts_position
        goal_a_x = self.goal_coordinates["A"][0]
        goal_a_y = self.goal_coordinates["A"][1]
        goal_b_x = self.goal_coordinates["B"][0]
        goal_b_y = self.goal_coordinates["B"][1]

        teams = np.zeros((self.grid.width, self.grid.height))

        for cell in self.grid.coord_iter():
            cell_content, x, y = cell
            if cell_content != None:
                if cell_content.team == "A":
                    teams[x][y] = 1
                if cell_content.team == "B":
                    teams[x][y] = 2

        # define colors
        cmap = colors.ListedColormap(["forestgreen", "forestgreen", "forestgreen"])
        bounds = [0, 5, 1.5, 10]
        norm = colors.BoundaryNorm(bounds, cmap.N)

        fig, ax = plt.subplots()
        im = ax.imshow(teams, interpolation="nearest", cmap=cmap)

        # plot pitch lines
        # TODO maybe move this to utils?
        plt.axvline((self.grid.height / 2) - 0.5, color="white", zorder=1)
        plt.axvline(-0.5, color="black", zorder=3)
        plt.axvline(self.grid.height - 0.5, color="black", zorder=3)

        plt.scatter(
            (self.grid.height / 2) - 0.5,
            (self.grid.width / 2) - 0.5,
            marker="o",
            s=3300,
            facecolors="none",
            edgecolors="white",
            zorder=1,
        )

        a_1_1, a_1_2 = (
            [-0.5, self.penalty_box_measures["length_c"] - 0.5],
            [
                self.penalty_box_measures["length_a"] - 0.5,
                self.penalty_box_measures["length_a"] - 0.5,
            ],
        )

        plt.plot(a_1_1, a_1_2, c="white", zorder=1)

        a_2_1, a_2_2 = (
            [
                self.penalty_box_measures["length_c"] - 0.5,
                self.penalty_box_measures["length_c"] - 0.5,
            ],
            [
                self.penalty_box_measures["length_a"] - 0.5,
                self.penalty_box_measures["length_b"] - 0.5,
            ],
        )

        plt.plot(a_2_1, a_2_2, c="white", zorder=1)

        a_3_1, a_3_2 = (
            [-0.5, self.penalty_box_measures["length_c"] - 0.5],
            [
                self.penalty_box_measures["length_b"] - 0.5,
                self.penalty_box_measures["length_b"] - 0.5,
            ],
        )

        plt.plot(a_3_1, a_3_2, c="white", zorder=1)

        b_1_1, b_1_2 = (
            [self.penalty_box_measures["length_d"] - 0.5, self.grid.height - 0.5,],
            [
                self.penalty_box_measures["length_a"] - 0.5,
                self.penalty_box_measures["length_a"] - 0.5,
            ],
        )

        plt.plot(b_1_1, b_1_2, c="white", zorder=1)

        b_2_1, b_2_2 = (
            [
                self.penalty_box_measures["length_d"] - 0.5,
                self.penalty_box_measures["length_d"] - 0.5,
            ],
            [
                self.penalty_box_measures["length_a"] - 0.5,
                self.penalty_box_measures["length_b"] - 0.5,
            ],
        )

        plt.plot(b_2_1, b_2_2, c="white", zorder=1)

        b_3_1, b_3_2 = (
            [self.penalty_box_measures["length_d"] - 0.5, self.grid.height - 0.5,],
            [
                self.penalty_box_measures["length_b"] - 0.5,
                self.penalty_box_measures["length_b"] - 0.5,
            ],
        )

        plt.plot(b_3_1, b_3_2, c="white", zorder=1)

        ax.set_xticks(np.arange(self.grid.height))
        ax.set_yticks(np.arange(self.grid.width))

        # plot player labels
        for player in self.schedule.agents:
            if player.team == "A":
                back_color = "white"
                color = "black"
            else:
                back_color = "black"
                color = "white"

            plt.scatter(
                player.pos[1], player.pos[0], marker="s", s=450, c=color, zorder=2
            )
            plt.annotate(
                player.unique_id,
                (player.pos[1] - 0.32, player.pos[0] + 0.1),
                c=back_color,
                size=8,
                weight="bold",
                zorder=2,
            )

        # plot ball
        plt.scatter(
            ball_coord_y,
            ball_coord_x,
            facecolors="none",
            edgecolors="orangered",
            marker="s",
            s=450,
            linewidth=3,
            zorder=10,
        )
        plt.scatter(goal_a_y, goal_a_x, c="orange", marker="d", s=250, zorder=4)
        plt.scatter(goal_b_y, goal_b_x, c="orange", marker="d", s=250, zorder=4)

        title = (
            (
                "A - "
                + str(self.result["A"])
                + " : "
                + str(self.result["B"])
                + " - B ("
                + str(minute)
                + "')"
                + "\n"
            )
            + "\n"
            + subtitle
        )

        plt.title(title, fontweight="bold")

        # save graph if desired
        if save_fig:
            if check_if_folder_exists(self.game_id, "Images") == False:
                os.mkdir(Path("Images") / self.game_id)
                logger.info("Created folder for game")
            plt.savefig(
                Path("Images/" + str(self.game_id))
                / (str(self.game_id) + "_" + "{:02d}".format(int(minute)) + ".png"),
                bbox_inches="tight",
            )
            logger.info("Image saved")

        plt.show()

    def step(self, minute=1, plot_outcome=True, save_plots=False):
        """Activate player that has the ball
        
        Args:
            plot_outcome (bool, optional): Should outcome of play be plotted.
                 Defaults to True.
        """
        plot_subtitle = self.who_has_ball()["player"].step()
        if plot_outcome:
            self.plot_grid(minute=minute, subtitle=plot_subtitle, save_fig=save_plots)

    def simulate_whole_game(self, plot_outcome=True, save_plots=False):
        """Simulate the length of the entire game
        
        Args:
            plot_outcome (bool, optional): Should outcome of play be plotted.
                 Defaults to True.
        """
        if save_plots:
            self.plot_grid(minute=0, subtitle="KICKOFF", save_fig=True)

        for i in range(self.game_length):
            self.step(minute=i + 1, plot_outcome=plot_outcome, save_plots=save_plots)
        self.final_result = self.result
        logger.info("FINAL RESULT")
        logger.info(
            "A - "
            + str(self.final_result["A"])
            + " : "
            + str(self.final_result["B"])
            + " - B"
        )

        self.plot_grid(
            minute=self.game_length, subtitle="FINAL WHISTLE!", save_fig=save_plots
        )
        return self.final_result

