SKSP_2022_strategicka_hra/klient/play.py

#!/usr/bin/env python3
import json
import sys
from typing import Callable, Iterable, List, Optional, Tuple


Coords = Tuple[int, int]


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


class Member:
    def __init__(self, team: int, id: int, remaining_rounds: int):
        self.team = team
        self.id = id
        self.remaining_rounds = remaining_rounds
        self.action: Optional[str] = None


class Field:
    def __init__(self, home_for_team: Optional[int],
                 occupied_by_team: Optional[int],
                 hill: bool, members: List[Member]):
        self.home_for_team = home_for_team
        self.occupied_by_team = occupied_by_team
        self.hill = hill
        self.members = members


state: State = None


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 = "up"

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


def find_fields(
    predicate: Callable[[Field], bool]
) -> List[Tuple[Field, Coords]]:
    """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[Tuple[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 build_turn(members: Iterable[Member]) -> dict:
    return {
        "members": [
            {"id": member.id, "action": member.action}
            for member in members
        ]
    }


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


if __name__ == '__main__':
    main()