DEVS-vs-ABMFrameworks/WolfSheep/PyPDEVS/Environment.py

import random
from pypdevs.simulator import AtomicDEVS

class Phase:
    IDLE = 0
    RESOLVING = 1


PHASES = ["IDLE", "RESOLVING"]


class EnvironmentState:
    def __init__(self, dims: tuple[int, int], regrowth_time: int, rng: random.Random):
        self.phase = Phase.IDLE
        self.width, self.height = dims
        self.regrowth_time = regrowth_time

        total_cells = self.width * self.height
        self.growing_grass = [
            rng.randint(1, regrowth_time) if rng.random() >= 0.5 else 0
            for _ in range(total_cells)
        ]

        self.sheep_grid = {}
        self.animal_positions = {}

        self.dead_animals = []
        self.successful_eaters = set()

    def __repr__(self):
        # alive_sheep = sum(len(cell) for row in self.sheep_grid for cell in row)
        # alive_wolves = sum(len(cell) for row in self.wolf_grid for cell in row)
        return f"[Environment | Phase: {PHASES[self.phase]}]"


class Environment(AtomicDEVS):
    def __init__(self, dims: tuple[int, int], regrowth_time: int, rng: random.Random):
        super().__init__("Environment")
        self.rng = rng
        self.width, self.height = dims
        self.state = EnvironmentState(dims, regrowth_time, rng)

        self.animal_born = self.addInPort("Animal Born")
        self.animal_died = self.addInPort("Animal Died")
        self.sheep_move = self.addInPort("Sheep Move")
        self.wolf_move = self.addInPort("Wolf Move")

        self.out_broadcast = self.addOutPort("Animal Broadcast")

    def extTransition(self, inputs):
        if self.animal_born in inputs:
            for species, aid, pos in inputs[self.animal_born]:
                self._register_animal(aid, species, pos)

        if self.animal_died in inputs:
            for species, aid, pos in inputs[self.animal_died]:
                self._deregister_animal(aid, species, pos)

        sheep_moves = inputs.get(self.sheep_move, [])
        wolf_moves = inputs.get(self.wolf_move, [])

        if sheep_moves or wolf_moves:
            self._process_sheep_turn(sheep_moves)
            self._process_wolf_turn(wolf_moves)
            self._grow_grass()
            self.state.phase = Phase.RESOLVING

        return self.state

    def intTransition(self):
        if self.state.phase == Phase.RESOLVING:
            self.state.successful_eaters = set()
            self.state.phase = Phase.IDLE
        else:
            self._grow_grass()
        return self.state

    def timeAdvance(self):
        return 0.0 if self.state.phase == Phase.RESOLVING else 1.0

    def outputFnc(self):
        out_bag = {}
        if self.state.phase == Phase.RESOLVING:
            out_bag[self.out_broadcast] = [self.state.successful_eaters]
        return out_bag

    def modelTransition(self, state):
        if self.state.dead_animals:
            if "Dead" not in state:
                state["Dead"] = []
            state["Dead"].extend(self.state.dead_animals)

            self.state.dead_animals.clear()
            return True
        return False

    def _register_animal(self, aid: int, species: int, pos: int):
        self.state.animal_positions[aid] = pos
        if species == 0: # SHEEP
            if pos not in self.state.sheep_grid:
                self.state.sheep_grid[pos] = set()
            self.state.sheep_grid[pos].add(aid)

    def _deregister_animal(self, aid: int, species: int, pos: int):
        if aid in self.state.animal_positions:
            del self.state.animal_positions[aid]
        if species == 0 and pos in self.state.sheep_grid: # SHEEP
            self.state.sheep_grid[pos].discard(aid)
            if not self.state.sheep_grid[pos]:
                del self.state.sheep_grid[pos]

    def _move_animal_on_grid(self, aid: int, species: int, new_pos: int):
        old_pos = self.state.animal_positions.get(aid)
        if old_pos == new_pos:
            return

        if old_pos and species == 0: # SHEEP
            self.state.sheep_grid[old_pos].discard(aid)
            if not self.state.sheep_grid[old_pos]:
                del self.state.sheep_grid[old_pos]

        self.state.animal_positions[aid] = new_pos
        if species == 0: # SHEEP
            if new_pos not in self.state.sheep_grid:
                self.state.sheep_grid[new_pos] = set()
            self.state.sheep_grid[new_pos].add(aid)

    def _process_sheep_turn(self, sheep_moves: list):
        self.rng.shuffle(sheep_moves)
        for aid, new_pos in sheep_moves:
            if aid not in self.state.animal_positions:
                continue

            self._move_animal_on_grid(aid, 0, new_pos)
            if self.state.growing_grass[new_pos] == 0:
                self.state.growing_grass[new_pos] = self.state.regrowth_time
                self.state.successful_eaters.add(aid)

    def _process_wolf_turn(self, wolf_moves: list):
        self.rng.shuffle(wolf_moves)
        for aid, new_pos in wolf_moves:
            if aid not in self.state.animal_positions:
                continue

            self._move_animal_on_grid(aid, 1, new_pos)
            sheep_at_pos = self.state.sheep_grid.get(new_pos)
            if sheep_at_pos:
                victim_id = self.rng.choice(sorted(list(sheep_at_pos)))
                self._deregister_animal(victim_id, 0, new_pos)
                self.state.dead_animals.append(victim_id)
                self.state.successful_eaters.add(aid)

    def _grow_grass(self):
        for pos in range(len(self.state.growing_grass)):
            if self.state.growing_grass[pos] > 0:
                self.state.growing_grass[pos] -= 1