Compressor/models/common/io.py

"""CSV output helpers shared by the reorganized algorithms."""
from __future__ import annotations

import csv
import re
from pathlib import Path

from .queries import unit_sort_key


def unit_sort_key_from_unit(value):  # retained for callers that only have a unit string
    return unit_sort_key(value)


def output_sort_key(row: dict):
    """Keep exported records grouped by unit, then ordered by their start time."""
    time_field = next(
        (field for field in ("小时", "开始小时", "周期开始小时", "开始日期", "日期") if row.get(field)),
        None,
    )
    time_value = str(row.get(time_field, "")) if time_field else ""
    return (
        unit_sort_key(row.get("机组")),
        time_value,
        int(row.get("批次") or 0),
        str(row.get("周期编号", "")),
    )


def write_csv(path: Path, rows: list[dict], fields: list[str], headers: dict[str, str]):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    rows = sorted(rows, key=output_sort_key)
    with path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writerow({field: headers.get(field, field) for field in fields})
        writer.writerows(rows)


def format_number(value):
    if value is None or value == "":
        return ""
    if isinstance(value, float):
        return round(value, 6)
    return value


def clean_rows(rows):
    cleaned = []
    for row in rows:
        item = {}
        for key, value in row.items():
            if key.startswith("_"):
                continue
            item[key] = format_number(value)
        cleaned.append(item)
    return cleaned


def write_thresholds(path: Path, configs: dict[int, dict]):
    rows = []
    for batch in sorted(configs):
        config = configs[batch]
        rows.extend([
            {
                "机组": config["unit"], "批次": batch, "点位": "YSJ_5",
                "点位描述": config["metadata"]["oil_description"], "单位": config["metadata"]["oil_units"],
                "低报警阈值": config["oil"]["low"], "低低报警阈值": config["oil"]["low_low"],
                "高报警阈值": config["oil"]["high"], "高高报警阈值": config["oil"]["high_high"],
                "低报警配置": config["alarm_sources"]["oil"]["low"],
                "低低报警配置": config["alarm_sources"]["oil"]["low_low"],
                "高报警配置": config["alarm_sources"]["oil"]["high"],
                "高高报警配置": config["alarm_sources"]["oil"]["high_high"],
            },
            {
                "机组": config["unit"], "批次": batch, "点位": "YSJ_10",
                "点位描述": config["metadata"]["coupling_description"], "单位": config["metadata"]["coupling_units"],
                "低报警阈值": "", "低低报警阈值": "",
                "高报警阈值": config["coupling"]["high"], "高高报警阈值": config["coupling"]["high_high"],
                "低报警配置": "", "低低报警配置": "",
                "高报警配置": config["alarm_sources"]["coupling"]["high"],
                "高高报警配置": config["alarm_sources"]["coupling"]["high_high"],
            },
            {
                "机组": config["unit"], "批次": batch, "点位": "YSJ_11",
                "点位描述": config["metadata"]["chain_description"], "单位": config["metadata"]["chain_units"],
                "低报警阈值": "", "低低报警阈值": "",
                "高报警阈值": config["chain"]["high"], "高高报警阈值": config["chain"]["high_high"],
                "低报警配置": "", "低低报警配置": "",
                "高报警配置": config["alarm_sources"]["chain"]["high"],
                "高高报警配置": config["alarm_sources"]["chain"]["high_high"],
            },
        ])
    fields = [
        "机组", "批次", "点位", "点位描述", "单位", "低报警阈值", "低低报警阈值",
        "高报警阈值", "高高报警阈值", "低报警配置", "低低报警配置", "高报警配置", "高高报警配置",
    ]
    write_csv(path, rows, fields, {field: field for field in fields})


__all__ = [
    "output_sort_key", "write_csv", "format_number", "clean_rows", "write_thresholds",
    "unit_sort_key",
]