schemas/input.py

"""Pydantic input models for FDAM AI Pipeline.

Uses Literal unions instead of Enums per project code style.
"""

from datetime import date
from typing import Literal, Optional

from pydantic import BaseModel, Field, field_validator, model_validator


# --- Type Definitions (Literal unions) ---

FacilityClassification = Literal["operational", "non-operational", "public-childcare"]
ConstructionEra = Literal["pre-1980", "1980-2000", "post-2000"]
ZoneType = Literal["burn", "near-field", "far-field"]
ConditionLevel = Literal["background", "light", "moderate", "heavy", "structural-damage"]

# Material categories
MaterialType = Literal[
    # Non-porous
    "steel",
    "concrete",
    "glass",
    "metal",
    "cmu",
    # Semi-porous
    "drywall-painted",
    "drywall-unpainted",
    "wood-sealed",
    "wood-unsealed",
    # Porous
    "carpet",
    "carpet-pad",
    "insulation-fiberglass",
    "insulation-other",
    "acoustic-tile",
    "upholstery",
    # HVAC
    "ductwork-rigid",
    "ductwork-flexible",
    "hvac-interior-insulation",
]

MaterialCategory = Literal["non-porous", "semi-porous", "porous", "hvac"]

Disposition = Literal["no-action", "clean", "evaluate", "remove", "remove-repair"]

OdorIntensity = Literal["none", "faint", "moderate", "strong"]
CharDensity = Literal["sparse", "moderate", "dense"]
SampleType = Literal["tape_lift", "surface_wipe", "both"]
Priority = Literal["high", "medium", "low"]


# --- Helper Functions ---

def get_material_category(material: MaterialType) -> MaterialCategory:
    """Get the category for a material type."""
    non_porous = {"steel", "concrete", "glass", "metal", "cmu"}
    semi_porous = {"drywall-painted", "drywall-unpainted", "wood-sealed", "wood-unsealed"}
    porous = {"carpet", "carpet-pad", "insulation-fiberglass", "insulation-other", "acoustic-tile", "upholstery"}
    hvac = {"ductwork-rigid", "ductwork-flexible", "hvac-interior-insulation"}

    if material in non_porous:
        return "non-porous"
    elif material in semi_porous:
        return "semi-porous"
    elif material in porous:
        return "porous"
    elif material in hvac:
        return "hvac"
    else:
        return "porous"  # Conservative default


# --- Project Level ---

class ProjectInfo(BaseModel):
    """Project-level information."""

    project_name: str = Field(..., min_length=1, description="Project or facility name")
    address: str = Field(..., min_length=1, description="Full street address")
    city: str = Field(..., min_length=1)
    state: str = Field(..., min_length=2, max_length=2)
    zip_code: str = Field(..., min_length=5)

    client_name: str = Field(..., min_length=1)
    client_contact: Optional[str] = None
    client_email: Optional[str] = None
    client_phone: Optional[str] = None

    fire_date: date = Field(..., description="Date of fire incident")
    assessment_date: date = Field(..., description="Date of assessment")

    facility_classification: FacilityClassification
    construction_era: ConstructionEra

    assessor_name: str = Field(..., min_length=1, description="Industrial hygienist name")
    assessor_credentials: Optional[str] = Field(None, description="CIH, CSP, etc.")


# --- Room/Area Level ---

class Dimensions(BaseModel):
    """Room dimensions for calculations."""

    length_ft: float = Field(..., gt=0, le=10000, description="Length in feet")
    width_ft: float = Field(..., gt=0, le=10000, description="Width in feet")
    ceiling_height_ft: float = Field(..., gt=0, le=500, description="Ceiling height in feet")

    @property
    def area_sf(self) -> float:
        """Calculate floor area in square feet."""
        return self.length_ft * self.width_ft

    @property
    def volume_cf(self) -> float:
        """Calculate volume in cubic feet."""
        return self.area_sf * self.ceiling_height_ft


class Surface(BaseModel):
    """Individual surface within a room."""

    id: str = Field(..., min_length=1, description="Unique surface identifier")
    material: MaterialType = Field(..., description="Material type")
    description: str = Field(..., min_length=1, description="e.g., 'North wall drywall'")
    area_sf: float = Field(..., gt=0, description="Surface area in square feet")

    zone: Optional[ZoneType] = Field(None, description="Can be set by AI or user")
    condition: Optional[ConditionLevel] = Field(None, description="Can be set by AI or user")
    disposition: Optional[Disposition] = Field(None, description="Calculated by system")

    ai_detected: bool = Field(False, description="Was this detected by AI from images?")
    confidence: Optional[float] = Field(None, ge=0, le=1, description="AI confidence score")

    @property
    def category(self) -> MaterialCategory:
        """Get the material category."""
        return get_material_category(self.material)


