counters.py

from __future__ import annotations

import logging
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,
    CookingEquipment,
    Meal,
)

log = logging.getLogger(__name__)


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, on_hands: bool = True):
        """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.

        """
        if on_hands:
            if self.occupied_by:
                occupied_by = self.occupied_by
                self.occupied_by = None
                return occupied_by
            return None
        if self.occupied_by and isinstance(self.occupied_by, CookingEquipment):
            return self.occupied_by.release()
        occupied_by = self.occupied_by
        self.occupied_by = None
        return occupied_by

    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):
            return 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
        log.debug(f"Drop off {item}")
        # TODO define rewards
        self.game_score.increment_score(reward)
        return None

    def can_score(self, item):
        if (
            isinstance(item, CookingEquipment)
            and "Plate" in item.name
            and item.content is not None
        ):
            if isinstance(item.content, Meal) and item.content.progressed_steps:
                return item.content.finished

            # TODO: Salad can be always be served. Check, if all needed parts are present.
            if item.content.name in ["Salad"]:
                return True
        return False

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

    def pick_up(self, on_hands: bool = True):
        pass


class Dispenser(Counter):
    def __init__(self, pos, dispensing):
        self.dispensing = dispensing
        super().__init__(
            pos,
            self.dispensing.create_item(),
        )

    def pick_up(self, on_hands: bool = True):
        return_this = self.occupied_by
        self.occupied_by = self.dispensing.create_item()
        return return_this

    def drop_off(self, item: Item) -> Item | None:
        if self.occupied_by.can_combine(item):
            return self.occupied_by.combine(item)

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

    def __repr__(self):
        return f"{self.dispensing.name}Dispenser"


class Trash(Counter):
    def pick_up(self, on_hands: bool = True):
        pass

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

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


class Stove(Counter):
    def can_drop_off(self, item: Item) -> bool:
        if self.occupied_by is None:
            return isinstance(item, CookingEquipment) and item.name in ["Pot", "Pan"]
        else:
            return self.occupied_by.can_combine(item)

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