from __future__ import annotations

from typing import ClassVar, Optional

from pydantic import Field, validator

from pyatlan.model.fields.atlan_fields import (
    BooleanField,
    KeywordField,
    NumericField,
    RelationField,
    TextField,
)

from .asset40 import Sigma

class SigmaWorkbook(Sigma):
    """Description"""

    type_name: str = Field("SigmaWorkbook", allow_mutation=False)

    @validator("type_name")
    def validate_type_name(cls, v):
        if v != "SigmaWorkbook":
            raise ValueError("must be SigmaWorkbook")
        return v

    def __setattr__(self, name, value):
        if name in SigmaWorkbook._convenience_properties:
            return object.__setattr__(self, name, value)
        super().__setattr__(name, value)

    SIGMA_PAGE_COUNT: ClassVar[NumericField] = NumericField(
        "sigmaPageCount", "sigmaPageCount"
    )
    """
    TBC
    """

    SIGMA_PAGES: ClassVar[RelationField] = RelationField("sigmaPages")
    """
    TBC
    """

    _convenience_properties: ClassVar[list[str]] = [
        "sigma_page_count",
        "sigma_pages",
    ]

    @property
    def sigma_page_count(self) -> Optional[int]:
        return None if self.attributes is None else self.attributes.sigma_page_count

    @sigma_page_count.setter
    def sigma_page_count(self, sigma_page_count: Optional[int]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_page_count = sigma_page_count

    @property
    def sigma_pages(self) -> Optional[list[SigmaPage]]:
        return None if self.attributes is None else self.attributes.sigma_pages

    @sigma_pages.setter
    def sigma_pages(self, sigma_pages: Optional[list[SigmaPage]]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_pages = sigma_pages

    class Attributes(Sigma.Attributes):
        sigma_page_count: Optional[int] = Field(
            None, description="", alias="sigmaPageCount"
        )
        sigma_pages: Optional[list[SigmaPage]] = Field(
            None, description="", alias="sigmaPages"
        )  # relationship

    attributes: "SigmaWorkbook.Attributes" = Field(
        default_factory=lambda: SigmaWorkbook.Attributes(),
        description="Map of attributes in the instance and their values. The specific keys of this map will vary by "
        "type, so are described in the sub-types of this schema.\n",
    )

class SigmaDataElementField(Sigma):
    """Description"""

    type_name: str = Field("SigmaDataElementField", allow_mutation=False)

    @validator("type_name")
    def validate_type_name(cls, v):
        if v != "SigmaDataElementField":
            raise ValueError("must be SigmaDataElementField")
        return v

    def __setattr__(self, name, value):
        if name in SigmaDataElementField._convenience_properties:
            return object.__setattr__(self, name, value)
        super().__setattr__(name, value)

    SIGMA_DATA_ELEMENT_FIELD_IS_HIDDEN: ClassVar[BooleanField] = BooleanField(
        "sigmaDataElementFieldIsHidden", "sigmaDataElementFieldIsHidden"
    )
    """
    TBC
    """
    SIGMA_DATA_ELEMENT_FIELD_FORMULA: ClassVar[TextField] = TextField(
        "sigmaDataElementFieldFormula", "sigmaDataElementFieldFormula"
    )
    """
    TBC
    """

    SIGMA_DATA_ELEMENT: ClassVar[RelationField] = RelationField("sigmaDataElement")
    """
    TBC
    """

    _convenience_properties: ClassVar[list[str]] = [
        "sigma_data_element_field_is_hidden",
        "sigma_data_element_field_formula",
        "sigma_data_element",
    ]

    @property
    def sigma_data_element_field_is_hidden(self) -> Optional[bool]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_field_is_hidden
        )

    @sigma_data_element_field_is_hidden.setter
    def sigma_data_element_field_is_hidden(
        self, sigma_data_element_field_is_hidden: Optional[bool]
    ):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_field_is_hidden = (
            sigma_data_element_field_is_hidden
        )

    @property
    def sigma_data_element_field_formula(self) -> Optional[str]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_field_formula
        )

    @sigma_data_element_field_formula.setter
    def sigma_data_element_field_formula(
        self, sigma_data_element_field_formula: Optional[str]
    ):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_field_formula = (
            sigma_data_element_field_formula
        )

    @property
    def sigma_data_element(self) -> Optional[SigmaDataElement]:
        return None if self.attributes is None else self.attributes.sigma_data_element

    @sigma_data_element.setter
    def sigma_data_element(self, sigma_data_element: Optional[SigmaDataElement]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element = sigma_data_element

    class Attributes(Sigma.Attributes):
        sigma_data_element_field_is_hidden: Optional[bool] = Field(
            None, description="", alias="sigmaDataElementFieldIsHidden"
        )
        sigma_data_element_field_formula: Optional[str] = Field(
            None, description="", alias="sigmaDataElementFieldFormula"
        )
        sigma_data_element: Optional[SigmaDataElement] = Field(
            None, description="", alias="sigmaDataElement"
        )  # relationship

    attributes: "SigmaDataElementField.Attributes" = Field(
        default_factory=lambda: SigmaDataElementField.Attributes(),
        description="Map of attributes in the instance and their values. The specific keys of this map will vary by "
        "type, so are described in the sub-types of this schema.\n",
    )

class SigmaPage(Sigma):
    """Description"""

    type_name: str = Field("SigmaPage", allow_mutation=False)

    @validator("type_name")
    def validate_type_name(cls, v):
        if v != "SigmaPage":
            raise ValueError("must be SigmaPage")
        return v

    def __setattr__(self, name, value):
        if name in SigmaPage._convenience_properties:
            return object.__setattr__(self, name, value)
        super().__setattr__(name, value)

    SIGMA_DATA_ELEMENT_COUNT: ClassVar[NumericField] = NumericField(
        "sigmaDataElementCount", "sigmaDataElementCount"
    )
    """
    TBC
    """

    SIGMA_DATA_ELEMENTS: ClassVar[RelationField] = RelationField("sigmaDataElements")
    """
    TBC
    """
    SIGMA_WORKBOOK: ClassVar[RelationField] = RelationField("sigmaWorkbook")
    """
    TBC
    """

    _convenience_properties: ClassVar[list[str]] = [
        "sigma_data_element_count",
        "sigma_data_elements",
        "sigma_workbook",
    ]

    @property
    def sigma_data_element_count(self) -> Optional[int]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_count
        )

    @sigma_data_element_count.setter
    def sigma_data_element_count(self, sigma_data_element_count: Optional[int]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_count = sigma_data_element_count

    @property
    def sigma_data_elements(self) -> Optional[list[SigmaDataElement]]:
        return None if self.attributes is None else self.attributes.sigma_data_elements

    @sigma_data_elements.setter
    def sigma_data_elements(
        self, sigma_data_elements: Optional[list[SigmaDataElement]]
    ):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_elements = sigma_data_elements

    @property
    def sigma_workbook(self) -> Optional[SigmaWorkbook]:
        return None if self.attributes is None else self.attributes.sigma_workbook

    @sigma_workbook.setter
    def sigma_workbook(self, sigma_workbook: Optional[SigmaWorkbook]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_workbook = sigma_workbook

    class Attributes(Sigma.Attributes):
        sigma_data_element_count: Optional[int] = Field(
            None, description="", alias="sigmaDataElementCount"
        )
        sigma_data_elements: Optional[list[SigmaDataElement]] = Field(
            None, description="", alias="sigmaDataElements"
        )  # relationship
        sigma_workbook: Optional[SigmaWorkbook] = Field(
            None, description="", alias="sigmaWorkbook"
        )  # relationship

    attributes: "SigmaPage.Attributes" = Field(
        default_factory=lambda: SigmaPage.Attributes(),
        description="Map of attributes in the instance and their values. The specific keys of this map will vary by "
        "type, so are described in the sub-types of this schema.\n",
    )

class SigmaDataElement(Sigma):
    """Description"""

    type_name: str = Field("SigmaDataElement", allow_mutation=False)

    @validator("type_name")
    def validate_type_name(cls, v):
        if v != "SigmaDataElement":
            raise ValueError("must be SigmaDataElement")
        return v

    def __setattr__(self, name, value):
        if name in SigmaDataElement._convenience_properties:
            return object.__setattr__(self, name, value)
        super().__setattr__(name, value)

    SIGMA_DATA_ELEMENT_QUERY: ClassVar[KeywordField] = KeywordField(
        "sigmaDataElementQuery", "sigmaDataElementQuery"
    )
    """
    TBC
    """
    SIGMA_DATA_ELEMENT_TYPE: ClassVar[KeywordField] = KeywordField(
        "sigmaDataElementType", "sigmaDataElementType"
    )
    """
    TBC
    """
    SIGMA_DATA_ELEMENT_FIELD_COUNT: ClassVar[NumericField] = NumericField(
        "sigmaDataElementFieldCount", "sigmaDataElementFieldCount"
    )
    """
    TBC
    """

    SIGMA_PAGE: ClassVar[RelationField] = RelationField("sigmaPage")
    """
    TBC
    """
    SIGMA_DATA_ELEMENT_FIELDS: ClassVar[RelationField] = RelationField(
        "sigmaDataElementFields"
    )
    """
    TBC
    """

    _convenience_properties: ClassVar[list[str]] = [
        "sigma_data_element_query",
        "sigma_data_element_type",
        "sigma_data_element_field_count",
        "sigma_page",
        "sigma_data_element_fields",
    ]

    @property
    def sigma_data_element_query(self) -> Optional[str]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_query
        )

    @sigma_data_element_query.setter
    def sigma_data_element_query(self, sigma_data_element_query: Optional[str]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_query = sigma_data_element_query

    @property
    def sigma_data_element_type(self) -> Optional[str]:
        return (
            None if self.attributes is None else self.attributes.sigma_data_element_type
        )

    @sigma_data_element_type.setter
    def sigma_data_element_type(self, sigma_data_element_type: Optional[str]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_type = sigma_data_element_type

    @property
    def sigma_data_element_field_count(self) -> Optional[int]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_field_count
        )

    @sigma_data_element_field_count.setter
    def sigma_data_element_field_count(
        self, sigma_data_element_field_count: Optional[int]
    ):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_field_count = sigma_data_element_field_count

    @property
    def sigma_page(self) -> Optional[SigmaPage]:
        return None if self.attributes is None else self.attributes.sigma_page

    @sigma_page.setter
    def sigma_page(self, sigma_page: Optional[SigmaPage]):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_page = sigma_page

    @property
    def sigma_data_element_fields(self) -> Optional[list[SigmaDataElementField]]:
        return (
            None
            if self.attributes is None
            else self.attributes.sigma_data_element_fields
        )

    @sigma_data_element_fields.setter
    def sigma_data_element_fields(
        self, sigma_data_element_fields: Optional[list[SigmaDataElementField]]
    ):
        if self.attributes is None:
            self.attributes = self.Attributes()
        self.attributes.sigma_data_element_fields = sigma_data_element_fields

    class Attributes(Sigma.Attributes):
        sigma_data_element_query: Optional[str] = Field(
            None, description="", alias="sigmaDataElementQuery"
        )
        sigma_data_element_type: Optional[str] = Field(
            None, description="", alias="sigmaDataElementType"
        )
        sigma_data_element_field_count: Optional[int] = Field(
            None, description="", alias="sigmaDataElementFieldCount"
        )
        sigma_page: Optional[SigmaPage] = Field(
            None, description="", alias="sigmaPage"
        )  # relationship
        sigma_data_element_fields: Optional[list[SigmaDataElementField]] = Field(
            None, description="", alias="sigmaDataElementFields"
        )  # relationship

    attributes: "SigmaDataElement.Attributes" = Field(
        default_factory=lambda: SigmaDataElement.Attributes(),
        description="Map of attributes in the instance and their values. The specific keys of this map will vary by "
        "type, so are described in the sub-types of this schema.\n",
    )

SigmaWorkbook.Attributes.update_forward_refs()

SigmaDataElementField.Attributes.update_forward_refs()

SigmaPage.Attributes.update_forward_refs()

SigmaDataElement.Attributes.update_forward_refs()


from copy import deepcopy
from enum import Enum, auto
from functools import wraps
from kaggle_environments.helpers import Point, group_by, Direction
from typing import *
import sys
import math
import random
import kaggle_environments.helpers

# region Data Model Classes
class Observation(kaggle_environments.helpers.Observation):
    """
    Observation primarily used as a helper to construct the Board from the raw observation.
    This provides bindings for the observation type described at https://github.com/Kaggle/kaggle-environments/blob/master/kaggle_environments/envs/kore/kore.json
    """
    @property
    def kore(self) -> List[float]:
        """Serialized list of available kore per cell on the board."""
        return self["kore"]

    @property
    def players(self) -> List[List[int]]:
        """List of players and their assets."""
        return self["players"]

    @property
    def player(self) -> int:
        """The current agent's player index."""
        return self["player"]

class Configuration(kaggle_environments.helpers.Configuration):
    """
    Configuration provides access to tunable parameters in the environment.
    This provides bindings for the configuration type described at https://github.com/Kaggle/kaggle-environments/blob/master/kaggle_environments/envs/kore/kore.json
    """
    @property
    def agent_timeout(self) -> float:
        """Maximum runtime (seconds) to initialize an agent."""
        return self["agentTimeout"]

    @property
    def starting_kore(self) -> int:
        """The starting amount of kore available on the board."""
        return self["startingKore"]

    @property
    def size(self) -> int:
        """The number of cells vertically and horizontally on the board."""
        return self["size"]

    @property
    def spawn_cost(self) -> int:
        """The amount of kore to spawn a new ship."""
        return self["spawnCost"]

    @property
    def convert_cost(self) -> int:
        """The amount of ships needed from a fleet to create a shipyard."""
        return self["convertCost"]

    @property
    def regen_rate(self) -> float:
        """The rate kore regenerates on the board."""
        return self["regenRate"]

    @property
    def max_cell_kore(self) -> int:
        """The maximum kore that can be in any cell."""
        return self["maxRegenCellKore"]

    @property
    def random_seed(self) -> int:
        """The seed to the random number generator (0 means no seed)."""
        return self["randomSeed"]

class ShipyardActionType(Enum):
    SPAWN = auto()
    LAUNCH = auto()

    def __str__(self) -> str:
        return self.name

class ShipyardAction:

    def __init__(self, type: ShipyardActionType, num_ships: Optional[int], flight_plan: Optional[str]) -> None:
        self._type = type
        assert num_ships >= 0, "must be a non-negative number"
        assert num_ships == int(num_ships), "must be an integer"
        self._num_ships = num_ships
        self._flight_plan = flight_plan

    def __str__(self) -> str:
        if self._type == ShipyardActionType.SPAWN:
            return f'{self._type.name}_{self._num_ships}'
        if self._type == ShipyardActionType.LAUNCH:
            return f'{self._type.name}_{self._num_ships}_{self._flight_plan}'

    @property
    def name(self):
        return str(self)

    @staticmethod
    def from_str(raw: str):
        if not raw:
            return None
        if raw.startswith(ShipyardActionType.SPAWN.name):
            _, ship_str = raw.split("_")
            try:
                num_ships = int(ship_str.split(".")[0])
            except _:
                num_ships = 0
            return ShipyardAction.spawn_ships(num_ships)
        if raw.startswith(ShipyardActionType.LAUNCH.name):
            _, ship_str, plan_str = raw.split("_")
            try:
                num_ships = int(ship_str.split(".")[0])
            except _:
                num_ships = 0
            return ShipyardAction.launch_fleet_with_flight_plan(num_ships, plan_str)

    @staticmethod
    def launch_fleet_in_direction(number_ships: int, direction: Direction):
        flight_plan = None
        if isinstance(direction, Direction):
            flight_plan = direction.to_char()
        else:
            flight_plan = flight_plan.upper()
        return ShipyardAction.launch_fleet_with_flight_plan(number_ships, flight_plan)

    @staticmethod
    def launch_fleet_with_flight_plan(number_ships: int, flight_plan: str):
        flight_plan = flight_plan.upper()
        assert number_ships > 0, "must be a positive number_ships"
        assert number_ships == int(number_ships), "must be an integer number_ships"
        assert flight_plan is not None and len(flight_plan) > 0, "flight_plan must be a str of len > 0"
        assert flight_plan[0].isalpha() and flight_plan[0] in "NESW", "flight_plan must start with a valid direciton NESW"
        assert all([c in "NESWC0123456789" for c in flight_plan]), "flight_plan (" + flight_plan + ")can only contain NESWC0-9"
        if len(flight_plan) > Fleet.max_flight_plan_len_for_ship_count(number_ships):
            print("flight plan will be truncated: flight plan for " + str(number_ships) + " must be at most " + str(Fleet.max_flight_plan_len_for_ship_count(number_ships)))
        return ShipyardAction(ShipyardActionType.LAUNCH, number_ships, flight_plan)

    @staticmethod
    def spawn_ships(number_ships: int):
        assert number_ships == int(number_ships), "must be an integer number_ships"
        return ShipyardAction(ShipyardActionType.SPAWN, number_ships, None)

    @property
    def action_type(self) -> ShipyardActionType:
        return self._type

    @property
    def num_ships(self) -> Optional[int]:
        return self._num_ships

    @property
    def flight_plan(self) -> Optional[str]:
        return self._flight_plan

FleetId = NewType('FleetId', str)
ShipyardId = NewType('ShipyardId', str)
PlayerId = NewType('PlayerId', int)

class Cell:
    def __init__(self, position: Point, kore: float, shipyard_id: Optional[ShipyardId], fleet_id: Optional[FleetId], board: 'Board') -> None:
        self._position = position
        self._kore = kore
        self._shipyard_id = shipyard_id
        self._fleet_id = fleet_id
        self._board = board

    @property
    def position(self) -> Point:
        return self._position

    @property
    def kore(self) -> float:
        return self._kore

    @property
    def shipyard_id(self) -> Optional[ShipyardId]:
        return self._shipyard_id

    @property
    def fleet_id(self) -> Optional[FleetId]:
        return self._fleet_id

    @property
    def fleet(self) -> Optional['Fleet']:
        """Returns the fleet on this cell if it exists and None otherwise."""
        return self._board.fleets.get(self.fleet_id)

    @property
    def shipyard(self) -> Optional['Shipyard']:
        """Returns the shipyard on this cell if it exists and None otherwise."""
        return self._board.shipyards.get(self.shipyard_id)

    def neighbor(self, offset: Point) -> 'Cell':
        """Returns the cell at self.position + offset."""
        (x, y) = self.position + offset
        return self._board[x, y]

    @property
    def north(self) -> 'Cell':
        """Returns the cell north of this cell."""
        return self.neighbor(Direction.NORTH.to_point())

    @property
    def south(self) -> 'Cell':
        """Returns the cell south of this cell."""
        return self.neighbor(Direction.SOUTH.to_point())

    @property
    def east(self) -> 'Cell':
        """Returns the cell east of this cell."""
        return self.neighbor(Direction.EAST.to_point())

    @property
    def west(self) -> 'Cell':
        """Returns the cell west of this cell."""
        return self.neighbor(Direction.WEST.to_point())

class Fleet:
    def __init__(self, fleet_id: FleetId, ship_count: int, direction: Direction, position: Point, kore: int, flight_plan: str, player_id: PlayerId, board: 'Board') -> None:
        self._id = fleet_id
        self._ship_count = ship_count
        self._direction = direction
        self._position = position
        self._flight_plan = flight_plan
        self._kore = kore
        self._player_id = player_id
        self._board = board

    @property
    def id(self) -> FleetId:
        return self._id

    @property
    def ship_count(self) -> int:
        return self._ship_count

    @property
    def direction(self) -> Direction:
        return self._direction

    @property
    def position(self) -> Point:
        return self._position

    @property
    def kore(self) -> int:
        return self._kore

    @property
    def player_id(self) -> PlayerId:
        return self._player_id

    @property
    def cell(self) -> Cell:
        """Returns the cell this fleet is on."""
        return self._board[self.position]

    @property
    def player(self) -> 'Player':
        """Returns the player that owns this ship."""
        return self._board.players[self.player_id]

    @property
    def flight_plan(self) -> str:
        """Returns the current flight plan of the fleet"""
        return self._flight_plan

    @property
    def collection_rate(self) -> float:
        """ln(ship_count) / 10"""
        return min(math.log(self.ship_count) / 20, .99)

    @staticmethod
    def max_flight_plan_len_for_ship_count(ship_count) -> int:
        """Returns the length of the longest possible flight plan this fleet can be assigned"""
        return math.floor(2 * math.log(ship_count)) + 1

    @property
    def _observation(self) -> List[int]:
        """Converts a fleet back to the normalized observation subset that constructed it."""
        return [self.position.to_index(self._board.configuration.size), self.kore, self.ship_count, self.direction.to_index(), self.flight_plan]

    def less_than_other_allied_fleet(self, other):
        if not self.ship_count == other.ship_count:
            return self.ship_count < other.ship_count
        if not self.kore == other.kore:
            return self.kore < other.kore
        return self.direction.to_index() > other.direction.to_index()

upgrade_times = [pow(i,2) + 1 for i in range(1, 10)]
SPAWN_VALUES = []
current = 0
for t in upgrade_times:
    current += t
    SPAWN_VALUES.append(current)

class Shipyard:
    def __init__(self, shipyard_id: ShipyardId, ship_count: int, position: Point, player_id: PlayerId, turns_controlled: int, board: 'Board', next_action: Optional[ShipyardAction] = None) -> None:
        self._id = shipyard_id
        self._ship_count = ship_count
        self._position = position
        self._player_id = player_id
        self._turns_controlled = turns_controlled
        self._board = board
        self._next_action = next_action

    @property
    def id(self) -> ShipyardId:
        return self._id

    @property
    def ship_count(self):
        return self._ship_count

    @property
    def position(self) -> Point:
        return self._position

    @property
    def player_id(self) -> PlayerId:
        return self._player_id

    @property
    def max_spawn(self) -> int:
        for idx, target in enumerate(SPAWN_VALUES):
            if self._turns_controlled < target:
                return idx + 1
        return len(SPAWN_VALUES) + 1

    @property
    def cell(self) -> Cell:
        """Returns the cell this shipyard is on."""
        return self._board[self.position]

    @property
    def player(self) -> 'Player':
        return self._board.players[self.player_id]

    @property
    def next_action(self) -> ShipyardAction:
        """Returns the action that will be executed by this shipyard when Board.next() is called (when the current turn ends)."""
        return self._next_action

    @next_action.setter
    def next_action(self, value: Optional[ShipyardAction]) -> None:
        """Sets the action that will be executed by this shipyard when Board.next() is called (when the current turn ends)."""
        self._next_action = value

    @property
    def _observation(self) -> List[int]:
        """Converts a shipyard back to the normalized observation subset that constructed it."""
        return [self.position.to_index(self._board.configuration.size), self.ship_count, self._turns_controlled]

class Player:
    def __init__(self, player_id: PlayerId, kore: int, shipyard_ids: List[ShipyardId], fleet_ids: List[FleetId], board: 'Board') -> None:
        self._id = player_id
        self._kore = kore
        self._shipyard_ids = shipyard_ids
        self._fleet_ids = fleet_ids
        self._board = board

    @property
    def id(self) -> PlayerId:
        return self._id

    @property
    def kore(self) -> int:
        return self._kore

    @property
    def shipyard_ids(self) -> List[ShipyardId]:
        return self._shipyard_ids

    @property
    def fleet_ids(self) -> List[FleetId]:
        return self._fleet_ids

    @property
    def shipyards(self) -> List[Shipyard]:
        """Returns all shipyards owned by this player."""
        return [
            self._board.shipyards[shipyard_id]
            for shipyard_id in self.shipyard_ids
        ]

    @property
    def fleets(self) -> List[Fleet]:
        """Returns all fleets owned by this player."""
        return [
            self._board.fleets[fleet_id]
            for fleet_id in self.fleet_ids
        ]

    @property
    def is_current_player(self) -> bool:
        """Returns whether this player is the current player (generally if this returns True, this player is you)."""
        return self.id == self._board.current_player_id

    @property
    def next_actions(self) -> Dict[str, str]:
        """Returns all queued fleet and shipyard actions for this player formatted for the kore interpreter to receive as an agent response."""
        shipyard_actions = {
            shipyard.id: shipyard.next_action.name
            for shipyard in self.shipyards
            if shipyard.next_action is not None
        }
        return {**shipyard_actions}

    @property
    def _observation(self):
        """Converts a player back to the normalized observation subset that constructed it."""
        shipyards = {shipyard.id: shipyard._observation for shipyard in self.shipyards}
        fleets = {fleet.id: fleet._observation for fleet in self.fleets}
        return [self.kore, shipyards, fleets]
# endregion

class Board:
    def __init__(
        self,
        raw_observation: Dict[str, Any],
        raw_configuration: Union[Configuration, Dict[str, Any]],
        next_actions: Optional[List[Dict[str, str]]] = None
    ) -> None:
        """
        Creates a board from the provided observation, configuration, and next_actions as specified by
        https://github.com/Kaggle/kaggle-environments/blob/master/kaggle_environments/envs/kore/kore.json
        Board tracks players (by id), fleets (by id), shipyards (by id), and cells (by position).
        Each entity contains both key values (e.g. fleet.player_id) as well as entity references (e.g. fleet.player).
        References are deep and chainable e.g.
            [fleet.kore for player in board.players for fleet in player.fleets]
            fleet.player.shipyards[0].cell.north.east.fleet
        Consumers should not set or modify any attributes except and Shipyard.next_action
        """
        observation = Observation(raw_observation)
        # next_actions is effectively a Dict[Union[[FleetId, FleetAction], [ShipyardId, ShipyardAction]]]
        # but that type's not very expressible so we simplify it to Dict[str, str]
        # Later we'll iterate through it once for each fleet and shipyard to pull all the actions out
        next_actions = next_actions or ([{}] * len(observation.players))

        self._step = observation.step
        self._remaining_overage_time = observation.remaining_overage_time
        self._configuration = Configuration(raw_configuration)
        self._current_player_id = observation.player
        self._players: Dict[PlayerId, Player] = {}
        self._fleets: Dict[FleetId, Fleet] = {}
        self._shipyards: Dict[ShipyardId, Shipyard] = {}
        self._cells: Dict[Point, Cell] = {}

        size = self.configuration.size
        # Create a cell for every point in a size x size grid
        for x in range(size):
            for y in range(size):
                position = Point(x, y)
                kore = observation.kore[position.to_index(size)]
                # We'll populate the cell's fleets and shipyards in _add_fleet and _add_shipyard
                self.cells[position] = Cell(position, kore, None, None, self)

        for (player_id, player_observation) in enumerate(observation.players):
            # We know the len(player_observation) == 3 based on the schema -- this is a hack to have a tuple in json
            [player_kore, player_shipyards, player_fleets] = player_observation
            # We'll populate the player's fleets and shipyards in _add_fleet and _add_shipyard
            self.players[player_id] = Player(player_id, player_kore, [], [], self)
            player_actions = next_actions[player_id] or {}

            for (fleet_id, [fleet_index, fleet_kore, ship_count, direction, flight_plan]) in player_fleets.items():
                fleet_position = Point.from_index(fleet_index, size)
                fleet_direction = Direction.from_index(direction)
                self._add_fleet(Fleet(fleet_id, ship_count, fleet_direction, fleet_position, fleet_kore, flight_plan, player_id, self))

            for (shipyard_id, [shipyard_index, ship_count, turns_controlled]) in player_shipyards.items():
                shipyard_position = Point.from_index(shipyard_index, size)
                raw_action = player_actions.get(shipyard_id)
                action = ShipyardAction.from_str(raw_action)
                self._add_shipyard(Shipyard(shipyard_id, ship_count, shipyard_position, player_id, turns_controlled, self, action))

    @property
    def configuration(self) -> Configuration:
        return self._configuration

    @property
    def players(self) -> Dict[PlayerId, Player]:
        return self._players

    @property
    def fleets(self) -> Dict[FleetId, Fleet]:
        """Returns all fleets on the current board."""
        return self._fleets

    @property
    def shipyards(self) -> Dict[ShipyardId, Shipyard]:
        """Returns all shipyards on the current board."""
        return self._shipyards

    @property
    def cells(self) -> Dict[Point, Cell]:
        """Returns all cells on the current board."""
        return self._cells

    @property
    def step(self) -> int:
        return self._step

    @property
    def current_player_id(self) -> PlayerId:
        return self._current_player_id

    @property
    def current_player(self) -> Player:
        """Returns the current player (generally this is you)."""
        return self._players[self.current_player_id]

    @property
    def opponents(self) -> List[Player]:
        """
        Returns all players that aren't the current player.
        You can get all opponent fleets with [fleet for fleet in player.fleets for player in board.opponents]
        """
        return [player for player in self.players.values() if not player.is_current_player]

    @property
    def observation(self) -> Dict[str, Any]:
        """Converts a Board back to the normalized observation that constructed it."""
        size = self.configuration.size
        kore = [self[Point.from_index(index, size)].kore for index in range(size * size)]
        players = [player._observation for player in self.players.values()]

        return {
            "kore": kore,
            "players": players,
            "player": self.current_player_id,
            "step": self.step,
            "remainingOverageTime": self._remaining_overage_time,
        }

    def __deepcopy__(self, _) -> 'Board':
        actions = [player.next_actions for player in self.players.values()]
        return Board(self.observation, self.configuration, actions)

    def __getitem__(self, point: Union[Tuple[int, int], Point]) -> Cell:
        """
        This method will wrap the supplied position to fit within the board size and return the cell at that location.
        e.g. on a 3x3 board, board[2, 1] is the same as board[5, 1]
        """
        if not isinstance(point, Point):
            (x, y) = point
            point = Point(x, y)
        return self._cells[point % self.configuration.size]

    def __str__(self) -> str:
        """
        The board is printed in a grid with the following rules:
        Capital letters are shipyards
        Lower case letters are fleets
        Digits are cell kore and scale from 0-9 directly proportional to a value between 0 and self.configuration.max_cell_kore
        Player 1 is letter a/A
        Player 2 is letter b/B
        etc.
        """
        size = self.configuration.size
        result = ''
        for y in range(size):
            for x in range(size):
                cell = self[(x, size - y - 1)]
                result += '|'
                result += (
                    chr(ord('a') + cell.ship.player_id)
                    if cell.fleet is not None
                    else ' '
                )
                # This normalizes a value from 0 to max_cell kore to a value from 0 to 9
                normalized_kore = int(9.0 * cell.kore / float(self.configuration.max_cell_kore))
                result += str(normalized_kore)
                result += (
                    chr(ord('A') + cell.shipyard.player_id)
                    if cell.shipyard is not None
                    else ' '
                )
            result += '|\n'
        return result

    def _add_fleet(self: 'Board', fleet: Fleet) -> None:
        fleet.player.fleet_ids.append(fleet.id)
        fleet.cell._fleet_id = fleet.id
        self._fleets[fleet.id] = fleet

    def _add_shipyard(self: 'Board', shipyard: Shipyard) -> None:
        shipyard.player.shipyard_ids.append(shipyard.id)
        shipyard.cell._shipyard_id = shipyard.id
        shipyard.cell._kore = 0
        self._shipyards[shipyard.id] = shipyard

    def _delete_fleet(self: 'Board', fleet: Fleet) -> None:
        fleet.player.fleet_ids.remove(fleet.id)
        if fleet.cell.fleet_id == fleet.id:
            fleet.cell._fleet_id = None
        del self._fleets[fleet.id]

    def _delete_shipyard(self: 'Board', shipyard: Shipyard) -> None:
        shipyard.player.shipyard_ids.remove(shipyard.id)
        if shipyard.cell.shipyard_id == shipyard.id:
            shipyard.cell._shipyard_id = None
        del self._shipyards[shipyard.id]

    def get_fleet_at_point(self: 'Board', position: Point) -> Optional[Fleet]:
        matches = [fleet for fleet in self.fleets.values() if fleet.position == position]
        if matches:
            assert len(matches) == 1
            return matches[0]
        return None

    def get_shipyard_at_point(self: 'Board', position: Point) -> Optional[Shipyard]:
        matches = [shipyard for shipyard in self.shipyards.values() if shipyard.position == position]
        if matches:
            assert len(matches) == 1
            return matches[0]
        return None

    def get_cell_at_point(self: 'Board', position: Point):
        return self.cells.get(position)

    def print(self: 'Board') -> None:
        size = self.configuration.size
        player_chars = {
            pid: alpha
            for pid, alpha in  zip(self.players, "abcdef"[:len(self.players)])
        }
        print(self.configuration.size * "=")
        for i in range(size):
            row = ""
            for j in range(size):
                pos = Point(j, size - 1 - i)
                curr_cell = self.cells[pos]
                if curr_cell.shipyard is not None:
                    row += player_chars[curr_cell.shipyard.player_id].upper()
                elif curr_cell.fleet is not None:
                    row += player_chars[curr_cell.fleet.player_id]
                elif curr_cell.kore <= 50:
                    row += " "
                elif curr_cell.kore <= 250:
                    row += "."
                elif curr_cell.kore <= 400:
                    row += "*"
                elif curr_cell.kore > 400:
                    row += "o"
            print(row)
        print(self.configuration.size * "=")

    def print_kore(self: 'Board') -> None:
        size = self.configuration.size
        print(self.configuration.size * "=")
        for i in range(size):
            row = ""
            for j in range(size):
                pos = Point(j, size - 1 - i)
                curr_cell = self.cells[pos]
                row += str(int(curr_cell.kore)) + ","
            print(row)
        print(self.configuration.size * "=")

    def next(self) -> 'Board':
        """
        Returns a new board with the current board's next actions applied.
        The current board is unmodified.
        This can form a kore interpreter, e.g.
            next_observation = Board(current_observation, configuration, actions).next().observation
        """
        # Create a copy of the board to modify so we don't affect the current board
        board = deepcopy(self)
        configuration = board.configuration
        convert_cost = configuration.convert_cost
        spawn_cost = configuration.spawn_cost
        uid_counter = 0

        # This is a consistent way to generate unique strings to form fleet and shipyard ids
        def create_uid():
            nonlocal uid_counter
            uid_counter += 1
            return f"{self.step + 1}-{uid_counter}"

        # this checks the validity of a flight plan
        def is_valid_flight_plan(flight_plan):
            return len([c for c in flight_plan if c not in "NESWC0123456789"]) == 0

        # Process actions and store the results in the fleets and shipyards lists for collision checking
        for player in board.players.values():
            for shipyard in player.shipyards:
                if shipyard.next_action == None:
                    pass
                elif shipyard.next_action.num_ships == 0:
                    pass
                elif (shipyard.next_action.action_type == ShipyardActionType.SPAWN
                        and player.kore >= spawn_cost * shipyard.next_action.num_ships
                        and shipyard.next_action.num_ships <= shipyard.max_spawn):
                    # Handle SPAWN actions
                    player._kore -= spawn_cost * shipyard.next_action.num_ships
                    shipyard._ship_count += shipyard.next_action.num_ships
                elif shipyard.next_action.action_type == ShipyardActionType.LAUNCH and shipyard.ship_count >= shipyard.next_action.num_ships:
                    flight_plan = shipyard.next_action.flight_plan
                    if not flight_plan or not is_valid_flight_plan(flight_plan):
                        continue
                    shipyard._ship_count -= shipyard.next_action.num_ships
                    direction = Direction.from_char(flight_plan[0])
                    max_flight_plan_len = Fleet.max_flight_plan_len_for_ship_count(shipyard.next_action.num_ships)
                    if len(flight_plan) > max_flight_plan_len:
                        flight_plan = flight_plan[:max_flight_plan_len]
                    board._add_fleet(Fleet(FleetId(create_uid()), shipyard.next_action.num_ships, direction, shipyard.position, 0, flight_plan, player.id, board))

            # Clear the shipyard's action so it doesn't repeat the same action automatically
            for shipyard in player.shipyards:
                shipyard.next_action = None
                shipyard._turns_controlled += 1

            def find_first_non_digit(candidate_str):
                for i in range(len(candidate_str)):
                    if not candidate_str[i].isdigit():
                        return i
                else:
                    return len(candidate_str) + 1
                return 0

            for fleet in player.fleets:
                # remove any errant 0s
                while fleet.flight_plan and fleet.flight_plan.startswith("0"):
                    fleet._flight_plan = fleet.flight_plan[1:]
                if fleet.flight_plan and fleet.flight_plan[0] == "C" and fleet.ship_count >= convert_cost and fleet.cell.shipyard_id is None:
                    player._kore += fleet.kore
                    fleet.cell._kore = 0
                    board._add_shipyard(Shipyard(ShipyardId(create_uid()), fleet.ship_count - convert_cost, fleet.position, player.id, 0, board))
                    board._delete_fleet(fleet)
                    continue

                while fleet.flight_plan and fleet.flight_plan[0] == "C":
                    # couldn't build, remove the Convert and continue with flight plan
                    fleet._flight_plan = fleet.flight_plan[1:]

                if fleet.flight_plan and fleet.flight_plan[0].isalpha():
                    fleet._direction = Direction.from_char(fleet.flight_plan[0])
                    fleet._flight_plan = fleet.flight_plan[1:]
                elif fleet.flight_plan:
                    idx = find_first_non_digit(fleet.flight_plan)
                    digits = int(fleet.flight_plan[:idx])
                    rest = fleet.flight_plan[idx:]
                    digits -= 1
                    if digits > 0:
                        fleet._flight_plan = str(digits) + rest
                    else:
                        fleet._flight_plan = rest

                # continue moving in the fleet's direction
                fleet.cell._fleet_id = None
                fleet._position = fleet.position.translate(fleet.direction.to_point(), configuration.size)
                # We don't set the new cell's fleet_id here as it would be overwritten by another fleet in the case of collision.

            def combine_fleets(fid1: FleetId, fid2: FleetId) -> FleetId:
                f1 = board.fleets[fid1]
                f2 = board.fleets[fid2]
                if f1.less_than_other_allied_fleet(f2):
                    f1, f2 = f2, f1
                    fid1, fid2 = fid2, fid1
                f1._kore += f2.kore
                f1._ship_count += f2._ship_count
                board._delete_fleet(f2)
                return fid1

            # resolve any allied fleets that ended up in the same square
            fleets_by_loc = group_by(player.fleets, lambda fleet: fleet.position.to_index(configuration.size))
            for value in fleets_by_loc.values():
                value.sort(key=lambda fleet: (fleet.ship_count, fleet.kore, -fleet.direction.to_index()), reverse=True)
                fid = value[0].id
                for i in range (1, len(value)):
                    fid = combine_fleets(fid, value[i].id)

            # Lets just check and make sure.
            assert player.kore >= 0

        def resolve_collision(fleets: List[Fleet]) -> Tuple[Optional[Fleet], List[Fleet]]:
            """
            Accepts the list of fleets at a particular position (must not be empty).
            Returns the fleet with the most ships or None in the case of a tie along with all other fleets.
            """
            if len(fleets) == 1:
                return fleets[0], []
            fleets_by_ships = group_by(fleets, lambda fleet: fleet.ship_count)
            most_ships = max(fleets_by_ships.keys())
            largest_fleets = fleets_by_ships[most_ships]
            if len(largest_fleets) == 1:
                # There was a winner, return it
                winner = largest_fleets[0]
                return winner, [fleet for fleet in fleets if fleet != winner]
            # There was a tie for most ships, all are deleted
            return None, fleets

        # Check for fleet to fleet collisions
        fleet_collision_groups = group_by(board.fleets.values(), lambda fleet: fleet.position)
        for position, collided_fleets in fleet_collision_groups.items():
            winner, deleted = resolve_collision(collided_fleets)
            shipyard = group_by(board.shipyards.values(), lambda shipyard: shipyard.position).get(position)
            if winner is not None:
                winner.cell._fleet_id = winner.id
                max_enemy_size = max([fleet.ship_count for fleet in deleted]) if deleted else 0
                winner._ship_count -= max_enemy_size
            for fleet in deleted:
                board._delete_fleet(fleet)
                if winner is not None:
                    # Winner takes deleted fleets' kore
                    winner._kore += fleet.kore
                elif winner is None and shipyard and shipyard[0].player:
                    # Desposit the kore into the shipyard
                    shipyard[0].player._kore += fleet.kore
                elif winner is None:
                    # Desposit the kore on the square
                    board.cells[position]._kore += fleet.kore

        # Check for fleet to shipyard collisions
        for shipyard in list(board.shipyards.values()):
            fleet = shipyard.cell.fleet
            if fleet is not None and fleet.player_id != shipyard.player_id:
                if fleet.ship_count > shipyard.ship_count:
                    count = fleet.ship_count - shipyard.ship_count
                    board._delete_shipyard(shipyard)
                    board._add_shipyard(Shipyard(ShipyardId(create_uid()), count, shipyard.position, fleet.player.id, 1, board))
                    fleet.player._kore += fleet.kore
                    board._delete_fleet(fleet)
                else:
                    shipyard._ship_count -= fleet.ship_count
                    shipyard.player._kore += fleet.kore
                    board._delete_fleet(fleet)

        # Deposit kore from fleets into shipyards
        for shipyard in list(board.shipyards.values()):
            fleet = shipyard.cell.fleet
            if fleet is not None and fleet.player_id == shipyard.player_id:
                shipyard.player._kore += fleet.kore
                shipyard._ship_count += fleet.ship_count
                board._delete_fleet(fleet)

        # apply fleet to fleet damage on all orthagonally adjacent cells
        incoming_fleet_dmg = DefaultDict(lambda: DefaultDict(int))
        for fleet in board.fleets.values():
            for direction in Direction.list_directions():
                curr_pos = fleet.position.translate(direction.to_point(), board.configuration.size)
                fleet_at_pos = board.get_fleet_at_point(curr_pos)
                if fleet_at_pos and not fleet_at_pos.player_id == fleet.player_id:
                    incoming_fleet_dmg[fleet_at_pos.id][fleet.id] = fleet.ship_count

        # dump 1/2 kore to the cell of killed fleets
        # mark the other 1/2 kore to go to surrounding fleets proportionally
        to_distribute = DefaultDict(lambda: DefaultDict(int))
        for fleet_id, fleet_dmg_dict in incoming_fleet_dmg.items():
            fleet = board.fleets[fleet_id]
            damage = sum(fleet_dmg_dict.values())
            if damage >= fleet.ship_count:
                fleet.cell._kore += fleet.kore / 2
                to_split = fleet.kore / 2
                for f_id, dmg in fleet_dmg_dict.items():
                    to_distribute[f_id][fleet.position.to_index(board.configuration.size)] = to_split * dmg/damage
                board._delete_fleet(fleet)
            else:
                fleet._ship_count -= damage

        # give kore claimed above to surviving fleets, otherwise add it to the kore of the tile where the fleet died
        for fleet_id, loc_kore_dict in to_distribute.items():
            fleet = board.fleets.get(fleet_id)
            if fleet:
                fleet._kore += sum(loc_kore_dict.values())
            else:
                for loc_idx, kore in loc_kore_dict.items():
                    board.cells.get(Point.from_index(loc_idx, board.configuration.size))._kore += kore

        # Collect kore from cells into fleets
        for fleet in board.fleets.values():
            cell = fleet.cell
            delta_kore = round(cell.kore * min(fleet.collection_rate, .99), 3)
            if delta_kore > 0:
                fleet._kore += delta_kore
                cell._kore -= delta_kore

        # Regenerate kore in cells
        for cell in board.cells.values():
            if cell.fleet_id is None and cell.shipyard_id is None:
                if cell.kore < configuration.max_cell_kore:
                    next_kore = round(cell.kore * (1 + configuration.regen_rate), 3)
                    cell._kore = next_kore

        board._step += 1

        # self.print()

        return board

def board_agent(agent: Callable[[Board], None]):
    """
    Decorator used to create an agent that modifies a board rather than an observation and a configuration
    Automatically returns the modified board's next actions

    @board_agent
    def my_agent(board: Board) -> None:
        ...
    """
    @wraps(agent)
    def agent_wrapper(obs, config) -> Dict[str, str]:
        board = Board(obs, config)
        agent(board)
        return board.current_player.next_actions

    if agent.__module__ is not None and agent.__module__ in sys.modules:
        setattr(sys.modules[agent.__module__], agent.__name__, agent_wrapper)
    return agent_wrapper


import os
import sys

sys.path.insert(0, os.path.abspath("../"))

from tokenization_kbalbert import KbAlbertCharTokenizer
from transformers import AutoModelForSequenceClassification, pipeline

kb_albert_model_path = "../kb-albert-char-base-v1"
model_output_path = "nsmc_outputs"
model = AutoModelForSequenceClassification.from_pretrained(model_output_path)
tokenizer = KbAlbertCharTokenizer.from_pretrained(kb_albert_model_path)

nsmc_classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer, framework="pt")

reviews = ["이 영화 최악이었어!", "볼거리가 많은 내 인생 영화 ㅎㅎ"]
results = nsmc_classifier(reviews)
for result in results:
    print(f"label: {result['label']}, with score: {round(result['score'], 4)}")


from copy import deepcopy

import torch
from torch import nn
from torch.nn import functional as F

class MLP(nn.Module):
    """
    Multi-layer perceptron (i.e. Multi-layer full-connection neural network)
    使用 nn.Conv1d 代替 nn.Linear
    (linear + BN + ReLU) * n + linear
    forward()输入 (N,C,L)
    :param in_channel:
    :param out_channels:
    :param do_in:
    :return:
        nn.Sequential
    """

    def __init__(self, in_channel: int, out_channels: list, do_in=True) -> None:
        super().__init__()
        self.layers = nn.ModuleList()
        for idx, out_channel in enumerate(out_channels):
            self.layers.append(nn.Conv1d(in_channel, out_channel, kernel_size=(1,), bias=True))
            if idx < len(out_channels) - 1:
                if do_in:
                    self.layers.append(nn.InstanceNorm1d(out_channel))
                self.layers.append(nn.ReLU())  # TODO: 默认 mlp 中使用 ReLU
            in_channel = out_channel
        self.out_channel = out_channels[-1]

    def forward(self, x):
        if x.shape[2] > 0:
            for layer in self.layers:
                x = layer(x)
            return x
        else:
            return torch.empty((1, self.out_channel, 0), device=x.device)

class ConvBlock(nn.Module):

    def __init__(self, in_channels, out_channels, conv_num_in_block, conv_norm='none'):
        super().__init__()
        self.layers = nn.ModuleList()
        for idx in range(conv_num_in_block):
            self.layers.append(nn.Conv2d(in_channels, out_channels, (3, 3), padding=(1, 1)))
            if conv_norm == 'batch':
                self.layers.append(nn.BatchNorm2d(out_channels))
            elif conv_norm == 'instance':
                self.layers.append(nn.InstanceNorm2d(out_channels))
            elif conv_norm == 'none':
                pass
            else:
                raise ValueError(f'conv_norm={conv_norm} in ConvBlock is invalid! use "none", "batch", "instance"')
            self.layers.append(nn.ReLU())  # TODO: in old using LReLU(0.2)
            in_channels = out_channels

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

class SENetBlock(nn.Module):

    def __init__(self, in_channels, dense_div):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(in_channels, max(in_channels // dense_div, 1)),
            nn.ReLU(),  # TODO: in old using LReLU(0.2)
            nn.Linear(max(in_channels // dense_div, 1), in_channels),
            nn.Sigmoid()
        ])

    def forward(self, x):
        x0 = x
        for layer in self.layers:
            x = layer(x)
        return torch.einsum('nchw,nc->nchw', x0, x)

class UNetEncoderBlock(nn.Module):

    def __init__(self, in_channels, out_channels, dense_div, conv_num_in_block, conv_norm, dropout_rate):
        super().__init__()

        self.conv_block = ConvBlock(in_channels, out_channels, conv_num_in_block, conv_norm)
        self.senet_block = SENetBlock(out_channels, dense_div)
        self.dropout2d = nn.Dropout2d(dropout_rate)
        self.maxpool2d = nn.MaxPool2d((2, 2))

    def forward(self, x):
        # x -- (N, C, H, W)
        x = self.conv_block(x)
        x = self.senet_block(x)
        x = self.dropout2d(x)
        x_pool = self.maxpool2d(x)
        # x -- (N, C_new, H, W)
        # x_pool -- (N, C_new, H/2, W/2)
        return x_pool, x

class UNetDecoderBlock(nn.Module):

    def __init__(self, in_channels, out_channels, dense_div, conv_num_in_block, conv_norm):
        super().__init__()

        self.conv_block = ConvBlock(in_channels, out_channels, conv_num_in_block, conv_norm)
        self.senet_block = SENetBlock(out_channels, dense_div)
        self.up_sample = nn.UpsamplingBilinear2d(scale_factor=(2, 2))

    def forward(self, x, x_cat):
        # x -- (N, C1, H, W)
        # x_cat -- (N, C2, H, W)
        x = self.conv_block(x)
        x_ori = self.senet_block(x)
        x = self.up_sample(x_ori)
        # re -- (N, C1_new + C2, H, W)
        return torch.cat((x_cat, x), dim=1), x_ori

class UNetEncoder(nn.Module):
    """
    ParticleDetector Encoder1&2
    """

    def __init__(self, in_channels, num_coders, basic_pow, dense_div, conv_num_in_block, conv_norm, dropout_rate):
        super().__init__()
        self.unet_encoder_list = nn.ModuleList()
        self.cat_channels_list = []
        for idx in range(num_coders):
            out_channels = 2 ** basic_pow
            self.unet_encoder_list.append(
                UNetEncoderBlock(in_channels, out_channels, dense_div, conv_num_in_block, conv_norm, dropout_rate)
            )
            self.cat_channels_list.append(out_channels)
            in_channels = out_channels

        self.out_channels = in_channels
        self.final_cat_channels = in_channels

    def forward(self, x):
        x_cat_list = []
        for unet_encoder in self.unet_encoder_list:
            x, x_cat = unet_encoder(x)
            x_cat_list.append(x_cat)  # encoder输出的未经maxpool的特征图，用于和之后decoder的结果cat
        return x, x_cat_list

class UNetDecoder(nn.Module):
    """
    ParticleDetector Decoder
    """

    def __init__(self, in_channels, num_coders, basic_pow, dense_div, conv_num_in_block, conv_norm, cat_channels_list):
        super().__init__()

        self.unet_decoder_list = nn.ModuleList()
        for idx in range(num_coders):
            out_channels = 2 ** basic_pow
            self.unet_decoder_list.append(
                UNetDecoderBlock(in_channels, out_channels, dense_div, conv_num_in_block, conv_norm)
            )
            if idx == 0:
                self.final_cat_channels = out_channels
            in_channels = out_channels + cat_channels_list.pop()

        self.out_channels = in_channels

    def forward(self, x, x_cat_list):
        x_cat_final2 = None
        for unet_decoder in self.unet_decoder_list:
            x, x_ = unet_decoder(x, x_cat_list.pop())
            if x_cat_final2 is None:
                x_cat_final2 = x_  # 第一个 decoder 在 up_sample 前的输出即为 x_cat_final2
        return x, x_cat_final2

class UNetFinal(nn.Module):  # TODO: 最终可能需要修改此处命名方式，重新训练一下

    def __init__(self, in_channels, basic_pow, final_channels, dense_div, conv_num_in_block, conv_norm):
        super().__init__()
        out_channels = 2 ** basic_pow
        self.conv_block = ConvBlock(in_channels, out_channels, conv_num_in_block, conv_norm)
        self.senet_block = SENetBlock(out_channels, dense_div)
        self.final_conv1 = nn.Conv2d(out_channels, out_channels, (1, 1), padding=(0, 0))
        self.final_conv2 = nn.Conv2d(out_channels, final_channels, (1, 1), padding=(0, 0))
        self.sigmoid = nn.Sigmoid()
        self.tanh = nn.Tanh()
        self.softmax = nn.Softmax(dim=1)
        self.final_channels = final_channels
        self.max_num_in_cell = int((final_channels - 1) / 3)

    def forward(self, x):
        # x -- (N, C, H, W)
        x = self.conv_block(x)
        x = self.senet_block(x)
        x = self.final_conv1(x)
        x = self.final_conv2(x)
        x_score = self.softmax(x[:, :self.max_num_in_cell + 1, :, :])
        x_loc_center = self.sigmoid(x[:, self.max_num_in_cell + 1:self.max_num_in_cell + 3, :, :])
        x_loc_bias = self.tanh(x[:, self.max_num_in_cell + 3:, :, :])
        x_loc = torch.clip(torch.cat([x_loc_center] * 3, dim=1) + x_loc_bias, 0., 1.)
        # x -- (N, final_channels, H, W)
        return torch.cat([x_score, x_loc], dim=1)

class CoordsTranslater(nn.Module):
    """
    ParticleDetector 输出的坐标和特征图转为 两个 list：归一化的绝对坐标（float） 以及 对应坐标上的特征描述子
    len(list) = batch_size
    """

    def __init__(self, cell_size, max_num_in_cell) -> None:
        super().__init__()
        self.cell_size = cell_size
        self.max_num_in_cell = max_num_in_cell
        self.cell_grid = None

    def forward(self, x_loc, x_fea):
        # x_loc -- (N, max_num_in_cell + 3 + max_num_in_cell [ID], H1, W1) -- H1,W1 = full_H,W / cell_size
        # x_fea -- (N, 2 ** basic_pow * 2, H2, W2) -- H2,W2 = full_H,W
        batch_size, loc_channels, cell_h, cell_w = x_loc.shape
        fea_channels = x_fea.shape[1]
        if loc_channels != 1 + 4 * self.max_num_in_cell:
            raise ValueError(f'ParticleDetector output x_loc has {x_loc.shape[1]} channels, '
                             f'but "opt.max_num_in_cell": {self.max_num_in_cell} * 4 + 1 does not match this!')

        if self.cell_grid is None:
            cell_grid_h, cell_grid_w = torch.meshgrid(
                torch.arange(0, cell_h, dtype=torch.long, device=x_loc.device),
                torch.arange(0, cell_w, dtype=torch.long, device=x_loc.device),
                indexing='ij'
            )
            self.cell_grid = torch.cat([cell_grid_h[None], cell_grid_w[None]] * self.max_num_in_cell, dim=0)

        batch_abs_coords = self.cell_size * (
                x_loc[:, self.max_num_in_cell + 1:3 * self.max_num_in_cell + 1, :, :] + self.cell_grid
        )
        batch_in_cell = torch.argmax(x_loc[:, :self.max_num_in_cell + 1, :, :], dim=1).long()
        batch_ids = x_loc[:, -self.max_num_in_cell:, :, :].long()

        # 从此开始，batch_size 存为 list，长度为N
        batch_abs_coords_list = list()
        batch_features_list = list()
        batch_ids_list = list()
        for n in range(batch_size):
            in_cell = batch_in_cell[n]
            abs_coords = torch.empty((0, 2), dtype=torch.float32, device=x_loc.device)
            features = torch.empty((0, fea_channels), dtype=torch.float32, device=x_loc.device)
            ids = torch.empty((0,), dtype=torch.long, device=x_loc.device)
            for num_in_cell in range(1, self.max_num_in_cell + 1):  # 对于含有 num_in_cell 个粒子的 cell 进行遍历
                cell_loc_matrix = torch.eq(in_cell, num_in_cell)

                if torch.sum(cell_loc_matrix) != 0:  # 含有 num_in_cell 个粒子的 cell 的数量不为 0
                    abs_coords_temp = batch_abs_coords[n][:2 * num_in_cell, cell_loc_matrix].permute(1, 0).reshape(-1,
                                                                                2)
                    abs_coords_floor_temp = torch.floor(abs_coords_temp).long()
                    features_temp = x_fea[n][:, abs_coords_floor_temp[:, 0], abs_coords_floor_temp[:, 1]].permute(1, 0)
                    ids_temp = batch_ids[n][:num_in_cell, cell_loc_matrix].permute(1, 0).reshape(-1)

                    # 若 cell 中有 num_in_cell 个粒子，则存储 num_in_cell 份该坐标，特征向量（存储的ID不同，但也是num_in_cell个）
                    abs_coords = torch.cat([abs_coords, abs_coords_temp], dim=0)
                    features = torch.cat([features, features_temp], dim=0)
                    ids = torch.cat([ids, ids_temp], dim=0)

            batch_abs_coords_list.append(abs_coords)
            batch_features_list.append(features)
            batch_ids_list.append(ids)
        return batch_abs_coords_list, batch_features_list, batch_ids_list

class CoordsEncoder(nn.Module):
    """
    输出特征描述子

    """

    def __init__(self, basic_pow, desc_pow) -> None:
        super().__init__()
        mlp_out_channels_list = [2 ** desc_pow, 2 ** desc_pow]
        self.fea_proj = MLP(in_channel=2 ** (basic_pow + 1), out_channels=[2 ** (basic_pow + 1)])
        self.encoder = MLP(in_channel=2 + 2 ** (basic_pow + 1), out_channels=mlp_out_channels_list)

    def forward(self, coords_list, fea_list, img_h, img_w):
        desc_list = []
        for coords, fea in zip(coords_list, fea_list):
            coords = normalize_coords(coords.transpose(0, 1)[None], (img_h, img_w))
            fea = self.fea_proj(fea.transpose(0, 1)[None])
            fea = F.normalize(fea, p=2, dim=1)
            encoder_inputs = torch.cat([coords, fea], dim=1)
            desc_list.append(self.encoder(encoder_inputs))
        return desc_list

def attention(query, key, value, softmax):
    dim = query.shape[1]
    scores = torch.einsum('bdhn,bdhm->bhnm', query, key) / dim ** .5
    prob = softmax(scores)
    return torch.einsum('bhnm,bdhm->bdhn', prob, value), prob

class MultiHeadedAttention(nn.Module):
    """ Multi-head attention to increase model expressivity """

    def __init__(self, num_heads: int, d_model: int):
        super().__init__()
        """
        d_model is just feature_dim
        d_model 必须能够被 num_heads 整除
        """
        assert d_model % num_heads == 0
        self.dim = d_model // num_heads
        self.num_heads = num_heads

        # 用 nn.Conv1d 代替 linear，此时作用于每个关键点的权重都是相同的
        self.merge = nn.Conv1d(d_model, d_model, kernel_size=(1,), bias=True)
        self.proj = nn.ModuleList([deepcopy(self.merge) for _ in range(3)])
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, query, key, value):
        batch_dim = query.size(0)
        query, key, value = [l(x).reshape(batch_dim, self.dim, self.num_heads, -1)
                             for l, x in zip(self.proj, (query, key, value))]
        x, _ = attention(query, key, value, self.softmax)
        return self.merge(x.reshape(batch_dim, self.dim * self.num_heads, -1))

class AttentionalPropagation(nn.Module):
    def __init__(self, feature_dim: int, num_heads: int):
        super().__init__()
        """
        (default usage) num_heads = 4
        """
        self.attn = MultiHeadedAttention(num_heads, feature_dim)
        self.mlp = MLP(feature_dim * 2, [feature_dim * 2, feature_dim])

    def forward(self, x, source):
        message = self.attn(x, source, source)
        return self.mlp(torch.cat([x, message], dim=1))

class AttentionalGNN(nn.Module):

    def __init__(self, desc_channels, gnn_layer_pairs):
        """
        :param desc_channels:
        :param gnn_layer_pairs:

        layers = ['self', 'cross', 'self', 'cross', ...,  'self', 'cross'] length = 2 * gnn_layer_pairs
        every layer: AttentionalPropagation = attn + mlp
        attn(MultiHeadedAttention) = merge + proj
        """

        super().__init__()
        layer_names = ['self', 'cross'] * gnn_layer_pairs
        self.layers = nn.ModuleList([
            AttentionalPropagation(desc_channels, num_heads=4)
            for _ in range(len(layer_names))])
        self.names = layer_names

    def forward(self, desc0, desc1):
        for layer, name in zip(self.layers, self.names):
            if name == 'cross':
                src0, src1 = desc1, desc0
            else:  # if name == 'self':
                src0, src1 = desc0, desc1
            delta0, delta1 = layer(desc0, src0), layer(desc1, src1)
            desc0, desc1 = (desc0 + delta0), (desc1 + delta1)
        return desc0, desc1

def log_sinkhorn_iterations(z: torch.Tensor, log_mu: torch.Tensor, log_nu: torch.Tensor, iters: int):
    """
    Perform Sinkhorn Normalization in Log-space for stability
    Z: 即经过 einsum 以及 / self.config['descriptor_dim'] ** .5 ，并加入了 bin_scores 行和列的 scores 张量（矩阵）
    epsilon 取值默认为 1？
    """
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(z + u.unsqueeze(2), dim=1)
    return z + u.unsqueeze(2) + v.unsqueeze(1)

def log_optimal_transport(scores: torch.Tensor, alpha, iters: int):
    """
    Perform Differentiable Optimal Transport in Log-space for stability
    变量均为tensor，因此可以进行梯度的后向传播，端到端训练

    alphe: 即 bin_score，论文中的 z
    """
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m * one).to(scores), (n * one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], dim=-1),
                           torch.cat([bins1, alpha], dim=-1)], dim=1)

    norm = - (ms + ns).log()  # TODO: 此处关于log_sinkhorn_iterations的前处理原理还需要理解
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)  # 只是按batch_size扩展张量

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z

