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 = {
        "stay": (0, 0),
        "left": (0, -1),
        "right": (0, 1),
        "up": (-1, 0),
        "down": (1, 0)
    }
    # Jaké území okolo domečku je chráněné
    PROTECTED_AREA = [
        [-1, -1], [-1, 0], [-1, 1],
        [0, -1],  [0, 0],  [0, 1],
        [1, -1],  [1, 0],  [1, 1]
    ]

    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.world_width = self.teams_width * configuration["width_per_team"]
        self.world_height = self.teams_height * configuration["height_per_team"]
        # Parametry k počtu nových vojáků
        self.rounds_total = configuration["initial_remaining_rounds"]
        self.spawn_price = configuration["spawn_price"]
        self.last_spawn = configuration["last_spawn"]
    
    def generate_soldier(self, team_id, soldier_id):
        return {
            "type": "soldier",
            "team": team_id,
            "id": soldier_id,
            "remaining_rounds": 1
        }

    def budget_for_round(self, round_id):
        return (self.last_spawn * self.spawn_price * (round_id + 1)**2) // (self.rounds_total**2)
    
    # Počáteční stav
    def zero_state(self) -> Any:
        world = [
            [
                {
                    "home_for_team": None,
                    "occupied_by_team": None,
                    "protected_for_team": None,
                    "hill": (self.configuration["hills"]
                             [j % self.configuration["height_per_team"]]
                             [i % self.configuration["width_per_team"]]
                             == "x"),
                    "members": []
                }
                for i in range(self.world_width)
            ]
            for j in range(self.world_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):
                world[home_y][home_x]["home_for_team"] = team_id
                world[home_y][home_x]["occupied_by_team"] = team_id
                for x, y in self.PROTECTED_AREA:
                    world[(home_y+y) % self.world_height][(home_x+x) % self.world_width]["protected_for_team"] = team_id
                # Initial stav vojáků:
                for soldier_id in range(max(1, self.budget_for_round(0) // self.spawn_price)):
                    world[home_y][home_x]["members"].append(self.generate_soldier(team_id, soldier_id))
                team_id += 1
                home_x += self.configuration["width_per_team"]
            home_y += self.configuration["height_per_team"]
        
        return {"world": world}

    # 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"
    #     ]
    # }
    # Předpokládám validitu `state`
    # 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
        # 2D array s `True`s tam, kde je hill
        hills = [[False] * self.world_width for _ in range(self.world_height)]
        # 2D array s ID týmu, který zde má chráněné území (okolo domečku)
        # (-1 tam, kde nikdo nemá domeček)
        protected_for = [[-1] * self.world_width for _ in range(self.world_height)]

        y = 0
        for row in state["world"]:
            x = 0
            for field in row:
                hills[y][x] = field["hill"]
                hft = field["home_for_team"]
                if hft is not None:
                    for p in self.PROTECTED_AREA:
                        protected_for[(y + p[0]) % self.world_height][(x + p[1]) % self.world_width] = hft
                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 team_id, move in enumerate(moves):
            if move is not None:
                for member in move["members"]:
                    if member["id"] not in id_positions[team_id]:
                        # Neplatné ID vojáka
                        continue
                    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.world_width)]
            for _ in range(self.world_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]]
                orig_y = id_positions[team_id][soldier_id][0]
                orig_x = id_positions[team_id][soldier_id][1]
                new_y = (orig_y + move_vect[0]) % self.world_height
                new_x = (orig_x + move_vect[1]) % self.world_width
                # Chráněné území nebo skála
                if (protected_for[new_y][new_x] > -1 and protected_for[new_y][new_x] != team_id) \
                   or hills[new_y][new_x]:
                    soldier_lists[orig_y][orig_x][team_id].append(soldier_id)
                else:
                    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_count = len(all_ids[team_id])
            spawn_count = max(1, (self.budget_for_round(round_id) - soldier_count) // self.spawn_price)
            for i in range(spawn_count):
                soldier_lists[home_pos[0]][home_pos[1]][team_id].append(highest_ids[team_id] + i + 1)

        # Vojáci se pomlátí:
        for y in range(self.world_height):
            for x in range(self.world_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_count = len(soldier_lists[y][x][tid])
                    if soldier_count <= team_strengths[-2]:
                        soldier_lists[y][x][tid] = []
                    else:
                        remaining_count = soldier_count - (team_strengths[-2]**2 // soldier_count)
                        soldier_lists[y][x][tid] = soldier_lists[y][x][tid][0:remaining_count]

        # Vygenerujeme výsledek:
        score = [0] * self.teams_count
        new_world = [[] for _ in range(self.world_height)]
        y = 0
        for row in state["world"]:
            x = 0
            for field in row:
                new_field = {
                    "home_for_team": field["home_for_team"],
                    "protected_for_team": field["protected_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_world[y].append(new_field)
                x += 1
            y += 1
        return ({"world": new_world}, 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.world_width for _ in range(self.world_height)]
        # 2D array `True` na chráněných územích ostatních týmů
        protected = [[False] * self.world_width for _ in range(self.world_height)]
        # ID soldierů, se kterými tým už hnul
        ids_moved = {}

        y = 0
        for row in state["world"]:
            x = 0
            for field in row:
                hills[y][x] = field["hill"]
                if field["home_for_team"] is not None and field["home_for_team"] != team_id:
                    for p in self.PROTECTED_AREA:
                        protected[(y + p[0]) % self.world_height][(x + p[1]) % self.world_width] = True
                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]) % self.world_height
            new_position_x = (id_positions[member["id"]][1] + move_vect[1]) % self.world_width
            if protected[new_position_y][new_position_x]:
                warns.append(self.generate_warn(member["id"], "Voják se snaží dostat do domečku cizího týmu. To není povoleno."))
            elif 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["world"]:
            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.world_height // 2) - home_yx[0]
        x_diff = (self.world_width // 2) - home_yx[1]

        # Vycentruj:
        world_personalized = [
            [
                state["world"][(y - y_diff) % self.world_height][(x - x_diff) % self.world_width]
                for x in range(self.world_width)
            ]
            for y in range(self.world_height)
        ]
        
        return {"world": world_personalized}