Compressor/models/common/queries.py

"""Hourly aggregation queries and the shared hourly intermediate table."""
from __future__ import annotations

import csv
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

from . import stats
from .constants import EXPECTED_PER_HOUR, MIN_HOUR_COVERAGE
from .db import connect
from .settings import BACKEND  # noqa: F401  (ensures backend path is importable)

HOURLY_SQL = """
SELECT
    import_batch_id,
    FROM_UNIXTIME((UNIX_TIMESTAMP(sample_time) DIV 3600) * 3600) AS hour_start,
    COUNT(*) AS samples,
    SUM(CASE WHEN YSJ_41 > 0 THEN 1 ELSE 0 END) AS running_samples,
    AVG(CASE WHEN YSJ_41 > 0 THEN YSJ_5 END) AS oil_avg,
    MIN(CASE WHEN YSJ_41 > 0 THEN YSJ_5 END) AS oil_min,
    MAX(CASE WHEN YSJ_41 > 0 THEN YSJ_5 END) AS oil_max,
    AVG(CASE WHEN YSJ_41 > 0 THEN YSJ_10 END) AS coupling_avg,
    MAX(CASE WHEN YSJ_41 > 0 THEN YSJ_10 END) AS coupling_max,
    AVG(CASE WHEN YSJ_41 > 0 THEN YSJ_11 END) AS chain_avg,
    MAX(CASE WHEN YSJ_41 > 0 THEN YSJ_11 END) AS chain_max,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_5 <= %s THEN 1 ELSE 0 END) AS oil_low_count,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_5 <= %s THEN 1 ELSE 0 END) AS oil_low_low_count,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_10 >= %s THEN 1 ELSE 0 END) AS coupling_high_count,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_10 >= %s THEN 1 ELSE 0 END) AS coupling_high_high_count,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_11 >= %s THEN 1 ELSE 0 END) AS chain_high_count,
    SUM(CASE WHEN YSJ_41 > 0 AND YSJ_11 >= %s THEN 1 ELSE 0 END) AS chain_high_high_count
FROM pks_long_sample
WHERE import_batch_id = %s
  AND sample_time >= %s
  AND sample_time < %s
GROUP BY import_batch_id, FROM_UNIXTIME((UNIX_TIMESTAMP(sample_time) DIV 3600) * 3600)
ORDER BY hour_start, import_batch_id
"""

NUMERIC_HOURLY_FIELDS = (
    "oil_avg", "oil_min", "oil_max", "coupling_avg", "coupling_max",
    "chain_avg", "chain_max", "oil_low_count", "oil_low_low_count",
    "coupling_high_count", "coupling_high_high_count", "chain_high_count",
    "chain_high_high_count",
)


def parse_dt(value: str) -> datetime:
    return datetime.strptime(value, "%Y-%m-%d %H:%M:%S")


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


def unit_sort_key(value):
    import re

    match = re.search(r"[789]", str(value or ""))
    return int(match.group()) if match else 999