def normalize_coords(coords, image_shape):
    """ Normalize locations based on image image_shape"""
    size = coords.new_tensor(image_shape)[None, :, None]
    center = (size / 2)
    scaling = size.max(dim=1, keepdim=True).values * 0.5  # TODO: 此处指定的keypoints坐标normalize scaling值可能需要修改
    return (coords - center) / scaling  # 将值缩到 -1 ~ 1 之间

class TrajectoryManagerV1(nn.Module):
    """
    提供两帧之间的匹配管理
    """

    def __init__(self, points0):
        super().__init__()

    def forward(self, points1, matches1):
        """
            track_dict = {
                last_keypoint_idx: [start_frame_idx, (x0, y0), (x1, y1), ...],
                l1: ... ,
                l2: ... ,
            }

            last_keypoint_idx = 'finished_track_id_x' 则轨迹已停止
        """
        self.frame_idx += 1
        for idx, match in enumerate(matches1):
            if match > -1:
                # 完成匹配的关键点添加到对应的原轨迹尾部
                self.tracking_dict['last:' + str(match)].append(tuple(points1[idx]))
                # 将 last_keypoint_idx 更新为 now_keypoint_idx
                self.tracking_dict.update({'now:' + str(idx): self.tracking_dict.pop('last:' + str(match))})
            else:
                self.tracking_dict['now:' + str(idx)] = [self.frame_idx, tuple(points1[idx])]

        for key in list(self.tracking_dict):
            if key.startswith('last'):
                self.track_finished.append(self.tracking_dict.pop(key))
        for key in list(self.tracking_dict):
            self.tracking_dict.update({key.replace('now', 'last'): self.tracking_dict.pop(key)})

    def stop(self, length_threshold=None):
        for key in list(self.tracking_dict):
            self.track_finished.append(self.tracking_dict.pop(key))
        if length_threshold is not None:
            self._length_filter(length_threshold)

    def _length_filter(self, length_threshold):
        for track in self.track_finished:
            if len(track) - 1 >= length_threshold:  # -1 为了排除列表中指定起始帧的元素（第一个数）
                self.track_length_filtered.append(track)


