schemas/output.py

"""Pydantic output models for FDAM AI Pipeline.

Contains vision analysis results, calculation outputs, and final assessment output.
"""

from typing import Optional

from pydantic import BaseModel, Field

from .input import (
    AssessmentInput,
    BoundingBox,
    ConditionLevel,
    MaterialCategory,
    MaterialType,
    Priority,
    SampleType,
    ZoneType,
)


# --- Vision Analysis Output ---

class ZoneAnalysis(BaseModel):
    """Zone classification from vision analysis."""

    classification: ZoneType
    confidence: float = Field(..., ge=0, le=1)
    reasoning: str


class ConditionAnalysis(BaseModel):
    """Condition assessment from vision analysis."""

    level: ConditionLevel
    confidence: float = Field(..., ge=0, le=1)
    reasoning: str


class DetectedMaterial(BaseModel):
    """Material detected in image by vision model."""

    type: MaterialType
    category: MaterialCategory
    confidence: float = Field(..., ge=0, le=1)
    location_description: Optional[str] = None
    bounding_box: Optional[BoundingBox] = None


class CombustionIndicators(BaseModel):
    """Combustion particle indicators from vision analysis."""

    soot_visible: bool = False
    soot_pattern: Optional[str] = None
    char_visible: bool = False
    char_description: Optional[str] = None
    ash_visible: bool = False
    ash_description: Optional[str] = None


class SamplingRecommendation(BaseModel):
    """Recommended sampling location from vision analysis."""

    description: str
    sample_type: SampleType
    priority: Priority


class VisionAnalysisResult(BaseModel):
    """Complete vision analysis result for a single image.

    Matches the VISION_OUTPUT_SCHEMA from the technical spec.
    """

    zone: ZoneAnalysis
    condition: ConditionAnalysis
    materials: list[DetectedMaterial] = Field(default_factory=list)
    combustion_indicators: CombustionIndicators
    structural_concerns: list[str] = Field(default_factory=list)
    access_issues: list[str] = Field(default_factory=list)
    recommended_sampling_locations: list[SamplingRecommendation] = Field(default_factory=list)
    flags_for_review: list[str] = Field(default_factory=list)


# --- Calculation Results ---

class RoomAreaSummary(BaseModel):
    """Area summary for a single room."""

    floor_area: float
    surface_area: float
    volume: float


class SurfaceAreas(BaseModel):
    """Surface area calculations by various groupings."""

    by_type: dict[str, float] = Field(default_factory=dict)
    by_disposition: dict[str, float] = Field(default_factory=dict)
    by_zone: dict[str, float] = Field(default_factory=dict)
    by_room: dict[str, RoomAreaSummary] = Field(default_factory=dict)
    total_floor_sf: float = 0
    total_surface_sf: float = 0
    total_volume_cf: float = 0


class AirFiltration(BaseModel):
    """Air filtration calculation results per NADCA ACR 2021."""

    total_volume_cf: float
    required_ach: int = 4
    unit_cfm: int = 2000
    units_required: int
    calculation: str
    standard_reference: str = "NADCA ACR 2021, Section 3.6"


class SampleDensity(BaseModel):
    """Sample density recommendations per FDAM 2.3."""

    total_sf: float
    size_category: str
    surface_types_count: int
    surface_types: list[str] = Field(default_factory=list)
    tape_lifts_per_type: str
    surface_wipes_per_type: str
    recommended_tape_lifts: int
    recommended_surface_wipes: int
    ceiling_deck_note: Optional[str] = None
    control_samples_recommended: bool = True
    control_sample_note: str = "Control samples from unaffected areas recommended for baseline comparison"


class LaborEstimate(BaseModel):
    """Labor hour estimates by task."""

    hepa_vacuum: float = 0
    wet_wipe: float = 0
    dry_sponge: float = 0
    power_wash: float = 0
    scrubber: float = 0
    removal: float = 0
    hvac_cleaning: float = 0
    total_hours: float = 0


class EquipmentRequirements(BaseModel):
    """Equipment requirements for the project."""

    air_scrubbers: int = 0
    hepa_vacuums: int = 0
    negative_air_machines: int = 0
    dehumidifiers: int = 0
    notes: list[str] = Field(default_factory=list)


class RegulatoryFlag(BaseModel):
    """Regulatory flag for potential hazards."""

    flag_type: str
    description: str
    recommendation: str
    reference: str


class RegulatoryFlags(BaseModel):
    """Regulatory flags based on construction era and facility type."""

    lead_paint_flag: Optional[RegulatoryFlag] = None
    acm_flag: Optional[RegulatoryFlag] = None
    other_flags: list[RegulatoryFlag] = Field(default_factory=list)


class CalculationResults(BaseModel):
    """All calculation results from FDAM logic engine."""

    surface_areas: SurfaceAreas
    air_filtration: AirFiltration
    sample_density: SampleDensity
    labor_estimate: LaborEstimate
    equipment: EquipmentRequirements
    regulatory_flags: RegulatoryFlags


# --- Document Output ---

class GeneratedDocuments(BaseModel):
    """Generated document outputs."""

    cleaning_specification_md: str = Field(..., description="Cleaning Specification / SOW in Markdown")
    sampling_plan_md: Optional[str] = Field(None, description="Sampling plan recommendations in Markdown")
    confidence_report_md: Optional[str] = Field(None, description="Confidence report in Markdown")


# --- Confidence Report ---

class FlaggedItem(BaseModel):
    """Item flagged for professional review."""

    type: str
    room: Optional[str] = None
    surface: Optional[str] = None
    image_id: Optional[str] = None
    confidence: Optional[float] = None
    recommendation: str


class ConfidenceReport(BaseModel):
    """Confidence report for assessment."""

    flagged_items: list[FlaggedItem] = Field(default_factory=list)
    overall_confidence: float = Field(..., ge=0, le=1)
    review_required: bool = False


# --- Complete Assessment Output ---

class AssessmentOutput(BaseModel):
    """Complete output from FDAM AI assessment pipeline."""

    # Original input (with AI-enriched fields)
    input: AssessmentInput

    # Vision analysis results (by image ID)
    vision_results: dict[str, VisionAnalysisResult] = Field(default_factory=dict)

    # Calculation results
    calculations: CalculationResults

    # Generated documents
    documents: GeneratedDocuments

    # Confidence report
    confidence_report: ConfidenceReport

    # Processing metadata
    processing_time_seconds: Optional[float] = None
    model_versions: dict[str, str] = Field(default_factory=dict)