from collections import defaultdict from typing import Any from psycopg import AsyncConnection from app.domain.measurement_units import display_unit_for_metric from app.domain.schemas.timeseries_history import ( ElementHistoryQuery, ElementHistoryResponse, ElementHistorySeries, HistoryElementType, HistoryMetric, HistoryPoint, HistorySource, ) from app.infra.db.postgresql.analysis import AnalysisRepository from app.infra.db.postgresql.network_settings import NetworkSettingsRepository from app.infra.db.postgresql.scada import ScadaInfoRepository from app.infra.db.timescaledb.repositories.analysis import AnalysisResultsRepository from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository from app.infra.db.timescaledb.repositories.scada import ScadaRepository def _target_metric(element_type: HistoryElementType) -> HistoryMetric: return ( HistoryMetric.FLOW if element_type == HistoryElementType.PIPE else HistoryMetric.PRESSURE ) def _device_metric(device_type: str) -> HistoryMetric | None: normalized = device_type.strip().lower() if normalized in {"pipe_flow", "flow"}: return HistoryMetric.FLOW if normalized == "pressure": return HistoryMetric.PRESSURE return None def _points( rows: list[dict[str, Any]], *, value_field: str, ) -> list[HistoryPoint]: return [ HistoryPoint( time=row["time"], value=( float(row[value_field]) if row.get(value_field) is not None else None ), ) for row in rows ] class TimeseriesHistoryService: @staticmethod async def query( timescale_conn: AsyncConnection, postgres_conn: AsyncConnection, query: ElementHistoryQuery, ) -> ElementHistoryResponse: pipe_ids = list( dict.fromkeys( target.element_id for target in query.elements if target.element_type == HistoryElementType.PIPE ) ) junction_ids = list( dict.fromkeys( target.element_id for target in query.elements if target.element_type == HistoryElementType.JUNCTION ) ) devices = await ScadaInfoRepository.get_scadas_for_elements( postgres_conn, junction_ids, pipe_ids ) devices_by_element: dict[tuple[HistoryElementType, str], list[dict[str, Any]]] = ( defaultdict(list) ) for device in devices: if device.get("link_id") in pipe_ids: devices_by_element[(HistoryElementType.PIPE, device["link_id"])].append( device ) if device.get("node_id") in junction_ids: devices_by_element[ (HistoryElementType.JUNCTION, device["node_id"]) ].append(device) selected_devices: dict[str, tuple[dict[str, Any], HistoryElementType, str]] = {} for target in query.elements: target_devices = devices_by_element.get( (target.element_type, target.element_id), [] ) if target.device_ids is not None: target_device_ids = {device["device_id"] for device in target_devices} unknown = set(target.device_ids) - target_device_ids if unknown: raise ValueError( f"SCADA devices do not belong to {target.element_type.value} " f"{target.element_id}: {', '.join(sorted(unknown))}" ) requested = set(target.device_ids) target_devices = [ device for device in target_devices if device["device_id"] in requested ] expected_metric = _target_metric(target.element_type) for device in target_devices: if _device_metric(device["device_type"]) != expected_metric: continue selected_devices[device["device_id"]] = ( device, target.element_type, target.element_id, ) scada_rows = await ScadaRepository.get_scada_by_ids_time_range( timescale_conn, list(selected_devices), query.start_time, query.end_time, ) scada_by_device: dict[str, list[dict[str, Any]]] = defaultdict(list) for row in scada_rows: scada_by_device[str(row["device_id"])].append(row) series: list[ElementHistorySeries] = [] for device_id, (device, element_type, element_id) in selected_devices.items(): rows = scada_by_device.get(device_id, []) if not rows: continue metric = _target_metric(element_type) unit = device["measurement_unit"] for source, field in ( (HistorySource.SCADA_RAW, "monitored_value"), (HistorySource.SCADA_CLEANED, "cleaned_value"), ): series.append( ElementHistorySeries( element_id=element_id, element_type=element_type, device_id=device_id, metric=metric, source=source, source_unit=unit, display_unit=display_unit_for_metric(metric.value), points=_points( rows, value_field=field, ), ) ) if query.mode.value == "observed": return ElementHistoryResponse(series=series) current_units = await NetworkSettingsRepository.get_result_units(postgres_conn) if query.mode.value == "analysis_comparison": run = await AnalysisRepository.get_run(postgres_conn, query.run_id) if run is None: raise ValueError(f"analysis run does not exist: {query.run_id}") parameters = run.get("parameters") or {} recorded_units = parameters.get("result_units") unit_inferred = not isinstance(recorded_units, dict) result_units = recorded_units if not unit_inferred else current_units link_values = await AnalysisResultsRepository.get_series_by_ids( timescale_conn, query.run_id, "link", pipe_ids, query.start_time, query.end_time, "flow", ) node_values = await AnalysisResultsRepository.get_series_by_ids( timescale_conn, query.run_id, "node", junction_ids, query.start_time, query.end_time, "pressure", ) source = HistorySource.ANALYSIS_SIMULATION else: unit_inferred = False result_units = current_units link_values = await RealtimeRepository.get_link_fields_by_ids_time_range( timescale_conn, query.start_time, query.end_time, pipe_ids, "flow", ) node_values = await RealtimeRepository.get_node_fields_by_ids_time_range( timescale_conn, query.start_time, query.end_time, junction_ids, "pressure", ) source = HistorySource.REALTIME_SIMULATION for target in query.elements: metric = _target_metric(target.element_type) rows = ( link_values.get(target.element_id, []) if target.element_type == HistoryElementType.PIPE else node_values.get(target.element_id, []) ) if not rows: continue source_unit = str(result_units.get(metric.value, "")).strip() series.append( ElementHistorySeries( element_id=target.element_id, element_type=target.element_type, metric=metric, source=source, source_unit=source_unit, display_unit=display_unit_for_metric(metric.value), unit_inferred=unit_inferred, points=_points( rows, value_field="value", ), ) ) return ElementHistoryResponse(series=series)