import functools

from flask import g

class Middleware:
    def __call__(self, func):
        @functools.wraps(func)
        def _wrapper(*args, **kwargs):
            if "subject" not in g:
                return {
                    "code": "authentication_failure",
                    "message": "Authentication required to access endpoint.",
                }, 401
            return func(*args, **kwargs)

        return _wrapper


from django.conf.urls import patterns, include, url
from django.contrib import admin
from django.views.generic.base import RedirectView

urlpatterns = patterns('',
	url(r'^$', RedirectView.as_view(url='/principal/login/')),
    url(r'^principal/', include('principal.urls', namespace="principal")),
    url(r'^empresas/', include('empresas.urls', namespace="empresas")),
    url(r'^cotizaciones/', include('cotizaciones.urls', namespace="cotizaciones")),
    url(r'^contactos/', include('contactos.urls', namespace="contactos")),
    url(r'^oauth2callback', 'contactos.views.auth_return'),
    url(r'^grappelli/', include('grappelli.urls')),
    url(r'^admin/', include(admin.site.urls)),
)


from dataclasses import dataclass, field
from typing import Optional

from .area_rando_types import AreaDoor
from .connection_data import SunkenNestL, area_doors
from .game import Game
from .item_data import Items
from .loadout import Loadout
from .location_data import Location, spacePortLocs
from .logic_goal import can_win
from .logic_shortcut_data import can_fall_from_spaceport
from .logic_updater import updateLogic
from .trick import Trick
from .trick_data import Tricks