class Room(BaseModel):
    """Room or area within the building."""

    id: str = Field(..., min_length=1, description="Unique room identifier")
    name: str = Field(..., min_length=1, description="e.g., 'Warehouse Bay A'")
    floor: Optional[str] = Field(None, description="e.g., 'Ground Floor'")

    dimensions: Dimensions

    zone_classification: Optional[ZoneType] = Field(None, description="AI-determined or user override")
    zone_confidence: Optional[float] = Field(None, ge=0, le=1)
    zone_user_override: bool = Field(False)

    surfaces: list[Surface] = Field(default_factory=list)
    image_ids: list[str] = Field(default_factory=list, description="Associated image IDs")


# --- Image Level ---

class BoundingBox(BaseModel):
    """Bounding box for detected elements in an image."""

    x: float = Field(..., ge=0, le=1, description="X coordinate (normalized 0-1)")
    y: float = Field(..., ge=0, le=1, description="Y coordinate (normalized 0-1)")
    width: float = Field(..., gt=0, le=1, description="Width (normalized 0-1)")
    height: float = Field(..., gt=0, le=1, description="Height (normalized 0-1)")


class ImageAnnotation(BaseModel):
    """Annotation for a detected element in an image."""

    label: str
    bounding_box: BoundingBox
    confidence: Optional[float] = Field(None, ge=0, le=1)


class ImageMetadata(BaseModel):
    """Metadata for uploaded image."""

    id: str = Field(..., min_length=1)
    filename: str = Field(..., min_length=1)
    room_id: str = Field(..., min_length=1, description="Associated room ID")
    description: Optional[str] = Field(None, description="User description of image")

    # AI-populated fields
    detected_materials: list[MaterialType] = Field(default_factory=list)
    detected_zone: Optional[ZoneType] = None
    zone_confidence: Optional[float] = Field(None, ge=0, le=1)
    detected_condition: Optional[ConditionLevel] = None
    condition_confidence: Optional[float] = Field(None, ge=0, le=1)

    # Bounding box annotations (for UI overlay)
    annotations: list[ImageAnnotation] = Field(default_factory=list)

    analysis_complete: bool = Field(False)


# --- Qualitative Observations ---

class QualitativeObservations(BaseModel):
    """Qualitative observation checklist per FDAM 2.3."""

    smoke_fire_odor: bool = Field(..., description="Smoke/fire odor present?")
    odor_intensity: Optional[OdorIntensity] = None

    visible_soot_deposits: bool = Field(..., description="Visible soot deposits?")
    soot_pattern_description: Optional[str] = None

    large_char_particles: bool = Field(..., description="Large char particles observed?")
    char_density_estimate: Optional[CharDensity] = None

    ash_like_residue: bool = Field(..., description="Ash-like residue present?")
    ash_color_texture: Optional[str] = None

    surface_discoloration: bool = Field(..., description="Surface discoloration?")
    discoloration_description: Optional[str] = None

    dust_loading_interference: bool = Field(..., description="Dust loading or interference?")
    dust_notes: Optional[str] = None

    wildfire_indicators: bool = Field(..., description="Burned soil/pollen/vegetation indicators?")
    wildfire_notes: Optional[str] = None

    additional_notes: Optional[str] = None


# --- Complete Assessment Input ---

class AssessmentInput(BaseModel):
    """Complete input for FDAM AI assessment."""

    project: ProjectInfo
    rooms: list[Room] = Field(..., min_length=1)
    images: list[ImageMetadata] = Field(default_factory=list, max_length=20)
    observations: QualitativeObservations

    @field_validator("rooms")
    @classmethod
    def validate_room_ids(cls, rooms: list[Room]) -> list[Room]:
        """Ensure room IDs are unique."""
        ids = [r.id for r in rooms]
        if len(ids) != len(set(ids)):
            raise ValueError("Room IDs must be unique")
        return rooms

    @model_validator(mode="after")
    def validate_image_rooms(self) -> "AssessmentInput":
        """Ensure all images reference valid room IDs."""
        room_ids = {r.id for r in self.rooms}
        for img in self.images:
            if img.room_id not in room_ids:
                raise ValueError(f"Image {img.id} references unknown room {img.room_id}")
        return self