1#!/usr/bin/env python3
  2"""
  3Cellular Automaton Framework
  4============================
  5
  6This module provides the base cellular automaton class and the
  7Predator-Prey (PP) implementation with Numba-accelerated kernels.
  8
  9Classes
 10-------
 11CA: Abstract base class for spatial cellular automata.
 12
 13PP: Predator-Prey model with configurable hunting behavior.
 14
 15Example
 16-------
 17```python
 18from models.CA import PP
 19
 20# Basic usage
 21model = PP(rows=100, cols=100, densities=(0.3, 0.15), seed=42)
 22model.run(steps=1000)
 23
 24# With evolution enabled
 25model = PP(rows=100, cols=100, seed=42)
 26model.evolve("prey_death", sd=0.05, min_val=0.01, max_val=0.15)
 27model.run(steps=500)
 28
 29# With directed hunting
 30model = PP(rows=100, cols=100, directed_hunting=True, seed=42)
 31```
 32"""
 33
 34from typing import Tuple, Dict, Optional
 35
 36import numpy as np
 37import logging
 38import sys
 39from pathlib import Path
 40
 41# Add parent directory to path for imports
 42sys.path.insert(0, str(Path(__file__).parent.parent))
 43
 44from models.numba_optimized import PPKernel, set_numba_seed
 45
 46# Module logger
 47logger = logging.getLogger(__name__)
 48
 49
 50class CA:
 51    """
 52    Base cellular automaton class for spatial simulations.
 53
 54    This class provides a framework for multi-species cellular automata with
 55    support for global parameters, per-cell evolving parameters, and
 56    grid initialization based on density.
 57
 58    Attributes
 59    ----------
 60    grid : np.ndarray
 61        2D numpy array containing integers in range [0, n_species].
 62    params : Dict[str, Any]
 63        Global parameters shared by all cells.
 64    cell_params : Dict[str, Any]
 65        Local per-cell parameters, typically stored as numpy arrays matching the grid shape.
 66    neighborhood : str
 67        The adjacency rule used ('neumann' or 'moore').
 68    generator : np.random.Generator
 69        The random number generator instance for reproducibility.
 70    species_names : Tuple[str, ...]
 71        Human-readable names for each species state.
 72    """
 73
 74    # Default colormap spec (string or sequence); resolved in `visualize` at runtime
 75    _default_cmap = "viridis"
 76
 77    # Read-only accessors for size/densities (protected attributes set in __init__)
 78    @property
 79    def rows(self) -> int:
 80        """int: Number of rows in the grid."""
 81        return getattr(self, "_rows")
 82
 83    @property
 84    def cols(self) -> int:
 85        """int: Number of columns in the grid."""
 86        return getattr(self, "_cols")
 87
 88    @property
 89    def densities(self) -> Tuple[float, ...]:
 90        """Tuple[float, ...]: Initial density fraction for each species."""
 91        return tuple(getattr(self, "_densities"))
 92
 93    # make n_species protected with read-only property
 94    @property
 95    def n_species(self) -> int:
 96        """int: Number of distinct species states (excluding empty state 0)."""
 97        return int(getattr(self, "_n_species"))
 98
 99    def __init__(
100        self,
101        rows: int,
102        cols: int,
103        densities: Tuple[float, ...],
104        neighborhood: str,
105        params: Dict[str, object],
106        cell_params: Dict[str, object],
107        seed: Optional[int] = None,
108    ) -> None:
109        """
110        Initialize the cellular automaton grid and configurations.
111
112        Parameters
113        ----------
114        rows : int
115            Number of rows in the grid (must be > 0).
116        cols : int
117            Number of columns in the grid (must be > 0).
118        densities : Tuple[float, ...]
119            Initial density for each species. Length defines `n_species`.
120            Values must sum to <= 1.0.
121        neighborhood : {'neumann', 'moore'}
122            Type of neighborhood connectivity.
123        params : Dict[str, Any]
124            Initial global parameter values.
125        cell_params : Dict[str, Any]
126            Initial local per-cell parameters.
127        seed : int, optional
128            Seed for the random number generator.
129        """
130        assert isinstance(rows, int) and rows > 0, "rows must be positive int"
131        assert isinstance(cols, int) and cols > 0, "cols must be positive int"
132        assert (
133            isinstance(densities, tuple) and len(densities) > 0
134        ), "densities must be a non-empty tuple"
135        for d in densities:
136            assert (
137                isinstance(d, (float, int)) and d >= 0
138            ), "each density must be non-negative"
139        total_density = float(sum(densities))
140        assert total_density <= 1.0 + 1e-12, "sum of densities must not exceed 1"
141        assert neighborhood in (
142            "neumann",
143            "moore",
144        ), "neighborhood must be 'neumann' or 'moore'"
145
146        self._n_species: int = len(densities)
147        # store protected size/density attributes (read-only properties exposed)
148        self._rows: int = rows
149        self._cols: int = cols
150        self._densities: Tuple[float, ...] = tuple(densities)
151        self.params: Dict[str, object] = dict(params) if params is not None else {}
152        self.cell_params: Dict[str, object] = (
153            dict(cell_params) if cell_params is not None else {}
154        )
155
156        # per-parameter evolve metadata and evolution state
157        # maps parameter name -> dict with keys 'sd','min','max','species'
158        self._evolve_info: Dict[str, Dict[str, float]] = {}
159        # when True, inheritance uses deterministic copy from parent (no mutation)
160        self._evolution_stopped: bool = False
161
162        # human-readable species names (useful for visualization). Default
163        # generates generic names based on n_species; subclasses may override.
164        self.species_names: Tuple[str, ...] = tuple(
165            f"species{i+1}" for i in range(self._n_species)
166        )
167        self.neighborhood: str = neighborhood
168        self.generator: np.random.Generator = np.random.default_rng(seed)
169
170        self.grid: np.ndarray = np.zeros((rows, cols), dtype=int)
171
172        total_cells = rows * cols
173        # Fill grid with species states 1..n_species according to densities.
174        for i, dens in enumerate(densities):
175            if dens <= 0:
176                continue
177            n_to_fill = int(round(total_cells * float(dens)))
178            if n_to_fill <= 0:
179                continue
180            empty_flat = np.flatnonzero(self.grid.ravel() == 0)
181            if len(empty_flat) == 0:
182                break
183            n_choice = min(n_to_fill, len(empty_flat))
184            chosen = self.generator.choice(empty_flat, size=n_choice, replace=False)
185            # assign chosen flattened indices to state i+1
186            r = chosen // cols
187            c = chosen % cols
188            self.grid[r, c] = i + 1
189
190    def validate(self) -> None:
191        """
192        Validate core CA invariants and grid dimensions.
193
194        Checks that the neighborhood is valid, the grid matches initialized dimensions,
195        and that local parameter arrays match the grid shape.
196
197        Raises
198        ------
199        ValueError
200            If any structural invariant is violated.
201        """
202        if self.neighborhood not in ("neumann", "moore"):
203            raise ValueError("neighborhood must be 'neumann' or 'moore'")
204
205        expected_shape = (int(getattr(self, "_rows")), int(getattr(self, "_cols")))
206        if self.grid.shape != expected_shape:
207            raise ValueError(
208                f"grid shape {self.grid.shape} does not match expected {expected_shape}"
209            )
210
211        # Ensure any array in cell_params matches grid shape
212        for k, v in (self.cell_params or {}).items():
213            if isinstance(v, np.ndarray) and v.shape != expected_shape:
214                raise ValueError(f"cell_params['{k}'] must have shape equal to grid")
215
216    def _infer_species_from_param_name(self, param_name: str) -> Optional[int]:
217        """
218        Infer the 1-based species index from a parameter name using `species_names`.
219
220        This method checks if the given parameter name starts with any of the
221        defined species names followed by an underscore (e.g., 'prey_birth').
222        It is used to automatically route global parameters to the correct
223        species' local parameter arrays.
224
225        Parameters
226        ----------
227        param_name : str
228            The name of the parameter to check.
229
230        Returns
231        -------
232        Optional[int]
233            The 1-based index of the species if a matching prefix is found;
234            otherwise, None.
235
236        Notes
237        -----
238        The method expects `self.species_names` to be a collection of strings.
239        If `param_name` is not a string or no match is found, it returns None.
240        """
241        if not isinstance(param_name, str):
242            return None
243        for idx, name in enumerate(self.species_names or ()):  # type: ignore
244            if isinstance(name, str) and param_name.startswith(f"{name}_"):
245                return idx + 1
246        return None
247
248    def evolve(
249        self,
250        param: str,
251        species: Optional[int] = None,
252        sd: float = 0.05,
253        min_val: Optional[float] = None,
254        max_val: Optional[float] = None,
255    ) -> None:
256        """
257        Enable per-cell evolution for a specific parameter on a given species.
258
259        This method initializes a spatial parameter array (local parameter map)
260        for a global parameter. It allows individual cells to carry their own
261        values for that parameter, which can then mutate and evolve during
262        the simulation.
263
264        Parameters
265        ----------
266        param : str
267            The name of the global parameter to enable for evolution.
268            Must exist in `self.params`.
269        species : int, optional
270            The 1-based index of the species to which this parameter applies.
271            If None, the method attempts to infer the species from the
272            parameter name prefix.
273        sd : float, default 0.05
274            The standard deviation of the Gaussian mutation applied during
275            inheritance/reproduction.
276        min_val : float, optional
277            The minimum allowable value for the parameter (clamping).
278            Defaults to 0.01 if not provided.
279        max_val : float, optional
280            The maximum allowable value for the parameter (clamping).
281            Defaults to 0.99 if not provided.
282
283        Raises
284        ------
285        ValueError
286            If the parameter is not in `self.params`, the species cannot be
287            inferred, or the species index is out of bounds.
288
289        Notes
290        -----
291        The local parameter is stored in `self.cell_params` as a 2D numpy
292        array initialized with the current global value for all cells of
293        the target species, and `NaN` elsewhere.
294        """
295        if min_val is None:
296            min_val = 0.01
297        if max_val is None:
298            max_val = 0.99
299        if param not in self.params:
300            raise ValueError(f"Unknown parameter '{param}'")
301        if species is None:
302            species = self._infer_species_from_param_name(param)
303            if species is None:
304                raise ValueError(
305                    "species must be provided or inferable from param name and species_names"
306                )
307        if not isinstance(species, int) or species <= 0 or species > self._n_species:
308            raise ValueError("species must be an integer between 1 and n_species")
309
310        arr = np.full(self.grid.shape, np.nan, dtype=float)
311        mask = self.grid == int(species)
312        arr[mask] = float(self.params[param])
313        self.cell_params[param] = arr
314        self._evolve_info[param] = {
315            "sd": float(sd),
316            "min": float(min_val),
317            "max": float(max_val),
318            "species": int(species),
319        }
320
321    def update(self) -> None:
322        """
323        Perform one update step of the cellular automaton.
324
325        This is an abstract method that defines the transition rules of the
326        system. It must be implemented by concrete subclasses to specify
327        how cell states and parameters change over time based on their
328        current state and neighborhood.
329
330        Raises
331        ------
332        NotImplementedError
333            If called directly on the base class instead of an implementation.
334
335        Returns
336        -------
337        None
338
339        Notes
340        -----
341        In a typical implementation, this method handles the logic for
342        stochastic transitions, movement, or predator-prey interactions.
343        """
344        raise NotImplementedError(
345            "Override update() in a subclass to define CA dynamics"
346        )
347
348    def run(
349        self,
350        steps: int,
351        stop_evolution_at: Optional[int] = None,
352        snapshot_iters: Optional[list] = None,
353    ) -> None:
354        """
355        Execute the cellular automaton simulation for a specified number of steps.
356
357        This method drives the simulation loop, calling `update()` at each
358        iteration. It manages visualization updates, directory creation for
359        data persistence, and handles the freezing of evolving parameters
360        at a specific time step.
361
362        Parameters
363        ----------
364        steps : int
365            The total number of iterations to run (must be non-negative).
366        stop_evolution_at : int, optional
367            The 1-based iteration index after which parameter mutation is
368            disabled. Useful for observing system stability after a period
369            of adaptation.
370        snapshot_iters : List[int], optional
371            A list of specific 1-based iteration indices at which to save
372            the current grid state to the results directory.
373
374        Returns
375        -------
376        None
377
378        Notes
379        -----
380        If snapshots are requested, a results directory is automatically created
381        using a timestamped subfolder (e.g., 'results/run-1700000000/').
382        Visualization errors are logged but do not terminate the simulation.
383        """
384        assert (
385            isinstance(steps, int) and steps >= 0
386        ), "steps must be a non-negative integer"
387
388        # normalize snapshot iteration list
389        snapshot_set = set(snapshot_iters) if snapshot_iters is not None else set()
390
391        for i in range(steps):
392            self.update()
393            # Update visualization if enabled every `interval` iterations
394            if getattr(self, "_viz_on", False):
395                # iteration number is 1-based for display
396                try:
397                    self._viz_update(i + 1)
398                except Exception:
399                    # Log visualization errors but don't stop the simulation
400                    logger.exception(
401                        "Visualization update failed at iteration %d", i + 1
402                    )
403
404            # create snapshots if requested at this iteration
405            if (i + 1) in snapshot_set:
406                try:
407                    # create snapshot folder if not present
408                    if (
409                        not hasattr(self, "_viz_snapshot_dir")
410                        or self._viz_snapshot_dir is None
411                    ):
412                        import os, time
413
414                        base = "results"
415                        ts = int(time.time())
416                        run_folder = f"run-{ts}"
417                        full = os.path.join(base, run_folder)
418                        os.makedirs(full, exist_ok=True)
419                        self._viz_snapshot_dir = full
420                    self._viz_save_snapshot(i + 1)
421                except (OSError, PermissionError):
422                    logger.exception(
423                        "Failed to create or write snapshot at iteration %d", i + 1
424                    )
425
426            # stop evolution at specified time-step (disable further evolution)
427            if stop_evolution_at is not None and (i + 1) == int(stop_evolution_at):
428                # mark evolution as stopped; do not erase evolve metadata so
429                # deterministic inheritance can still use parent values
430                self._evolution_stopped = True
431
432
433class PP(CA):
434    """
435    Predator-Prey Cellular Automaton model with Numba-accelerated kernels.
436
437    This model simulates a stochastic predator-prey system where species
438    interact on a 2D grid. It supports evolving per-cell death rates,
439    periodic boundary conditions, and both random and directed hunting
440    behaviors.
441
442    Parameters
443    ----------
444    rows : int, default 10
445        Number of rows in the simulation grid.
446    cols : int, default 10
447        Number of columns in the simulation grid.
448    densities : Tuple[float, ...], default (0.2, 0.1)
449        Initial population densities for (prey, predator).
450    neighborhood : {'moore', 'neumann'}, default 'moore'
451        The neighborhood type for cell interactions.
452    params : Dict[str, object], optional
453        Global parameters: "prey_death", "predator_death", "prey_birth",
454        "predator_birth".
455    cell_params : Dict[str, object], optional
456        Initial local parameter maps (2D arrays).
457    seed : int, optional
458        Random seed for reproducibility.
459    synchronous : bool, default True
460        If True, updates the entire grid at once. If False, updates
461        cells asynchronously.
462    directed_hunting : bool, default False
463        If True, predators selectively hunt prey rather than choosing
464        neighbors at random.
465
466    Attributes
467    ----------
468    species_names : Tuple[str, ...]
469        Labels for the species ('prey', 'predator').
470    synchronous : bool
471        Current update mode.
472    directed_hunting : bool
473        Current hunting strategy logic.
474    """
475
476    # Default colors: 0=empty black, 1=prey green, 2=predator red
477    _default_cmap = ("black", "green", "red")
478
479    def __init__(
480        self,
481        rows: int = 10,
482        cols: int = 10,
483        densities: Tuple[float, ...] = (0.2, 0.1),
484        neighborhood: str = "moore",
485        params: Dict[str, object] = None,
486        cell_params: Dict[str, object] = None,
487        seed: Optional[int] = None,
488        synchronous: bool = True,
489        directed_hunting: bool = False,  # New directed hunting option
490    ) -> None:
491        """
492        Initialize the Predator-Prey CA with validated parameters and kernels.
493        """
494        # Allowed params and defaults
495        _defaults = {
496            "prey_death": 0.05,
497            "predator_death": 0.1,
498            "prey_birth": 0.25,
499            "predator_birth": 0.2,
500        }
501
502        # Validate user-supplied params: only allowed keys
503        if params is None:
504            merged_params = dict(_defaults)
505        else:
506            if not isinstance(params, dict):
507                raise TypeError("params must be a dict or None")
508            extra = set(params.keys()) - set(_defaults.keys())
509            if extra:
510                raise ValueError(f"Unexpected parameter keys: {sorted(list(extra))}")
511            # Do not override user-set values: start from defaults then update with user values
512            merged_params = dict(_defaults)
513            merged_params.update(params)
514
515        # Validate numerical ranges
516        for k, v in merged_params.items():
517            if not isinstance(v, (int, float)):
518                raise TypeError(f"Parameter '{k}' must be a number between 0 and 1")
519            if not (0.0 <= float(v) <= 1.0):
520                raise ValueError(f"Parameter '{k}' must be between 0 and 1")
521
522        # Call base initializer with merged params
523        super().__init__(
524            rows, cols, densities, neighborhood, merged_params, cell_params, seed
525        )
526
527        self.synchronous: bool = bool(synchronous)
528        self.directed_hunting: bool = bool(directed_hunting)
529
530        # set human-friendly species names for PP
531        self.species_names = ("prey", "predator")
532
533        if seed is not None:
534            # This sets the seed for all @njit functions globally
535            set_numba_seed(seed)
536
537        self._kernel = PPKernel(
538            rows, cols, neighborhood, directed_hunting=directed_hunting
539        )
540
541    # Remove PP-specific evolve wrapper; use CA.evolve with optional species
542
543    def validate(self) -> None:
544        """
545        Validate Predator-Prey specific invariants and spatial parameter arrays.
546
547        Extends the base CA validation to ensure that numerical parameters are
548        within the [0, 1] probability range and that evolved parameter maps
549        (e.g., prey_death) correctly align with the species locations.
550
551        Raises
552        ------
553        ValueError
554            If grid shapes, parameter ranges, or species masks are inconsistent.
555        TypeError
556            If parameters are non-numeric.
557        """
558        super().validate()
559
560        # Validate global params
561        for k, v in (self.params or {}).items():
562            if not isinstance(v, (int, float)):
563                raise TypeError(f"Parameter '{k}' must be numeric")
564            if not (0.0 <= float(v) <= 1.0):
565                raise ValueError(f"Parameter '{k}' must be between 0 and 1")
566
567        # Validate per-cell evolve arrays
568        for pname, meta in (self._evolve_info or {}).items():
569            arr = self.cell_params.get(pname)
570            if not isinstance(arr, np.ndarray):
571                # absent or non-array per-cell params are allowed; skip
572                continue
573            # shape already checked in super().validate(), but be explicit
574            if arr.shape != self.grid.shape:
575                raise ValueError(f"cell_params['{pname}'] must match grid shape")
576            # expected non-NaN positions correspond to species stored in metadata
577            species = None
578            if isinstance(meta, dict) and "species" in meta:
579                species = int(meta.get("species"))
580            else:
581                # try to infer species from parameter name using species_names
582                species = self._infer_species_from_param_name(pname)
583                if species is None:
584                    raise ValueError(
585                        f"cell_params['{pname}'] missing species metadata and could not infer from name"
586                    )
587            nonnan = ~np.isnan(arr)
588            expected = self.grid == species
589            if not np.array_equal(nonnan, expected):
590                raise ValueError(
591                    f"cell_params['{pname}'] non-NaN entries must match positions of species {species}"
592                )
593            # values must be within configured range where not NaN
594            mn = float(meta.get("min", 0.0))
595            mx = float(meta.get("max", 1.0))
596            vals = arr[~np.isnan(arr)]
597            if vals.size > 0:
598                if np.any(vals < mn) or np.any(vals > mx):
599                    raise ValueError(
600                        f"cell_params['{pname}'] contains values outside [{mn}, {mx}]"
601                    )
602
603    def update_async(self) -> None:
604        """
605        Execute an asynchronous update using the optimized Numba kernel.
606
607        This method retrieves the evolved parameter maps and delegates the
608        stochastic transitions to the `PPKernel`. Asynchronous updates
609        typically handle cell-by-cell logic where changes can be
610        immediately visible to neighbors.
611        """
612        # Get the evolved prey death map
613        # Fallback to a full array of the global param if it doesn't exist yet
614        p_death_arr = self.cell_params.get("prey_death")
615        if p_death_arr is None:
616            p_death_arr = np.full(
617                self.grid.shape, self.params["prey_death"], dtype=np.float64
618            )
619
620        meta = self._evolve_info.get(
621            "prey_death", {"sd": 0.05, "min": 0.001, "max": 0.1}
622        )
623
624        # Call the optimized kernel (uses pre-allocated buffers)
625        self._kernel.update(
626            self.grid,
627            p_death_arr,
628            float(self.params["prey_birth"]),
629            float(self.params["prey_death"]),
630            float(self.params["predator_birth"]),
631            float(self.params["predator_death"]),
632            float(meta["sd"]),
633            float(meta["min"]),
634            float(meta["max"]),
635            self._evolution_stopped,
636        )
637
638    def update(self) -> None:
639        """
640        Dispatch the simulation step based on the configured update mode.
641        """
642        self.update_async()