from typing import Any, Optional, List
from dataclasses import dataclass
import jax
import jax.numpy as jnp

@dataclass(frozen=True)
class Entity:
    id: Any
    x: int
    y: int

@dataclass(frozen=True)
class Agent(Entity):
    max_load: int
    curr_load: int
    bag: jnp.ndarray
    beliefs: Any  # Beliefs can be managed externally in a JAX-friendly structure

    @staticmethod
    def create(id: Any, x: int, y: int, n_objectives: int, max_load: Optional[int] = 2) -> 'Agent':
        return Agent(
            id=id,
            x=x,
            y=y,
            max_load=max_load,
            curr_load=0,
            bag=jnp.zeros(n_objectives, dtype=jnp.int32),
            beliefs=None  # Can initialize beliefs if needed in a JAX-compatible form
        )

    def add_to_bag(self, index: int, amount: int) -> 'Agent':
        """Add items to the agent's bag."""
        updated_bag = self.bag.at[index].add(amount)
        return Agent(
            id=self.id,
            x=self.x,
            y=self.y,
            max_load=self.max_load,
            curr_load=self.curr_load + amount,
            bag=updated_bag,
            beliefs=self.beliefs
        )

@dataclass(frozen=True)
class Plate(Entity):
    pressed: bool
    pressing_agent: Optional[int]

    @staticmethod
    def create(id: Any, x: int, y: int) -> 'Plate':
        return Plate(id=id, x=x, y=y, pressed=False, pressing_agent=None)

    def press(self, agent_id: int) -> 'Plate':
        return Plate(id=self.id, x=self.x, y=self.y, pressed=True, pressing_agent=agent_id)

@dataclass(frozen=True)
class Door(Entity):
    open: bool

    @staticmethod
    def create(id: Any, x: int, y: int) -> 'Door':
        return Door(id=id, x=x, y=y, open=False)

    def toggle(self, state: bool) -> 'Door':
        return Door(id=self.id, x=self.x, y=self.y, open=state)

@dataclass(frozen=True)
class Wall(Entity):
    @staticmethod
    def create(id: Any, x: int, y: int) -> 'Wall':
        return Wall(id=id, x=x, y=y)

@dataclass(frozen=True)
class Goal(Entity):
    achieved: bool

    @staticmethod
    def create(id: Any, x: int, y: int) -> 'Goal':
        return Goal(id=id, x=x, y=y, achieved=False)

    def achieve(self) -> 'Goal':
        return Goal(id=self.id, x=self.x, y=self.y, achieved=True)
