counters.py

from __future__ import annotations

from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from overcooked_simulator.overcooked_environment import GameScore

import numpy as np
import numpy.typing as npt

from overcooked_simulator.game_items import (
    CuttableItem,
    Item,
    ProgressibleItem,
    Plate,
    Tomato,
    Pot,
    Pan,
)


class Counter:
    """Simple class for a counter at a specified position (center of counter). Can hold things on top."""

    def __init__(self, pos: npt.NDArray[float], occupied_by: Optional[Item] = None):
        self.pos: npt.NDArray[float] = pos
        self.occupied_by: Optional[Item] = occupied_by

    def pick_up(self):
        """Gets called upon a player performing the pickup action. If the counter can give something to
        the player, it does so. In the standard counter this is when an item is on the counter.

        Returns: The item which the counter is occupied by. None if nothing is there.

        """
        give_player = self.occupied_by
        self.occupied_by = None
        return give_player

    def can_drop_off(self, item: Item) -> bool:
        """Checks whether an item by the player can be dropped of. More relevant for example with
        ingredient dispensers, which should always be occupied and cannot take an item.

        Args:
            item: The item for which to check, if it can be placed on the counter.

        Returns: True if the item can be placed on the counter, False if not.

        """
        return self.occupied_by is None or self.occupied_by.can_combine(item)

    def drop_off(self, item: Item) -> Item | None:
        """Takes the thing dropped of by the player.

        Args:
            item: The item to be placed on the counter.

        Returns: TODO Return information, whether the score is affected (Serving Window?)

        """
        if self.occupied_by is None:
            self.occupied_by = item
        elif self.occupied_by.can_combine(item):
            self.occupied_by.combine(item)
        return None

    def interact_start(self):
        """Starts an interaction by the player. Nothing happens for the standard counter."""
        pass

    def interact_stop(self):
        """Stops an interaction by the player. Nothing happens for the standard counter."""
        pass

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(pos={self.pos},occupied_by={self.occupied_by})"
        )


class CuttingBoard(Counter):
    def __init__(self, pos: np.ndarray):
        self.progressing = False
        super().__init__(pos)

    def progress(self):
        """Called by environment step function for time progression"""
        if self.progressing:
            if isinstance(self.occupied_by, CuttableItem):
                self.occupied_by.progress()

    def start_progress(self):
        """Starts the cutting process."""
        self.progressing = True

    def pause_progress(self):
        """Pauses the cutting process"""
        self.progressing = False

    def interact_start(self):
        """Handles player interaction, starting to hold key down."""
        self.start_progress()

    def interact_stop(self):
        """Handles player interaction, stopping to hold key down."""
        self.pause_progress()


class ServingWindow(Counter):
    def __init__(self, pos, game_score: GameScore):
        self.game_score = game_score
        super().__init__(pos)

    def drop_off(self, item) -> Item | None:
        reward = 5
        # TODO define rewards
        self.game_score.increment_score(reward)
        return None

    def can_score(self, item):
        if isinstance(item, Plate) and isinstance(item.holds, ProgressibleItem):
            return item.holds.finished

    def can_drop_off(self, item: Item) -> bool:
        return self.can_score(item)

    def pick_up(self):
        return None


class PlateReturn(Counter):
    def __init__(self, pos):
        super().__init__(pos)
        self.occupied_by = [Plate()]

    def pick_up(self):
        """Gets called upon a player performing the pickup action. Gives back a plate (possibly with ingredient.

        Returns: A plate possibly with an ingredient on it.

        """
        give_player = self.occupied_by.pop()
        if not self.occupied_by:
            self.occupied_by.append(Plate())
        return give_player

    def drop_off(self, item: Item) -> Item | None:
        """Takes the ingredient dropped of by the player.

        Args:
            item: The ingredient to be placed on the counter.
        """
        if isinstance(item, Plate):
            if self.occupied_by[-1].holds:
                return item
            self.occupied_by.append(item)
            return None
        if self.occupied_by[-1].can_combine(item):
            self.occupied_by[-1].combine(item)
            return None
        return item

    def can_drop_off(self, item: Item) -> bool:
        """Checks whether an ingredient by the player can be dropped of.

        Args:
            item: The ingredient for which to check, if it can be placed on the counter.

        Returns: True if the ingredient can be placed on the counter, False if not.

        """
        # possibility to drop off empty plate on empty plate return
        return (
            isinstance(self.occupied_by[-1], Plate) and isinstance(item, Plate)
        ) or self.occupied_by[-1].can_combine(item)


class TomatoDispenser(Counter):
    def __init__(self, pos):
        super().__init__(pos)

    def pick_up(self):
        return Tomato()

    def drop_off(self, item: Item) -> Item | None:
        return None

    def can_drop_off(self, item: Item) -> bool:
        return False


class Trash(Counter):
    def pick_up(self):
        pass

    def drop_off(self, item: Item) -> Item | None:
        if isinstance(item, Plate):
            item.holds = None
            return item
        return None

    def can_drop_off(self, item: Item) -> bool:
        return True


class Stove(Counter):
    def __init__(self, pos: npt.NDArray[float], occupied_by: Optional[Item] = ...):
        if occupied_by is ...:
            occupied_by = Pot()
        super().__init__(pos, occupied_by)

    def drop_off(self, item) -> Item | None:
        if isinstance(item, (Pot, Pan)):
            self.occupied_by = item
        else:
            self.occupied_by.combine(item)
        return None

    def can_drop_off(self, item) -> bool:
        if isinstance(item, (Pot, Pan)):
            return not self.occupied_by
        if self.occupied_by:
            return self.occupied_by.can_combine(item)
        return False