@dataclass
class Sphere:
    fallen: bool
    pickups: dict[str, str] = field(default_factory=dict)
    new_doors: list[str] = field(default_factory=list)

@dataclass
class PlayThrough:
    spheres: list[Sphere] = field(default_factory=list)

def solve(game: Game,
          starting_items: Optional[Loadout] = None,
          excluded_door: Optional[AreaDoor] = None) -> tuple[bool, PlayThrough, list[Location]]:
    """ returns (whether completable, play through, accessible locations) """
    for loc in game.all_locations.values():
        loc['inlogic'] = False

    unused_locations: list[Location] = list(game.all_locations.values())
    used_locs: set[str] = set()
    doors_accessed: set[AreaDoor] = set()

    loadout = Loadout(game, starting_items)

    play_through = PlayThrough()

    def check_for_new_area_doors() -> None:
        new_area_doors: list[str] = []
        for thing in loadout:
            if isinstance(thing, AreaDoor) and thing not in doors_accessed:
                new_area_doors.append(thing.name)
                doors_accessed.add(thing)
        if len(new_area_doors):
            # print(new_area_doors)
            if len(play_through.spheres) == 0:
                play_through.spheres.append(Sphere(False))
            play_through.spheres[-1].new_doors.extend(new_area_doors)

    # this loop just for spaceport
    stuck = False
    while not stuck:
        prev_loadout_count = len(loadout)
        updateLogic(unused_locations, loadout, excluded_door)
        check_for_new_area_doors()
        play_through.spheres.append(Sphere(False))
        for loc in unused_locations:
            if loc['inlogic']:
                loc_name = loc['fullitemname']
                if loc_name not in spacePortLocs:
                    # debug
                    # print(f"found {loc_name} in logic while still in spaceport")
                    continue
                item = loc['item']
                if item:
                    loadout.append(item)
                    play_through.spheres[-1].pickups[loc_name] = item[0]
                used_locs.add(loc_name)
        # remove used locations
        unused_locations = [loc for loc in unused_locations if loc['fullitemname'] not in used_locs]
        stuck = len(loadout) == prev_loadout_count

    while (
        len(play_through.spheres) > 0 and
        len(play_through.spheres[-1].new_doors) + len(play_through.spheres[-1].pickups) == 0
    ):
        play_through.spheres.pop()

    if can_fall_from_spaceport not in loadout:
        # print("solver: couldn't get out of spaceport")
        # for loc in unused_locations:
        #     if loc['inlogic'] and loc['fullitemname'] not in spacePortLocs:
        #         print("solver: found another way out of spaceport besides Ridley")
        #         print(loadout)
        #         print("but logic doesn't support that yet")
        return False, play_through, [loc for loc in game.all_locations.values() if loc["fullitemname"] in used_locs]
    loadout.append(Items.spaceDrop)
    loadout.append(SunkenNestL)  # assuming this is where we land

    stuck = False
    while not stuck:
        prev_loadout_count = len(loadout)
        updateLogic(unused_locations, loadout, excluded_door)
        check_for_new_area_doors()
        play_through.spheres.append(Sphere(True))
        for loc in unused_locations:

            # special case: major/minor can put missiles or grav boots in sandy cache even though it's not in logic
            # old major minor code obsolete, but it doesn't hurt
            if game.options.fill_choice == "MM" and loc['fullitemname'] == "Sandy Cache" and loc['item'] in {
                Items.GravityBoots, Items.Missile
            }:
                loc['inlogic'] = True

            if loc['inlogic']:
                loc_name = loc['fullitemname']
                item = loc['item']
                if item:
                    loadout.append(item)
                    play_through.spheres[-1].pickups[loc_name] = item[0]
                used_locs.add(loc_name)
        # remove used locations
        unused_locations = [loc for loc in unused_locations if loc['fullitemname'] not in used_locs]
        stuck = len(loadout) == prev_loadout_count

    while (
        len(play_through.spheres) > 0 and
        len(play_through.spheres[-1].new_doors) + len(play_through.spheres[-1].pickups) == 0
    ):
        play_through.spheres.pop()

    # note: the reason for making a new list from all_locations rather than used_locs,
    # is that used_locs is a `set`, so iterating through it is not deterministic, so seeds wouldn't be reproducible
    return (
        (can_win in loadout),
        # len(unused_locations) == 0,
        play_through,
        [loc for loc in game.all_locations.values() if loc["fullitemname"] in used_locs]
    )

