TJWaterServerBinary/app/services/timeseries_history.py

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)