from typing import Any, List, Tuple, Union, Optional

logic_by_mode = {}
def add_logic(cls):
    logic_by_mode[cls.__name__.lower()] = cls

class Logic:
    def __init__(self, teams_count: int, configuration: Any):
        self.teams_count = teams_count
        self.configuration = configuration

    def zero_state(self) -> Any:
        return {}

    def step(self, state: Any, moves: List[Optional[Any]], round_id: int) -> Tuple[Any, List[int]]:
        return {}, [0] * self.teams_count # new_state, add_points for each team

    def validate_move(self, state: Any, team_id: int, move: Any, round_id: int) -> Union[None, Any]:
        return None # Bez chyby
        # return {"status": "warning", ... } # Drobná chyba, ale tah provedu
        # throw Exception("Chybí povinná ...") # Zásadní chyba, tah neuznán 
                # Když používají našeho klienta, tak by se toto nemělo stát

    def personalize_state(self, state: Any, team_id: int, round_id: int) -> Any:
        return state

@add_logic
class Occupy(Logic):
    MOVE_VECTORS = {
        None: (0, 0),
        "left": (0, -1),
        "right": (0, 1),
        "up": (-1, 0),
        "down": (1, 0)
    }

    def __init__(self, teams_count: int, configuration: Any):
        super().__init__(teams_count, configuration)
        self.teams_width = configuration["teams_width"]
        self.teams_height = configuration["teams_height"]
        assert(teams_count == self.teams_width * self.teams_height)
        self.map_width = self.teams_width * configuration["width_per_team"]
        self.map_height = self.teams_height * configuration["height_per_team"]
    
    def generate_soldier(self, team_id, soldier_id):
        return {
            "type": "soldier",
            "team": team_id,
            "id": soldier_id,
            "remaining_rounds": 1
        }
    
    # Počáteční stav
    def zero_state(self) -> Any:
        # TODO: Zatím negeneruju "hills"
        map = [
            [
                {
                    "home_for_team": None,
                    "occupied_by_team": None,
                    "hill": False,
                    "members": []
                }
                for _ in range(self.map_width)
            ]
            for _ in range(self.map_height)
        ]
        team_id = 0
        home_y = self.configuration["height_per_team"] // 2
        for _ in range(self.teams_height):
            home_x = self.configuration["width_per_team"] // 2
            for _ in range(self.teams_width):
                map[home_y][home_x]["home_for_team"] = team_id
                map[home_y][home_x]["occupied_by_team"] = team_id
                map[home_y][home_x]["members"].append(self.generate_soldier(team_id, 0))
                team_id += 1
                home_x += self.configuration["width_per_team"]
            home_y += self.configuration["height_per_team"]
        
        return {"map": map}

    # Zatím očekávám, že moves má takovéto prvky:
    # {
    #     team_id: <int>
    #     members: [
    #         id: <int>
    #         action: <Optional[str]>
    #             # None -> čekej
    #             # "left", "right", "up", "down"
    #     ]
    # }
    # VALIDITU MOVES NEOVĚŘUJU
    # Metoda funguje velmi primitivně po krocích:
    # * zaznamená se poloha objektů na mapě
    # * vojáky se pohne podle moves
    # * přidají se vojáci v homes
    # * odstraní se vojáci podle "soubojů"
    # * vygeneruje se json se správnými hodnotami "occupied" a spočítá se skóre
    def step(self, state: Any, moves: List[Optional[Any]], round_id: int) -> Tuple[Any, List[int]]:
        # souřadnice IDs pro všechny týmy
        id_positions = [{} for _ in range(self.teams_count)]
        # pohyby IDs pro všechny týmy
        id_moves = [{} for _ in range(self.teams_count)]
        # všechna IDs týmů
        all_ids = [[] for _ in range(self.teams_count)]
        # nejvyšší ID každého týmu
        highest_ids = [0] * self.teams_count
        # homes jednotlivých týmů
        home_positions = [None] * self.teams_count

        y = 0
        for row in state["map"]:
            x = 0
            for field in row:
                if field["home_for_team"] is not None:
                    home_positions[field["home_for_team"]] = [y, x]
                for member in field["members"]:
                    highest_ids[member["team"]] = max(highest_ids[member["team"]], member["id"])
                    all_ids[member["team"]].append(member["id"])
                    id_positions[member["team"]][member["id"]] = [y, x]
                    id_moves[member["team"]][member["id"]] = None
                x += 1
            y += 1

        for move in moves:
            team_id = move["team_id"]
            for member in move["members"]:
                id_moves[team_id][member["id"]] = member["action"]
        
        # Místo mapy tabulka seznamů vojáků pro každý tým
        soldier_lists = [
            [[[] for _ in range(self.teams_count)] for _ in range(self.map_width)]
            for _ in range(self.map_height)
        ]
        for team_id in range(self.teams_count):
            for soldier_id in all_ids[team_id]:
                move_vect = self.MOVE_VECTORS[id_moves[team_id][soldier_id]]
                new_y = (id_positions[team_id][soldier_id][0] + move_vect[0]) % self.map_height
                new_x = (id_positions[team_id][soldier_id][1] + move_vect[1]) % self.map_width
                soldier_lists[new_y][new_x][team_id].append(soldier_id)
        
        # Přidáme vojáky v homes:
        for team_id in range(self.teams_count):
            home_pos = home_positions[team_id]
            soldier_lists[home_pos[0]][home_pos[1]][team_id].append(highest_ids[team_id] + 1)

        # Vojáci se pomlátí:
        for y in range(self.map_height):
            for x in range(self.map_width):
                team_strengths = [len(soldier_lists[y][x][tid]) for tid in range(self.teams_count)]
                team_strengths.sort()
                for tid in range(self.teams_count):
                    soldier_lists[y][x][tid] = soldier_lists[y][x][tid][team_strengths[-2]:]

        # Vygenerujeme výsledek:
        score = [0] * self.teams_count
        new_map = [[] for _ in range(self.map_height)]
        y = 0
        for row in state["map"]:
            x = 0
            for field in row:
                new_field = {
                    "home_for_team": field["home_for_team"],
                    "occupied_by_team": field["occupied_by_team"],
                    "hill": field["hill"],
                    "members": []
                }
                for team_id in range(self.teams_count):
                    for soldier_id in soldier_lists[y][x][team_id]:
                        new_field["members"].append(self.generate_soldier(team_id, soldier_id))
                        new_field["occupied_by_team"] = team_id
                if new_field["occupied_by_team"] is not None:
                    score[new_field["occupied_by_team"]] += 1
                new_map[y].append(new_field)
                x += 1
            y += 1
        return ({"map": new_map}, score)
    
    def generate_warn(self, member_id, description):
        return {
            "id": member_id,
            "description": description
        }

    # Předpokládám validní `state`, vadit může
    # * "hill"
    # * nevalidní ID
    # * více tahů jedním soldierem
    def validate_move(self, state: Any, team_id: int, move: Any, round_id: int) -> Union[None, Any]:
        # nakumulované warningy
        warns = []
        # souřadnice IDs daného týmu
        id_positions = {}
        # 2D array s `True`s tam, kde je hill
        hills = [[False] * self.map_width for _ in range(self.map_height)]
        # ID soldierů, se kterými tým už hnul
        ids_moved = {}

        y = 0
        for row in state["map"]:
            x = 0
            for field in row:
                hills[y][x] = field["hill"]
                for member in field["members"]:
                    if member["team"] != team_id:
                        continue
                    id_positions[member["id"]] = [y, x]
                x += 1
            y += 1
        
        for member in move["members"]:
            if member["id"] not in id_positions:
                warns.append(self.generate_warn(member["id"], "Neplatné ID!"))
                continue
            if member["id"] in ids_moved:
                warns.append(self.generate_warn(member["id"], "Duplicitní pohyb vojákem. Jen jeho poslední pohyb se počítá."))
                continue
            ids_moved[member["id"]] = True
            if member["action"] is None:
                continue
            move_vect = self.MOVE_VECTORS[member["action"]]
            new_position_y = id_positions[member["id"]][0] + move_vect[0]
            new_position_x = id_positions[member["id"]][1] + move_vect[1]
            if hills[new_position_y][new_position_x]:
                warns.append(self.generate_warn(member["id"], "Voják se snaží vylézt na skálu, což nejde."))
        
        if len(warns) == 0:
            return None
        else:
            return {"status": "warning", "members": warns}

    def home_position(self, state: Any, team_id):
        y = 0
        for row in state["map"]:
            x = 0
            for field in row:
                if field["home_for_team"] == team_id:
                    return [y, x]
                x += 1
            y += 1

    # Akorát vycentruju home týmu.
    # TODO: Chceme přečíslovat týmy ("můj tým je vždy 0")?
    def personalize_state(self, state: Any, team_id: int, round_id: int) -> Any:
        # Zjisti pozici domečku tohoto týmu:
        home_yx = self.home_position(state, team_id)
        y_diff = (self.map_height // 2) - home_yx[0]
        x_diff = (self.map_width // 2) - home_yx[1]

        # Vycentruj:
        map_personalized = [
            [
                state["map"][(y - y_diff) % self.map_height][(x - x_diff) % self.map_width]
                for x in range(self.map_width)
            ]
            for y in range(self.map_height)
        ]
        
        return {"map": map_personalized}