next-tjwater-server/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.postgresql import read_queries
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
    ]


async def _query_page(timescale_conn, postgres_conn, query, selected_devices, limit, offset):
    descriptors = []
    series = []

    def add(*, source, element_type, element_id, device_id=None, field, unit, unit_inferred=False):
        descriptors.append({
            "source": source, "element_type": element_type,
            "id": device_id if source == "scada" else element_id, "field": field,
        })
        metric = _target_metric(element_type)
        series.append(ElementHistorySeries(
            element_id=element_id, element_type=element_type, device_id=device_id,
            metric=metric, source=(
                HistorySource.SCADA_RAW if field == "monitored_value" else HistorySource.SCADA_CLEANED
            ) if source == "scada" else (
                HistorySource.ANALYSIS_SIMULATION if source == "analysis" else HistorySource.REALTIME_SIMULATION
            ),
            source_unit=unit, display_unit=display_unit_for_metric(metric.value),
            unit_inferred=unit_inferred, points=[],
        ))

    # Stable series ordering makes a page crossing raw/cleaned/model boundaries repeatable.
    for device_id, (device, element_type, element_id) in sorted(selected_devices.items()):
        for field in ("monitored_value", "cleaned_value"):
            add(source="scada", element_type=element_type, element_id=element_id,
                device_id=device_id, field=field, unit=device["measurement_unit"])
    if query.mode.value != "observed":
        units = await NetworkSettingsRepository.get_result_units(postgres_conn)
        source = "analysis" if query.mode.value == "analysis_comparison" else "realtime"
        inferred = False
        if source == "analysis":
            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}")
            recorded = (run.get("parameters") or {}).get("result_units")
            inferred = not isinstance(recorded, dict)
            if not inferred:
                units = recorded
        targets = {(t.element_type, t.element_id) for t in query.elements}
        for element_type, element_id in sorted(targets):
            metric = _target_metric(element_type)
            add(source=source, element_type=element_type, element_id=element_id,
                field=metric.value, unit=str(units.get(metric.value, "")).strip(), unit_inferred=inferred)
    rows = await read_queries.history_rows(
        timescale_conn, descriptors, start=query.start_time, end=query.end_time,
        run_id=query.run_id, limit=limit + 1, offset=offset,
    )
    for row in rows[:limit]:
        series[row["series_index"]].points.append(HistoryPoint(time=row["time"], value=row["value"]))
    return ElementHistoryResponse(series=[s for s in series if s.points], truncated=len(rows) > limit)


class TimeseriesHistoryService:
    @staticmethod
    async def query(
        timescale_conn: AsyncConnection,
        postgres_conn: AsyncConnection,
        query: ElementHistoryQuery,
        *, row_limit: int | None = None, row_offset: int = 0,
    ) -> ElementHistoryResponse:
        if row_offset < 0 or (row_offset and row_limit is None) or (row_limit is not None and row_limit < 1):
            raise ValueError("row_offset requires a bounded page")

        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, **({"row_limit": 1001} if row_limit is not None else {})
        )
        if row_limit is not None and len(devices) > 1000:
            raise ValueError("too many device mappings; narrow the element selection")
        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,
                )

        if row_limit is not None and len(selected_devices) > 20:
            raise ValueError("select at most 20 devices explicitly")
        if row_limit is not None:
            return await _query_page(
                timescale_conn, postgres_conn, query, selected_devices, row_limit, row_offset,
            )
        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)