SKSP_2022_strategicka_hra/klient/play.py

#!/usr/bin/env python3
from __future__ import annotations
from typing import Callable, Iterable, Optional, List, Set
import collections
import enum
import json
import sys


class Action(enum.Enum):
    UP = enum.auto()
    DOWN = enum.auto()
    LEFT = enum.auto()
    RIGHT = enum.auto()
    STAY = enum.auto()

    def invert(self) -> Action:
        if self == Action.UP:
            return Action.DOWN
        if self == Action.DOWN:
            return Action.UP
        if self == Action.LEFT:
            return Action.RIGHT
        if self == Action.RIGHT:
            return Action.LEFT
        return Action.STAY

    def to_json(self) -> Optional[str]:
        if self == Action.STAY:
            return None
        return self.name.lower()


class State:
    def __init__(self, state: dict) -> None:
        self.world = parse_world(state["state"]["world"])
        self.my_team_id: int = state["team_id"]
        self.round_number: int = state["round"]


state: State


class Member:
    def __init__(self, field: Field, team: int, id: int) -> None:
        self.field = field
        self.team = team
        self.id = id
        self.action = Action.STAY


class Field:
    members: List[Member]

    def __init__(self, i: int, j: int, hill: bool,
                 home_for_team: Optional[int],
                 occupied_by_team: Optional[int]) -> None:
        self.i = i
        self.j = j
        self.hill = hill
        self.home_for_team = home_for_team
        self.occupied_by_team = occupied_by_team

    def __eq__(self, other: Field) -> bool:
        return (self.i, self.j) == (other.i, other.j)

    def __hash__(self) -> int:
        return hash((self.i, self.j))

    def get_neighbour_field(self, action: Action) -> Field:
        if action == Action.UP:
            neighbour_i = self.i - 1
            neighbour_j = self.j
        elif action == Action.DOWN:
            neighbour_i = self.i + 1
            neighbour_j = self.j
        elif action == Action.LEFT:
            neighbour_i = self.i
            neighbour_j = self.j - 1
        elif action == Action.RIGHT:
            neighbour_i = self.i
            neighbour_j = self.j + 1
        else:
            neighbour_i = self.i
            neighbour_j = self.j
        # ensure coords are in bounds
        neighbour_i %= len(state.world)
        neighbour_j %= len(state.world[0])
        return state.world[neighbour_i][neighbour_j]

    def is_accessible(self, team: int) -> bool:
        """If this field is accessible for the given team."""

        if self.hill:
            return False
        if self.home_for_team is not None and self.home_for_team != team:
            return False
        return True


def main() -> None:
    global state
    state = State(json.loads(sys.stdin.read()))
    my_members = find_team_members(state.my_team_id)

    # TODO: set actions for all members
    # example: all members go up
    for member in my_members:
        member.action = Action.UP

    print(json.dumps(build_turn(my_members)))


def find_fields(predicate: Callable[[Field], bool]) -> List[Field]:
    """Find all fields that satisfy the given predicate."""

    result = []
    for row in state.world:
        for field in row:
            if predicate(field):
                result.append(field)
    return result


def find_team_members(team_id: int) -> List[Member]:
    """Find all members that belong to the given team."""

    result = []
    for row in state.world:
        for field in row:
            for member in field.members:
                if member.team == team_id:
                    result.append((member))
    return result


def pathfind(member: Member, goal: Field) -> Action:
    """Find the shortest path from member position to goal.

    Returns the first action to take.
    If there is no path, returns Action.STAY.
    """

    # BFS from goal to member
    # this direction makes it easier to realize that no path exists
    explored: Set[Field] = set()
    queue = collections.deque([goal])
    while queue:
        field = queue.popleft()
        for action in Action:
            neighbour = field.get_neighbour_field(action)
            if (neighbour in explored or
                not neighbour.is_accessible(member.team)):
                continue
            if neighbour == member.field:
                return action.invert()
            queue.append(neighbour)
            explored.add(neighbour)
    return Action.STAY


def build_turn(members: Iterable[Member]) -> dict:
    return {
        "members": [
            {"id": member.id, "action": member.action.to_json()}
            for member in members
        ]
    }


def parse_world(world: dict) -> List[List[Field]]:
    fields = []
    for i, row in enumerate(world):
        fields_row = []
        for j, field in enumerate(row):
            parsed_field = Field(
                i, j,
                field["hill"],
                field["home_for_team"],
                field["occupied_by_team"],
            )
            members = []
            for member in field["members"]:
                members.append(Member(
                    parsed_field,
                    member["team"],
                    member["id"],
                ))
            parsed_field.members = members
            fields_row.append(parsed_field)

        fields.append(fields_row)
    return fields


if __name__ == '__main__':
    main()