def spoil_play_through(play_through: PlayThrough) -> list[str]:
    items_to_display = {
        "Missile": 999,
        "Morph Ball": 999,
        "Gravity Boots": 999,
        "Super Missile": 999,
        "Grapple Beam": 999,
        "Power Bomb": 999,
        "Speed Ball": 999,
        "Bombs": 999,
        "HiJump": 999,
        "Aqua Suit": 999,
        "Dark Visor": 999,
        "Wave Beam": 999,
        "Speed Booster": 999,
        "Spazer": 999,
        "Varia Suit": 999,
        "Ice Beam": 999,
        "Metroid Suit": 999,
        "Plasma Beam": 999,
        "Screw Attack": 999,
        "Space Jump": 999,
        "Charge Beam": 999,
        "Hypercharge": 999,
        "X-Ray Scope": 999,
        "Energy Tank": 5,
        "Small Ammo": 4,
        "Large Ammo": 6,
        "Space Jump Boost": 1,
        "Damage Amp": 1,
        "Accel Charge": 1,
        "Refuel Tank": 0,
    }

    seen_fallen = False
    log_lines = [" - spaceport -"]
    for sphere in play_through.spheres:
        if sphere.fallen and not seen_fallen:
            log_lines.append(" - fall from spaceport -")
            seen_fallen = True

        log_lines.append("sphere:")
        for loc_name, item in sphere.pickups.items():
            if items_to_display[item] > 0:  # TODO: or hard required
                log_lines.append(f"    get {item} from {loc_name}")
                items_to_display[item] -= 1
        # TODO: if didn't display any pickups in sphere, display all pickups in sphere (unless last sphere)
        if len(sphere.new_doors):
            log_lines.append(f"  new area doors: {', '.join(sphere.new_doors)}")

    # TODO: delete "sphere:" from end
    return log_lines

