"""
Data models for the Vector Borne Disease Control Environment.

City is represented as a graph of N zones (25-75), randomly generated each episode.

Action space:
    action_type 0 = spray  (zone_id)
    action_type 1 = trap   (zone_id)
    action_type 2 = wait

Zone features (per zone):
    infestation_rate — dynamic, continuous (0.0=clean, 1.0=fully infested)
    population       — fixed at reset (randomly generated per zone)
    rainfall         — dynamic, drifts per zone each step (0.0-1.0)
    humidity         — dynamic, drifts per zone each step (0.0-1.0)
    rained           — dynamic, flips probabilistically each step (0=no, 1=yes)
    temperature      — dynamic, drifts per zone each step (15.0-45.0 Celsius)
"""

from enum import IntEnum
from typing import Dict, List, Optional

from openenv.core.env_server.types import Action, Observation
from pydantic import BaseModel, Field

# agent can either spray, trap in a zone or wait, for a step.
class ActionType(IntEnum):
    SPRAY = 0
    TRAP = 1
    WAIT = 2


class ZoneState(BaseModel):
    """State of a single zone in the city graph."""

    infestation_rate: float = Field(..., description="Mosquito infestation rate (0.0=clean, 1.0=fully infested)")
    population: int = Field(..., description="Number of people in this zone; fixed for the episode")
    humidity: float = Field(..., description="Humidity level (0.0-1.0); drifts each step")
    rained: int = Field(..., description="Whether it rained this step (0=no, 1=yes)")
    temperature: float = Field(..., description="Temperature in Celsius (15.0-45.0); drifts each step")
    has_trap: bool = Field(default=False, description="Whether a trap is deployed in this zone")
    trap_steps_remaining: int = Field(default=0, description="Steps until the trap stops working")
    is_treated: bool = Field(default=False, description="Whether this zone has been sprayed and treatment is active")
    treatment_steps_remaining: int = Field(default=0, description="Steps until spray treatment wears off")


class VectorBorneDiseaseControlAction(Action):
    """Action for the Vector Borne Disease Control environment."""

    # One spray or trap resource covers the entire target zone.
    action_type: ActionType = Field(..., description="0=spray, 1=trap, 2=wait")
    zone_id: Optional[int] = Field(default=None, description="Target zone ID; ignored for wait")


class VectorBorneDiseaseControlObservation(Observation):
    """Observation from the Vector Borne Disease Control environment."""

    zones: Dict[int, ZoneState] = Field(
        default_factory=dict,
        description="Per-zone state keyed by zone ID",
    )
    adjacency: Dict[int, List[int]] = Field(
        default_factory=dict,
        description="Adjacency list — zone ID to list of neighboring zone IDs; fixed for the episode",
    )
    remaining_spray: int = Field(default=0, description="Spray resources remaining")
    remaining_traps: int = Field(default=0, description="Trap resources remaining")
    step: int = Field(default=0, description="Current step number")
    cumulative_avg_infestation_rate: float = Field(default=0.0, description="Average infestation_rate across all zones")