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


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
    ):
        """
        Initialize the grid with terrain and initial fires.
        
        Args:
            seed: Random seed for reproducibility
            fire_count: Number of initial fire points (1-25)
            fire_intensity: Initial fire intensity (0.0-1.0)
            building_count: Number of buildings to place (1-25)
        """
        if seed is not None:
            random.seed(seed)
            self.seed = seed
        
        # Clamp values to valid ranges
        fire_count = max(1, min(25, fire_count))
        fire_intensity = max(0.0, min(1.0, fire_intensity))
        building_count = max(1, min(25, 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,
        }