def _locations_in_prog_order(play_through: PlayThrough) -> dict[str, int]:
    """
    { location name: sphere number }

    0 sphere means not in log lines

    If locations aren't in the solve log lines, they're at the end of this list.
    """
    locations: dict[str, int] = {}

    for sphere_i, sphere in enumerate(play_through.spheres):
        for loc_name in sphere.pickups:
            locations[loc_name] = sphere_i

    return locations

def hard_required_locations(game: Game) -> tuple[dict[str, int], PlayThrough]:
    """ list of names of hard required locations in progression order """
    completable, play_through, _ = solve(game)
    if not completable:
        # not sure what I want to do with this function if I pass a game that isn't completable
        # maybe exception
        return {}, play_through

    locations = _locations_in_prog_order(play_through)

    req_locs: dict[str, int] = {}
    for excluded_loc_name in locations:
        excluded_loc = game.all_locations[excluded_loc_name]
        saved_item = excluded_loc['item']
        excluded_loc['item'] = None
        completable, _, _ = solve(game)
        if not completable:
            req_locs[excluded_loc_name] = locations[excluded_loc_name]
        excluded_loc['item'] = saved_item

    return req_locs, play_through

def _excluded_tricks(t: Trick) -> set[Trick]:
    """
    all of the tricks that should be excluded to exclude this trick

    "If you can't do the easier one, then you can't do the harder one."
    """
    data = {
        Tricks.hell_run_easy:
            {Tricks.hell_run_easy, Tricks.hell_run_medium, Tricks.hell_run_hard},
        Tricks.hell_run_medium:
            {Tricks.hell_run_medium, Tricks.hell_run_hard},

        Tricks.dark_easy:
            {Tricks.dark_easy, Tricks.dark_medium, Tricks.dark_hard},
        Tricks.dark_medium:
            {Tricks.dark_medium, Tricks.dark_hard},

        Tricks.movement_moderate:
            {Tricks.movement_moderate, Tricks.movement_zoast},

        Tricks.sbj_underwater_w_hjb:
            {Tricks.sbj_underwater_w_hjb, Tricks.sbj_underwater_no_hjb},

        Tricks.wall_jump_precise:
            {Tricks.wall_jump_precise, Tricks.wall_jump_delayed},

        Tricks.morph_jump_4_tile:
            {Tricks.morph_jump_4_tile, Tricks.morph_jump_3_tile},

        Tricks.crouch_or_downgrab:
            {Tricks.crouch_or_downgrab, Tricks.crouch_precise},

        Tricks.short_charge_2:
            {Tricks.short_charge_2, Tricks.short_charge_3, Tricks.short_charge_4},

        Tricks.uwu_2_tile:
            {Tricks.uwu_2_tile, Tricks.uwu_2_tile_surface}
    }
    if t in data:
        return data[t]
    return {t}

