Compressor/_tmp_temperature/data.py

"""Hourly temperature aggregation + robust run/stop inference for 4#~9#."""
from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timedelta

from models.common.db import connect

from .config import EXPECTED_PER_HOUR, TEMP_MIN_HOUR_COVERAGE, UnitConfig


def chunks(start: datetime, end: datetime, days: int):
    cursor = start
    while cursor < end:
        step = min(cursor + timedelta(days=days), end)
        yield cursor, step
        cursor = step


def _run_condition(run: dict) -> str:
    if run["type"] == "rpm":
        return f"`{run['column']}` > 0"
    return run["sql"]


def _hourly_sql(config: UnitConfig) -> tuple[str, list]:
    run = _run_condition(config.run)
    select = [
        "import_batch_id",
        "FROM_UNIXTIME((UNIX_TIMESTAMP(sample_time) DIV 3600) * 3600) AS hour_start",
        "COUNT(*) AS samples",
        f"SUM(CASE WHEN {run} THEN 1 ELSE 0 END) AS run_samples",
    ]
    if config.run["type"] == "cumulative":
        select.append(f"AVG(`{config.run['column']}`) AS cum_run")
    params: list = []
    for column in config.columns:
        limits = config.limits.get(column, {})
        high = limits.get("high")
        high_high = limits.get("high_high")
        high = 1e9 if high is None else float(high)
        high_high = 1e9 if high_high is None else float(high_high)
        select.append(f"AVG(CASE WHEN {run} AND `{column}` > 0 THEN `{column}` END) AS `{column}__avg`")
        select.append(f"MIN(CASE WHEN {run} AND `{column}` > 0 THEN `{column}` END) AS `{column}__min`")
        select.append(f"MAX(CASE WHEN {run} AND `{column}` > 0 THEN `{column}` END) AS `{column}__max`")
        select.append(f"STDDEV(CASE WHEN {run} AND `{column}` > 0 THEN `{column}` END) AS `{column}__std`")
        select.append(f"MIN(CASE WHEN {run} THEN `{column}` END) AS `{column}__minraw`")
        select.append(f"SUM(CASE WHEN {run} AND `{column}` = 0 THEN 1 ELSE 0 END) AS `{column}__zeros`")
        select.append(f"SUM(CASE WHEN ({run}) AND `{column}` >= %s THEN 1 ELSE 0 END) AS `{column}__high`")
        select.append(f"SUM(CASE WHEN ({run}) AND `{column}` >= %s THEN 1 ELSE 0 END) AS `{column}__high_high`")
        params.extend([high, high_high])
    sql = (
        "SELECT " + ",\n    ".join(select) +
        "\nFROM pks_long_sample\n"
        "WHERE import_batch_id = %s AND sample_time >= %s AND sample_time < %s\n"
        "GROUP BY import_batch_id, FROM_UNIXTIME((UNIX_TIMESTAMP(sample_time) DIV 3600) * 3600)\n"
        "ORDER BY hour_start"
    )
    return sql, params


def _apply_run_state(rows: list[dict], config: UnitConfig) -> None:
    """Infer run/stop with hysteresis and drop start/stop transition hours."""
    grouped: dict[int, list[dict]] = defaultdict(list)
    for row in rows:
        grouped[row["import_batch_id"]].append(row)

    for batch_rows in grouped.values():
        batch_rows.sort(key=lambda row: row["hour_start"])
        raw: list[bool] = []
        if config.run["type"] == "cumulative":
            previous_value = previous_dt = None
            for row in batch_rows:
                value = row.get("cum_run")
                state = None
                if value is not None and previous_value is not None:
                    gap = (row["hour_start"] - previous_dt).total_seconds() / 3600
                    delta = value - previous_value
                    if gap <= 1.5 and -1.0 < delta < 2.0:
                        if delta >= 0.5:
                            state = True
                        elif abs(delta) < 0.05:
                            state = False
                if state is None:
                    state = (row.get("coverage") or 0.0) >= TEMP_MIN_HOUR_COVERAGE
                raw.append(state)
                previous_value, previous_dt = value, row["hour_start"]
        elif config.run["type"] == "multi":
            raw = [(row.get("coverage") or 0.0) >= 0.5 for row in batch_rows]
        else:
            raw = [(row.get("coverage") or 0.0) >= TEMP_MIN_HOUR_COVERAGE for row in batch_rows]

        # 2-hour hysteresis: switch state only after two consecutive opposite hours.
        states: list[bool] = []
        current = raw[0] if raw else False
        pending = 0
        for value in raw:
            if value == current:
                pending = 0
            else:
                pending += 1
                if pending >= 2:
                    current = value
                    pending = 0
            states.append(current)

        transition = [False] * len(states)
        for index, state in enumerate(states):
            if not state:
                continue
            if index == 0 or not states[index - 1]:
                transition[index] = True
                if index + 1 < len(states):
                    transition[index + 1] = True
            if index == len(states) - 1 or not states[index + 1]:
                transition[index] = True

        for row, state, trans in zip(batch_rows, states, transition):
            row["run_state"] = bool(state)
            row["transition"] = bool(trans)
            row["valid"] = bool(state and not trans)


def fetch_unit_hourly(config: UnitConfig, start: datetime, end: datetime, chunk_days: int = 30) -> list[dict]:
    sql, params = _hourly_sql(config)
    rows: list[dict] = []
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for batch in config.batches:
                for chunk_start, chunk_end in chunks(start, end, chunk_days):
                    cursor.execute(sql, [*params, batch, chunk_start, chunk_end])
                    for row in cursor.fetchall():
                        row["机组"] = config.name
                        row["coverage"] = (row.get("run_samples") or 0) / EXPECTED_PER_HOUR
                        rows.append(row)
    finally:
        connection.close()
    rows.sort(key=lambda row: (row["import_batch_id"], row["hour_start"]))
    _apply_run_state(rows, config)
    return rows


__all__ = ["fetch_unit_hourly"]