models.py

"""
Pydantic models — exact mirrors of Phoebe iOS Swift DTOs.

Field names = snake_case matching CodingKeys from:
  - APIModels.swift (HeadacheLogSnapshotDTO, DailySnapshotDTO)
  - InsightPayloadDTO.swift (HealthKitMetricsDTO, WeatherDataDTO, UserContextDTO)
"""

from __future__ import annotations
from pydantic import BaseModel, Field
from typing import List, Optional
from enum import Enum


# ── Enums ────────────────────────────────────────────────────────────

class MedicationResponseEnum(str, Enum):
    horrible = "horrible"
    worse = "worse"
    same = "same"
    better = "better"
    excellent = "excellent"


# ── Diary sub-payloads ───────────────────────────────────────────────

class SymptomsPayload(BaseModel):
    symptoms: List[str] = []

class TriggersPayload(BaseModel):
    triggers: List[str] = []

class MedicationPayload(BaseModel):
    medication_taken: List[str] = []

class TherapeuticPayload(BaseModel):
    therapeutic_activities: List[str] = []


# ── HealthKit (from InsightPayloadDTO.swift) ─────────────────────────

class SleepAnalysisDTO(BaseModel):
    total_duration_hours: Optional[float] = None
    deep_sleep_minutes: Optional[float] = None
    rem_sleep_minutes: Optional[float] = None

class HRVSummaryDTO(BaseModel):
    average_ms: Optional[float] = None

class HealthKitMetricsDTO(BaseModel):
    resting_heart_rate: Optional[float] = None
    sleep_analysis: Optional[SleepAnalysisDTO] = None
    hrv_summary: Optional[HRVSummaryDTO] = None
    workout_minutes: Optional[int] = None
    had_menstrual_flow: Optional[bool] = None


# ── Weather (from InsightPayloadDTO.swift) ───────────────────────────

class WeatherDataDTO(BaseModel):
    barometric_pressure_mb: float = 1013.25
    pressure_change_24h_mb: float = 0.0
    humidity_percent: float = 50.0
    temperature_celsius: float = 15.0


# ── Headache log (from APIModels.swift) ──────────────────────────────

class HeadacheLogSnapshotDTO(BaseModel):
    severity: int = 0
    duration_hours: float = 0.0
    symptoms: SymptomsPayload = Field(default_factory=SymptomsPayload)
    triggers: TriggersPayload = Field(default_factory=TriggersPayload)
    notes: Optional[str] = None
    input_date: str = ""
    input_time: Optional[str] = None
    end_date: Optional[str] = None
    end_time: Optional[str] = None
    pain_frontal: bool = False
    pain_temporal_left: bool = False
    pain_temporal_right: bool = False
    pain_occipital: bool = False
    pain_parietal: bool = False
    pain_ocular_left: bool = False
    pain_ocular_right: bool = False
    pain_sinus: bool = False
    mood: Optional[str] = None
    medication_taken: Optional[MedicationPayload] = None
    medication_response: Optional[MedicationResponseEnum] = None
    therapeutic_activities: Optional[TherapeuticPayload] = None


# ── DailySnapshotDTO (top-level) ────────────────────────────────────

class DailySnapshotDTO(BaseModel):
    headache_log: Optional[HeadacheLogSnapshotDTO] = None
    health_kit_metrics: Optional[HealthKitMetricsDTO] = None
    weather_data: Optional[WeatherDataDTO] = None


# ── UserContextDTO ──────────────────────────────────────────────────

class UserContextDTO(BaseModel):
    age_range: Optional[str] = None
    location_region: Optional[str] = None


# ── API request / response ──────────────────────────────────────────

class PredictionRequest(BaseModel):
    user_context: Optional[UserContextDTO] = None
    daily_snapshots: List[DailySnapshotDTO]

class DayPrediction(BaseModel):
    day: int
    date: Optional[str] = None
    prediction: int
    probability: float
    risk_level: str
    top_risk_factors: List[str] = []

class PredictionResponse(BaseModel):
    predictions: List[DayPrediction]
    model_version: str
    threshold: float

# Legacy
class SinglePredictionRequest(BaseModel):
    features: List[float]

class SinglePredictionResponse(BaseModel):
    prediction: int
    probability: float