def obsoletes(t: str) -> set[str]:
    data = {
        "hell_run_hard": {"hell_run_medium", "hell_run_easy"},
        "hell_run_medium": {"hell_run_easy"},
        "dark_hard": {"dark_medium", "dark_easy"},
        "dark_medium": {"dark_easy"},
        "movement_zoast": {"movement_moderate"},
    }
    if t in data:
        return data[t]
    return set()

def required_tricks(game: Game) -> tuple[list[str], list[str]]:
    """ lists of names of (tricks required to win, tricks required to get all locations) """
    completable, _, _ = solve(game)
    if not completable:
        # not sure what I want to do with this function if I pass a game that isn't completable
        # maybe exception
        return [], []

    req_for_win: list[str] = []
    req_for_locs: list[str] = []
    all_tricks = {t: n for n, t in vars(Tricks).items() if isinstance(t, Trick)}
    tricks_allowed = game.options.logic
    for excluded_trick in tricks_allowed:
        game.options.logic = tricks_allowed - _excluded_tricks(excluded_trick)
        completable, _, locs = solve(game)
        if not completable:
            req_for_win.append(all_tricks[excluded_trick])
        if len(locs) != 122:
            req_for_locs.append(all_tricks[excluded_trick])
    game.options.logic = tricks_allowed

    obsoleted: set[str] = set()
    for trick_name in req_for_win:
        obsoleted.update(obsoletes(trick_name))
    req_for_win = [trick_name for trick_name in req_for_win if trick_name not in obsoleted]

    obsoleted = set()
    for trick_name in req_for_locs:
        obsoleted.update(obsoletes(trick_name))
    req_for_locs = [trick_name for trick_name in req_for_locs if trick_name not in obsoleted]

    return req_for_win, req_for_locs

def required_doors(loadout: Loadout, loc_name: str) -> list[str]:
    """ what doors are required to get to this location with this loadout """
    temp_loadout = Loadout(loadout.game, loadout)
    _, play_through, acc_locs = solve(temp_loadout.game, temp_loadout)

    assert temp_loadout.game.all_locations[loc_name] in acc_locs, f"{loc_name} not in logic"

    doors: list[str] = []
    for sphere in play_through.spheres:
        for door_name in sphere.new_doors:
            doors.append(door_name)
            temp_loadout.contents[area_doors[door_name]] = 0

    tr: list[str] = []
    for excluded_door in doors:
        # empty doors from loadout
        for thing in temp_loadout:
            if isinstance(thing, AreaDoor):
                temp_loadout.contents[thing] = 0

        _, _, acc_locs = solve(temp_loadout.game, temp_loadout, area_doors[excluded_door])
        if temp_loadout.game.all_locations[loc_name] not in acc_locs:
            tr.append(excluded_door)

    return tr


#!/usr/bin/env python3
import tApi.tonlib as TL
import asyncio
import json

import time

from loguru import logger

async def tick():
    try:
        while True:
            logger.warning('Tick')
            await asyncio.sleep(1)
    except asyncio.CancelledError:
        logger.warning(f"Tick function canceled")

async def set_verbosity_level(tonlib, level=0):
    # set verbosity
    request = {
        '@type': 'setLogVerbosityLevel',
        'new_verbosity_level': level
    }
    await tonlib.execute(request)

async def init_tonlib(tonlib):
    with open('settings/testnet.json', 'r') as f:
        tonlib_config = json.load(f)

    # setup config
    keystore_obj = {
        '@type': 'keyStoreTypeDirectory',
        'directory': './ton_keystore/'
    }
    request = {
        '@type': 'init',
        'options': {
            '@type': 'options',
            'config': {
                '@type': 'config',
                'config': json.dumps(tonlib_config),
                'use_callbacks_for_network': False,
                'blockchain_name': '',
                'ignore_cache': False
            },
            'keystore_type': keystore_obj
        }
    }
    await tonlib.execute(request)
    logger.info(F"TonLib inited successfully")

async def main():
    loop = asyncio.get_running_loop()
    tick_future = asyncio.ensure_future(tick(), loop=loop)

    # init tonlibjson
    tonlib = TL.TonLib(loop, -1, None)

    await set_verbosity_level(tonlib, level=0)
    await init_tonlib(tonlib)

    # send test request
    futures = []
    for i in range(10):
        masterchain_request = {'@type': 'blocks.getMasterchainInfo'}
        future = asyncio.ensure_future(tonlib.execute(masterchain_request), loop=loop)
        logger.info(f"Request {i} dispatched")
        futures.append(future)
        await asyncio.sleep(0.2)

    for future in futures:
        res = await future
        logger.info(res)

    logger.info('Finished')
    tick_future.cancel()
    return

if __name__ == "__main__":
    start = time.time()
    asyncio.run(main())
    elapsed = time.time() - start
    logger.error(F"Elapsed time: {elapsed}")


from abc import ABC, abstractmethod
from sqlalchemy import select
from sqlalchemy.orm import aliased

from src.settings.db import async_session_maker
from src.apps.models import *

class AbstractRepository(ABC):
    @abstractmethod
    async def get_all():
        raise NotImplementedError

    @abstractmethod
    async def get_deffirence(self, city):
        raise NotImplementedError

class SQLAlchemyRepository(AbstractRepository):
    model = None

    async def get_all(self):
        async with async_session_maker() as session:
            query = select(self.model)
            result = await session.execute(query)
            return result.all()

    async def get_deffirence(self, city):
        raise NotImplementedError

class SQLAlcemyRepositoryWithTwoModels(AbstractRepository):
    model_1 = None
    model_2 = None

    async def get_all():
        raise NotImplementedError

    async def get_deffirence(self, city):
        async with async_session_maker() as session:
                cities = aliased(self.model_1)
                cities_2 = aliased(self.model_1)
                Roads = aliased(self.model_2)

                neighboring_cities = (select(
                    cities.name.label("city"),
                    cities_2.name.label("target_city"),
                    Roads.distance.label("distance")
                )
                .join(Roads, Roads.previous_city == cities.id)
                .join(cities_2, cities_2.id == Roads.next_city)
                .filter(cities.name == city)).cte("neighboring_cities")

                query = (
                    select(
                            neighboring_cities,
                            sa.func.lead(neighboring_cities.c.distance)
                            .over(
                                    partition_by=neighboring_cities.c.city,
                                    order_by=neighboring_cities.c.target_city
                                )
                            - neighboring_cities.c.distance
                        )

                    )

                result = await session.execute(query)
                return result.all()


from pathlib import Path
from typing import List