def load_batch_ranges(batches: list[int], start: datetime, end: datetime) -> dict[int, tuple[datetime, datetime]]:
    placeholders = ",".join(["%s"] * len(batches))
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                SELECT import_batch_id, MIN(sample_time) AS data_start, MAX(sample_time) AS data_end
                FROM pks_long_sample
                WHERE import_batch_id IN ({placeholders})
                  AND sample_time >= %s AND sample_time < %s
                GROUP BY import_batch_id
                """,
                (*batches, start, end),
            )
            return {
                int(row["import_batch_id"]): (
                    max(row["data_start"], start),
                    min(row["data_end"] + timedelta(seconds=5), end),
                )
                for row in cursor.fetchall()
            }
    finally:
        connection.close()


def query_params(batches: list[int], start: datetime, end: datetime, config: dict) -> tuple:
    return (
        config["oil"]["low"], config["oil"]["low_low"],
        config["coupling"]["high"], config["coupling"]["high_high"],
        config["chain"]["high"], config["chain"]["high_high"],
        *batches, start, end,
    )


def fetch_hourly(unit: str, batches: list[int], start: datetime, end: datetime, chunk_days: int, config: dict):
    rows = []
    placeholders = ",".join(["%s"] * len(batches))
    sql = HOURLY_SQL.replace("import_batch_id = %s", f"import_batch_id IN ({placeholders})")
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for index, (chunk_start, chunk_end) in enumerate(chunks(start, end, chunk_days), 1):
                params = query_params(batches, chunk_start, chunk_end, config)
                cursor.execute("EXPLAIN " + sql, params)
                plan = cursor.fetchone()
                if not plan or plan.get("key") not in ("PRIMARY",):
                    raise RuntimeError(
                        f"安全检查失败：unit={unit} batches={batches} chunk={chunk_start} "
                        f"EXPLAIN 未使用 PRIMARY: {plan}"
                    )
                cursor.execute(sql, params)
                chunk_count = 0
                for raw in cursor.fetchall():
                    batch = int(raw["import_batch_id"])
                    row = {
                        "批次": batch,
                        "机组": unit,
                        "小时": raw["hour_start"].strftime("%Y-%m-%d %H:%M:%S"),
                        "样本数": int(raw["samples"] or 0),
                        "运行样本数": int(raw["running_samples"] or 0),
                    }
                    row["运行覆盖率"] = row["运行样本数"] / EXPECTED_PER_HOUR
                    for field in NUMERIC_HOURLY_FIELDS:
                        row[field] = stats.number(raw[field])
                    row["_dt"] = raw["hour_start"]
                    rows.append(row)
                    chunk_count += 1
                print(
                    f"unit={unit} batches={batches} chunk={index} {chunk_start}..{chunk_end}: "
                    f"{chunk_count:,} hourly rows"
                )
    finally:
        connection.close()
    return rows


def fetch_all_hourly(start: datetime, end: datetime, chunk_days: int, batch_to_unit: dict[int, str], configs: dict[int, dict], selected_batches=None):
    """Fetch hourly rows for every selected batch, grouped by unit like the scanner."""
    selected = selected_batches or sorted(batch_to_unit)
    batch_ranges = load_batch_ranges(selected, start, end)
    grouped = defaultdict(list)
    for batch in selected:
        if batch in batch_ranges:
            grouped[batch_to_unit[batch]].append(batch)
    rows = []
    for unit, batches in sorted(grouped.items(), key=lambda item: unit_sort_key(item[0])):
        unit_start = min(batch_ranges[batch][0] for batch in batches)
        unit_end = max(batch_ranges[batch][1] for batch in batches)
        rows.extend(fetch_hourly(unit, sorted(batches), unit_start, unit_end, chunk_days, configs[batches[0]]))
    rows.sort(key=lambda row: (unit_sort_key(row["机组"]), row["_dt"], row["批次"]))
    return rows


HOURLY_CSV_FIELDS = ["机组", "批次", "小时", "样本数", "运行样本数", "运行覆盖率", *NUMERIC_HOURLY_FIELDS]


def write_hourly_csv(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=HOURLY_CSV_FIELDS)
        writer.writeheader()
        for row in rows:
            writer.writerow({field: row.get(field) for field in HOURLY_CSV_FIELDS})


def load_hourly_csv(path: Path) -> list[dict]:
    """Rebuild the hourly row structure from a saved hourly.csv without querying the DB."""
    with Path(path).open(encoding="utf-8-sig", newline="") as handle:
        raw_rows = list(csv.DictReader(handle))
    result = []
    for row in raw_rows:
        item = {
            "机组": row["机组"],
            "批次": int(row["批次"]),
            "小时": row["小时"],
            "_dt": parse_dt(row["小时"]),
            "样本数": int(row["样本数"]),
            "运行样本数": int(row["运行样本数"]),
            "运行覆盖率": float(row["运行覆盖率"]),
        }
        for field in NUMERIC_HOURLY_FIELDS:
            item[field] = stats.finite(row.get(field))
        result.append(item)
    result.sort(key=lambda item: (unit_sort_key(item["机组"]), item["批次"], item["_dt"]))
    return result


__all__ = [
    "HOURLY_SQL", "NUMERIC_HOURLY_FIELDS", "HOURLY_CSV_FIELDS",
    "parse_dt", "chunks", "unit_sort_key", "load_batch_ranges", "query_params",
    "fetch_hourly", "fetch_all_hourly", "write_hourly_csv", "load_hourly_csv",
    "MIN_HOUR_COVERAGE",
]