models.py

"""
Fire-Rescue - Data Models

Defines core data structures for the fire rescue simulation.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
import random

from config import SCENARIO_DEFAULTS

FIRE_COUNT_DEFAULTS = SCENARIO_DEFAULTS["fire_count"]
BUILDING_COUNT_DEFAULTS = SCENARIO_DEFAULTS["building_count"]


class UnitType(str, Enum):
    """Types of firefighting units available."""
    FIRE_TRUCK = "fire_truck"
    HELICOPTER = "helicopter"


class CellType(str, Enum):
    """Types of terrain cells in the grid."""
    EMPTY = "empty"
    BUILDING = "building"
    FOREST = "forest"


class SimulationStatus(str, Enum):
    """Status of the simulation."""
    IDLE = "idle"
    RUNNING = "running"
    SUCCESS = "success"
    FAIL = "fail"


class FireLevel(str, Enum):
    """Initial fire intensity levels."""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"


@dataclass
class Fire:
    """Represents a fire cell in the grid."""
    x: int
    y: int
    intensity: float  # 0.0 to 1.0
    
    def to_dict(self) -> dict:
        return {
            "x": self.x,
            "y": self.y,
            "intensity": round(self.intensity, 2)
        }


@dataclass
class Unit:
    """Represents a firefighting unit."""
    id: str
    unit_type: UnitType
    owner: str  # "player" or "ai"
    x: int
    y: int
    cooldown: int = 0  # Ticks until next action
    
    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "type": self.unit_type.value,
            "owner": self.owner,
            "x": self.x,
            "y": self.y
        }


@dataclass
class Event:
    """Represents a simulation event."""
    tick: int
    event_type: str
    details: dict = field(default_factory=dict)
    
    def to_dict(self) -> dict:
        return {
            "tick": self.tick,
            "type": self.event_type,
            **self.details
        }


@dataclass
class Cell:
    """Represents a cell in the grid."""
    x: int
    y: int
    cell_type: CellType
    fire_intensity: float = 0.0  # 0.0 = no fire, 1.0 = max fire
    damage: float = 0.0  # Accumulated damage (0.0 to 1.0)
    
    def is_on_fire(self) -> bool:
        return self.fire_intensity > 0.0
    
    def is_destroyed(self) -> bool:
        return self.damage >= 1.0


@dataclass
class WorldState:
    """
    Represents the complete state of the simulation world.
    Uses a 2D grid system.
    """
    width: int
    height: int
    tick: int = 0
    status: SimulationStatus = SimulationStatus.IDLE
    
    # Grid cells
    grid: list[list[Cell]] = field(default_factory=list)
    
    # Units on the field
    units: list[Unit] = field(default_factory=list)
    
    # Recent events for logging
    recent_events: list[Event] = field(default_factory=list)
    
    # Global metrics
    building_integrity: float = 1.0  # Average building health (0.0 to 1.0)
    forest_burn_ratio: float = 0.0   # Percentage of forest burned (0.0 to 1.0)
    
    # Configuration
    max_ticks: int = 200
    max_units: int = 10
    seed: Optional[int] = None
    
    # Building positions (for dynamic placement)
    building_positions: list[tuple[int, int]] = field(default_factory=list)
    
    # Unit ID counter
    _unit_counter: int = 0
    
    def initialize_grid(
        self, 
        seed: Optional[int] = None, 
        fire_count: int = 4,
        fire_intensity: float = 0.6,
        building_count: int = 16
    ):
        f"""
        Initialize the grid with terrain and initial fires.
        
        Args:
            seed: Random seed for reproducibility
            fire_count: Number of initial fire points ({int(FIRE_COUNT_DEFAULTS.minimum)}-{int(FIRE_COUNT_DEFAULTS.maximum)})
            fire_intensity: Initial fire intensity (0.0-1.0)
            building_count: Number of buildings to place ({int(BUILDING_COUNT_DEFAULTS.minimum)}-{int(BUILDING_COUNT_DEFAULTS.maximum)})
        """
        if seed is not None:
            random.seed(seed)
            self.seed = seed
        
        # Clamp values using shared configuration
        fire_count = int(FIRE_COUNT_DEFAULTS.clamp(fire_count))
        fire_intensity = max(0.0, min(1.0, fire_intensity))
        building_count = int(BUILDING_COUNT_DEFAULTS.clamp(building_count))
        
        # Generate random building positions (connected cluster)
        self.building_positions = self._generate_building_positions(building_count)
        building_set = set(self.building_positions)
        
        self.grid = []
        for y in range(self.height):
            row = []
            for x in range(self.width):
                # Default to forest
                cell_type = CellType.FOREST
                
                # Place buildings at generated positions
                if (x, y) in building_set:
                    cell_type = CellType.BUILDING
                
                row.append(Cell(x=x, y=y, cell_type=cell_type))
            self.grid.append(row)
        
        # Place initial fires
        fires_placed = 0
        attempts = 0
        max_attempts = 100
        
        while fires_placed < fire_count and attempts < max_attempts:
            x = random.randint(0, self.width - 1)
            y = random.randint(0, self.height - 1)
            
            # Only place fire on forest initially (not buildings)
            cell = self.grid[y][x]
            if cell.cell_type == CellType.FOREST and cell.fire_intensity == 0:
                cell.fire_intensity = fire_intensity
                fires_placed += 1
            
            attempts += 1
    
    def _generate_building_positions(self, count: int) -> list[tuple[int, int]]:
        """
        Generate random building positions ensuring connectivity.
        
        Buildings grow as a connected cluster from a random starting point.
        At least 2 buildings will be adjacent (if count >= 2).
        """
        if count <= 0:
            return []
        
        positions = []
        
        # Start with a random position (avoid edges for better growth)
        start_x = random.randint(2, self.width - 3)
        start_y = random.randint(2, self.height - 3)
        positions.append((start_x, start_y))
        
        if count == 1:
            return positions
        
        # Grow the cluster by adding adjacent cells
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 4-directional adjacency
        
        while len(positions) < count:
            # Get all possible adjacent positions to existing buildings
            candidates = set()
            for (px, py) in positions:
                for dx, dy in directions:
                    nx, ny = px + dx, py + dy
                    # Check bounds and not already a building
                    if 0 <= nx < self.width and 0 <= ny < self.height:
                        if (nx, ny) not in positions:
                            candidates.add((nx, ny))
            
            if not candidates:
                # No more valid positions (unlikely but handle it)
                break
            
            # Randomly select one candidate
            new_pos = random.choice(list(candidates))
            positions.append(new_pos)
        
        return positions
    
    def get_cell(self, x: int, y: int) -> Optional[Cell]:
        """Get cell at position, returns None if out of bounds."""
        if 0 <= x < self.width and 0 <= y < self.height:
            return self.grid[y][x]
        return None
    
    def get_fires(self) -> list[Fire]:
        """Get list of all active fires."""
        fires = []
        for row in self.grid:
            for cell in row:
                if cell.is_on_fire():
                    fires.append(Fire(x=cell.x, y=cell.y, intensity=cell.fire_intensity))
        return fires
    
    def generate_unit_id(self, unit_type: UnitType) -> str:
        """Generate a unique unit ID."""
        self._unit_counter += 1
        prefix = "truck" if unit_type == UnitType.FIRE_TRUCK else "heli"
        return f"{prefix}_{self._unit_counter}"
    
    def add_unit(self, unit_type: UnitType, x: int, y: int, source: str) -> Optional[Unit]:
        """Add a new unit to the world. Returns None if limit reached or position invalid."""
        if len(self.units) >= self.max_units:
            return None
        
        if not (0 <= x < self.width and 0 <= y < self.height):
            return None
        
        # Check cell conditions - cannot deploy on fire or buildings
        cell = self.get_cell(x, y)
        if cell is None:
            return None
        
        # Cannot deploy on burning cells
        if cell.fire_intensity > 0:
            return None
        
        # Cannot deploy on buildings
        if cell.cell_type == CellType.BUILDING:
            return None
        
        unit = Unit(
            id=self.generate_unit_id(unit_type),
            unit_type=unit_type,
            owner="player",
            x=x,
            y=y
        )
        self.units.append(unit)
        
        # Record event
        self.recent_events.append(Event(
            tick=self.tick,
            event_type="deploy_unit",
            details={
                "by": source,
                "unit_type": unit_type.value,
                "x": x,
                "y": y
            }
        ))
        
        # Keep only recent events (last 20)
        if len(self.recent_events) > 20:
            self.recent_events = self.recent_events[-20:]
        
        return unit
    
    def calculate_metrics(self):
        """Recalculate global metrics (building damage ratio)."""
        total_buildings = 0
        burning_buildings = 0
        
        for row in self.grid:
            for cell in row:
                if cell.cell_type == CellType.BUILDING:
                    total_buildings += 1
                    # Building is burning if it has fire on it
                    if cell.fire_intensity > 0:
                        burning_buildings += 1
        
        # Building integrity: ratio of non-burning buildings
        if total_buildings > 0:
            self.building_integrity = (total_buildings - burning_buildings) / total_buildings
        else:
            self.building_integrity = 1.0
        
        # Store total buildings for reference
        self._total_buildings = total_buildings
        self._burning_buildings = burning_buildings
        
        # Forest burn ratio is no longer used (replaced by active fires count)
        self.forest_burn_ratio = 0.0
    
    def to_dict(self) -> dict:
        """Serialize world state to dictionary."""
        return {
            "tick": self.tick,
            "status": self.status.value,
            "width": self.width,
            "height": self.height,
            "fires": [f.to_dict() for f in self.get_fires()],
            "units": [u.to_dict() for u in self.units],
            "building_integrity": round(self.building_integrity, 2),
            "forest_burn_ratio": round(self.forest_burn_ratio, 2),
            "recent_events": [e.to_dict() for e in self.recent_events[-5:]],
            "buildings": [{"x": x, "y": y} for x, y in self.building_positions],
            "max_units": self.max_units,
        }