from models.custom_model import CustomModel, collection_name
from models.data_models.person_model import Person

@collection_name("theses")
class Thesis(CustomModel):
    title: str
    authors: List[str]
    advisors: List[str]
    keywords: List[str]
    version: int = 0
    balance: int = 2022

    @classmethod
    def from_advisors(cls, advisors: List[Person]):
        names = set([a.name for a in advisors])

        for thesis in cls.all():
            if set(thesis.advisors) & names:
                yield thesis

    def check(self):
        if not self.title:
            raise ValueError("El título no puede ser vacío.")

        if len(self.authors) == 0:
            raise ValueError("Debe tener al menos un autor.")

        if len(self.advisors) == 0:
            raise ValueError("Debe tener al menos un tutor.")

        if len(self.keywords) < 3:
            raise ValueError("Debe tener al menos 3 palabras clave.")

        return True

    def save_thesis_pdf(self, pdf):
        self.version += 1
        name_pdf = f"{self.uuid}_v{self.version}.pdf"
        path: Path = Path(f"/src/data/Thesis/files/{name_pdf}")
        path.parent.mkdir(exist_ok=True, parents=True)

        with path.open("wb") as f:
            f.write(pdf.getbuffer())


from playsound import playsound
import random

while(1):
    seq_len = input("请输入序列长度：")
    if len(seq_len) > 0:
        seq_len = int(seq_len)
        break

while(1):
    sequence = []
    for i in range(seq_len):
        num = random.randint(0, 9)
        sequence.append(num)

    for i in sequence:
        filename = "number_audios/audio_%d.mp3"%(i)
        playsound(filename)

    sequence_test= list(map(int, list(input("请输入听到的数字序列："))))

    print("实际值：" + "".join(list(map(str, sequence))))
    print("测量值：" + "".join(list(map(str, sequence_test))))

    print("测试结果：" + str(sequence_test == sequence))

    control = input("继续Y/退出n:")
    if control == "" or control == "Y" or control == "y":
        continue
    else:
        break


#!/usr/bin/env python

from astropy import log
from astropy.io import fits
from astroquery.mast import Observations
from astropy import units
import numpy as np
import os

desc="""
Download tess data from MAST for QLP and SPOC pipelines
"""

from . import tessutils

data_path = os.environ.get('QUICKTESS_DIR', None)

if data_path is None:
    print("Please set the QUICKTESS_DIR environment variable to path ot the downloaded data")

def download(TIC=None, ra=None, dec=None,
             pipeline='QLP',
             download_dir=data_path,
             verbose=False):

    full_path = None
    if TIC is None:
        assert(ra is not None)
        assert(dec is not None)

        obs_table = Observations.query_region(
                f'{ra} {dec}',
                radius=5.0*units.arcsecond, verbose=verbose)
        idx = np.where(obs_table['provenance_name'] == pipeline)[0]
        full_path = []
        if len(idx) != 0:
            TIC_arr = np.array(obs_table['target_name'][idx])
            if len(np.unique(TIC_arr)) != 1:
                idx = [i for i in idx
                       if obs_table['target_name'][i]==obs_table['target_name'][idx[0]] ]

            TIC = obs_table['target_name'][idx[0]]
            for i in range(len(idx)):
                print(obs_table['obs_id'][idx[i]])
                full_path.append(
                        os.path.join(
                                data_path,
                                'mastDownload/HLSP',
                                obs_table['obs_id'][idx[i]],
                                obs_table['obs_id'][idx[i]]+'.fits'))
            if pipeline == 'TESS-SPOC':
                full_path = [ x.replace('_tp.fits', '_lc.fits')
                              for x in full_path ]
        else:
            log.info(f'No TESS data on MAST for {ra} {dec}')
            return None

    obsTable = Observations.query_criteria(provenance_name=pipeline,
                                           target_name=TIC, verbose=verbose)

    if len(obsTable) == 0:
        return None

    data = Observations.get_product_list(obsTable, verbose=verbose)
    with tessutils.HiddenPrints():
        download_lc = Observations.download_products(
                data, download_dir=download_dir)

    if full_path is None:
        if pipeline == 'QLP':
            full_path = [ os.path.join(download_dir, 'mastDownload/HLSP',
                                       x, x+'.fits')
                          for x in obsTable['obs_id'] ]
        elif pipeline == 'TESS-SPOC':
            full_path = [ os.path.join(download_dir, 'mastDownload/HLSP',
                                       x, x.replace('_tp', '_lc.fits'))
                          for x in obsTable['obs_id'] ]
        else:
            raise ValueError(
                    f'pipeline name {pipeline} not recognized')

    return tessutils.sort_paths(full_path)


import logging
from pathlib import Path
from typing import List

import numpy as np
import pandas as pd
import pytorch_lightning as pl
import torch
import torchvision
from torch.utils.data.sampler import SubsetRandomSampler

logger = logging.getLogger(__name__)

class FolderDataset(torchvision.datasets.ImageFolder):
    def __init__(self, root: str, transform=[], target_transform=[]):
        transform = torchvision.transforms.Compose(transform)
        target_transform = torchvision.transforms.Compose(target_transform)
        super().__init__(root=root, transform=transform, target_transform=target_transform)

    def to_csv(self, filename, indices: List[int] = None):
        df = pd.DataFrame(self.samples, columns=["filename", "label"])
        df["class"] = df["label"].apply(lambda x: self.classes[x])
        if indices:
            df = df.iloc[indices]
        df.to_csv(filename, index=False)

class FolderDataModule(pl.LightningDataModule):
    def __init__(self, root_dir, train_transforms=[], val_transforms=[], val_size=0.1, batch_size=16):
        super().__init__()
        self.root_dir = root_dir
        self.transforms = train_transforms
        if val_transforms:
            logger.info("val_transforms not supported for this dataset, using train_transforms for all datasets")

        self.val_size = val_size
        self.batch_size = batch_size

    def prepare_data(self):
        self.dataset = FolderDataset(self.root_dir, transform=self.transforms)
        train_idx, valid_idx = self.get_sampler_indices(self.val_size)
        self.train_sampler = SubsetRandomSampler(train_idx)
        self.val_sampler = SubsetRandomSampler(valid_idx)

    def get_sampler_indices(self, val_size=0.1, shuffle=True, random_seed=14):
        num_train = len(self.dataset)
        indices = list(range(num_train))
        split = int(np.floor(val_size * num_train))

        if shuffle:
            np.random.seed(random_seed)
            np.random.shuffle(indices)

        train_idx, valid_idx = indices[split:], indices[:split]
        return train_idx, valid_idx

    def train_dataloader(self):
        return torch.utils.data.DataLoader(
            self.dataset, batch_size=self.batch_size, sampler=self.train_sampler, num_workers=4
        )

    def val_dataloader(self):
        return torch.utils.data.DataLoader(
            self.dataset, batch_size=self.batch_size, sampler=self.val_sampler, num_workers=4
        )

    @property
    def len_train(self):
        return len(self.train_sampler)

    @property
    def len_valid(self):
        return len(self.val_sampler)

    def to_csv(self, output_dir):
        train_csv = Path(output_dir, "train.csv")
        self.dataset.to_csv(train_csv, indices=self.train_sampler.indices)

        val_csv = Path(output_dir, "val.csv")
        self.dataset.to_csv(val_csv, indices=self.val_sampler.indices)


import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, pre_trained_enbedding, filter_sizes, num_filters, dropout, output_size):
        super(CNN, self).__init__()
        # filter_sizes = [1,2,3,4,5]
        # num_filters = 20
        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.word_embeddings.weight.data.copy_(pre_trained_enbedding)
        self.word_embeddings.weight.requires_grad = False

        self.convs1 = nn.ModuleList([nn.Conv2d(1, num_filters, (K, embedding_dim)) for K in filter_sizes])
        self.dropout = nn.Dropout(dropout)
        self.fc1 = nn.Linear(len(filter_sizes)*num_filters, output_size)

    def forward(self,inputs):
        embeds = self.word_embeddings(inputs)
        x = embeds.view(embeds.size(0), 1, embeds.size(1), embeds.size(2))
        x = [F.relu(conv(x)).squeeze(3) for conv in self.convs1]
        x = [F.max_pool1d(i, i.size(2)).squeeze(2) for i in x]
        x = torch.cat(x, 1)
        x = self.dropout(x)
        out = self.fc1(x)
        probs = F.log_softmax(out, dim=1)

        return probs

Sentiment Analysis Deep Learning/model/TextCNN.py


import pennylane as qml
from pennylane import numpy as np

num_wires = 5
dev1 = qml.device('default.qubit', wires=num_wires)

@qml.qnode(dev1)
def circuit():
	output = []

	# Encoding the data (Variational Encoding)
	qml.RX(1.01, wires=[0])
	qml.RX(0.81, wires=[1])
	qml.RX(0.81, wires=[2])
	qml.RX(1.3e-2, wires=[3])
	qml.RX(0.16, wires=[4])

	qml.RZ(1.01, wires=[0])
	qml.RZ(0.81, wires=[1])
	qml.RZ(0.81, wires=[2])
	qml.RZ(1.3e-2, wires=[3])
	qml.RZ(0.16, wires=[4])

	# Making the Parameterised QUantum Circuit (PQC)
	qml.RX(1.10, wires=[0])
	qml.RZ(0.75, wires=[0])
	qml.CNOT(wires=[0, 1])
	qml.RX(0.82, wires=[0])
	qml.RX(0.84, wires=[1])
	qml.RZ(0.53, wires=[0])
	qml.RZ(0.81, wires=[1])
	qml.CNOT(wires=[1, 2])
	qml.CNOT(wires=[0, 1])
	qml.RX(0.60, wires=[2])

	qml.RX(0.56, wires=[1])
	qml.RZ(0.98, wires=[2])
	qml.RZ(0.93, wires=[1])
	qml.CNOT(wires=[2, 3])
	qml.CNOT(wires=[1, 2])
	qml.RX(0.34, wires=[3])

	qml.RX(0.21, wires=[2])
	qml.RZ(1.20, wires=[3])
	qml.RZ(8.7e-2, wires=[2])
	qml.CNOT(wires=[3, 4])
	qml.CNOT(wires=[2, 3])
	qml.RX(0.90, wires=[4])

	qml.RX(-0.18, wires=[3])
	qml.RZ(0.49, wires=[4])
	qml.RZ(0.83, wires=[3])
	qml.CNOT(wires=[3, 4])

	qml.RX(0.72, wires=[4])
	qml.RZ(87, wires=[4])

	for i in range(num_wires):
		output.append(qml.expval(qml.PauliZ(i)))

	return output

print(circuit())


import socket
import tqdm
import os

SEPERATOR = "[SEPARATOR]"
BUFFER_SIZE = 4096 #How many data sent each step

#HomeServer ip adress
host = "192.168.1.5"

# port
port = 5001

filename = "test.fnk"

filesize = os.path.getsize(filename)

s= socket.socket()

print(f"[+] Connecting to {host}:{port}")
try:
	s.connect((host,port))
	print("[+] Connection Sucessfull!")
	pass
except Exception as e:
	print(f"[!] Connection unsucessfull!")
	print(f"[!] error code: {e}")
	exit()

s.send(f"{filename}{SEPERATOR}{filesize}".encode("utf-8"))

progess = tqdm.tqdm(range(filesize), f"Sending{filename}", unit = "B", unit_scale = True, unit_divisor= 1024)

with open(filename, "rb") as f:
	while True:
		bytes_read = f.read(BUFFER_SIZE)
		if not bytes_read:
			break

		s.sendall(bytes_read)

		progess.update(len(bytes_read))

s.close()