Compressor/backend/model_assistant_generalized/build_datasets.py

#!/usr/bin/env python3
"""Create reviewable configurable-window PKS samples from scanner hourly data.

This module deliberately imports the scanner's read and aggregation helpers.
It must not create an independent interpretation of PKS timestamps, coverage,
thresholds, or batch-to-unit mapping.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import shutil
import sys
from collections import Counter, defaultdict
from datetime import datetime, timedelta
from pathlib import Path
from statistics import median

ROOT = Path(__file__).resolve().parents[2]
BACKEND = ROOT / "backend"
if str(BACKEND) not in sys.path:
    sys.path.insert(0, str(BACKEND))

import scan_pks_fault_evolution as scanner  # noqa: E402

WINDOW_HOURS = 6
ISOLATION_HOURS = 12
RESTART_GAP_HOURS = 12
DEFAULT_START = "2025-01-01 00:00:00"
DEFAULT_END = "2026-09-01 00:00:00"


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


def finite(value) -> float | None:
    if value is None or value == "":
        return None
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if math.isfinite(number) else None


def mean(values: list[float | None]) -> float | None:
    valid = [value for value in values if value is not None]
    return sum(valid) / len(valid) if valid else None


def value_range(values: list[float | None]) -> float | None:
    valid = [value for value in values if value is not None]
    return max(valid) - min(valid) if valid else None


def slope(values: list[float | None]) -> float | None:
    return scanner.linear_slope(values)


def ratio(count, running) -> float:
    return float(count or 0) / max(float(running or 0), 1.0)


def csv_rows(path: Path) -> list[dict]:
    with path.open(encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def stage_ranges(path: Path, stage_field: str) -> list[dict]:
    ranges = []
    for index, row in enumerate(csv_rows(path), 1):
        ranges.append({
            "stage_id": f"{stage_field}_{row['机组']}_{row['批次']}_{index:04d}",
            "unit": row["机组"],
            "batch": int(row["批次"]),
            "level": row[stage_field],
            "start": parse_dt(row["开始小时"]),
            "end": parse_dt(row["结束小时"]),
            "confirmed": row.get("是否形成确认等级", ""),
        })
    return ranges


def fetch_hourly(start: datetime, end: datetime, chunk_days: int) -> tuple[list[dict], dict[int, dict]]:
    batch_to_unit = scanner.load_pks_batches()
    configs = scanner.load_alarm_config(batch_to_unit)
    ranges = scanner.load_batch_ranges(sorted(batch_to_unit), start, end)
    grouped = defaultdict(list)
    for batch in sorted(ranges):
        grouped[batch_to_unit[batch]].append(batch)
    rows = []
    for unit, batches in sorted(grouped.items(), key=lambda item: scanner._unit_sort_key(item[0])):
        unit_start = min(ranges[batch][0] for batch in batches)
        unit_end = max(ranges[batch][1] for batch in batches)
        rows.extend(scanner.fetch_hourly(unit, batches, unit_start, unit_end, chunk_days, configs[batches[0]]))
    rows.sort(key=lambda row: (scanner._unit_sort_key(row["机组"]), row["批次"], row["_dt"]))
    return rows, configs


def load_hourly_csv(path: Path) -> list[dict]:
    """Rebuild the hourly row structure from a saved hourly.csv without querying the DB."""
    result = []
    for row in csv_rows(path):
        item = {
            "机组": row["机组"],
            "批次": int(row["批次"]),
            "小时": row["小时"],
            "_dt": parse_dt(row["小时"]),
            "样本数": int(row["样本数"]),
            "运行样本数": int(row["运行样本数"]),
            "运行覆盖率": float(row["运行覆盖率"]),
        }
        for field in scanner.NUMERIC_HOURLY_FIELDS:
            item[field] = finite(row.get(field))
        result.append(item)
    result.sort(key=lambda item: (scanner._unit_sort_key(item["机组"]), item["批次"], item["_dt"]))
    return result


def assign_run_cycles(rows: list[dict], model_type: str) -> list[dict]:
    grouped = defaultdict(list)
    for row in rows:
        if row["运行覆盖率"] < scanner.MIN_HOUR_COVERAGE:
            continue
        if model_type == "pressure" and row["oil_avg"] is None:
            continue
        if model_type == "vibration" and (row["coupling_avg"] is None or row["chain_avg"] is None):
            continue
        grouped[(row["机组"], row["批次"])].append(row)

    result = []
    for (unit, batch), series in grouped.items():
        series.sort(key=lambda row: row["_dt"])
        cycle = 0
        previous = None
        for row in series:
            restart = previous is None or row["_dt"] - previous > timedelta(hours=RESTART_GAP_HOURS)
            if restart:
                cycle += 1
            item = dict(row)
            item["run_cycle_id"] = f"{unit}_B{batch}_R{cycle:04d}"
            item["restart_hour"] = restart
            result.append(item)
            previous = row["_dt"]
    return sorted(result, key=lambda row: (scanner._unit_sort_key(row["机组"]), row["批次"], row["_dt"]))


def previous_baseline(series: list[dict], index: int, field: str, hours: int) -> tuple[float | None, float | None]:
    now = series[index]["_dt"]
    prior = [row[field] for row in series[:index] if row["_dt"] >= now - timedelta(hours=hours) and row[field] is not None]
    return scanner.robust_baseline(prior)


def source_hour(row: dict) -> dict:
    running = row["运行样本数"]
    return {
        "hour": row["小时"],
        "running_samples": running,
        "running_coverage": row["运行覆盖率"],
        "oil_avg": row["oil_avg"], "oil_min": row["oil_min"], "oil_max": row["oil_max"],
        "oil_low_ratio": ratio(row["oil_low_count"], running),
        "oil_low_low_ratio": ratio(row["oil_low_low_count"], running),
        "coupling_avg": row["coupling_avg"], "coupling_max": row["coupling_max"],
        "chain_avg": row["chain_avg"], "chain_max": row["chain_max"],
        "coupling_high_ratio": ratio(row["coupling_high_count"], running),
        "coupling_high_high_ratio": ratio(row["coupling_high_high_count"], running),
        "chain_high_ratio": ratio(row["chain_high_count"], running),
        "chain_high_high_ratio": ratio(row["chain_high_high_count"], running),
    }


def cycle_baseline(series: list[dict], days: int = scanner.BASELINE_DAYS) -> float | None:
    """Median of daily oil medians over the first `days` calendar days of a cycle.

    Mirrors the scanner's cycle baseline: the baseline is taken from the start of
    the run cycle, so it never uses information from the window being scored.
    """
    by_day: dict = {}
    for row in series:
        if row["oil_avg"] is None:
            continue
        day = row["_dt"].date()
        if day not in by_day and len(by_day) >= days:
            break
        by_day.setdefault(day, []).append(row["oil_avg"])
    daily = [scanner.median_value(values) for _, values in sorted(by_day.items())[:days]]
    daily = [value for value in daily if value is not None]
    return scanner.median_value(daily)


def decline_hours(series: list[dict], end_index: int) -> int:
    """Consecutive valid hours of oil-pressure decline ending at end_index."""
    count = 0
    for index in range(end_index, 0, -1):
        current, previous = series[index]["oil_avg"], series[index - 1]["oil_avg"]
        if current is None or previous is None or current >= previous:
            break
        count += 1
    return count


def pressure_features(series: list[dict], end_index: int, window: list[dict]) -> dict:
    oil = [row["oil_avg"] for row in window]
    coverages = [row["运行覆盖率"] for row in window]
    baseline = scanner.median_value([row["oil_avg"] for row in window[:-1]])
    current = oil[-1]
    oil_mean, oil_min, oil_max = mean(oil), min(oil), max(oil)
    oil_range = value_range(oil)
    oil_slope = slope(oil)
    cycle_base = cycle_baseline(series)
    recent_center, _ = previous_baseline(series, end_index, "oil_avg", 24)
    return {
        "oil_mean": oil_mean, "oil_min": oil_min, "oil_max": oil_max, "oil_range": oil_range,
        "oil_slope": oil_slope,
        "oil_last_relative_to_window": (current / baseline - 1) if baseline else None,
        "oil_low_ratio_mean": mean([ratio(row["oil_low_count"], row["运行样本数"]) for row in window]),
        "oil_low_low_ratio_mean": mean([ratio(row["oil_low_low_count"], row["运行样本数"]) for row in window]),
        "running_coverage_mean": mean(coverages), "running_coverage_min": min(coverages),
        "oil_rel_cycle_mean": (oil_mean / cycle_base - 1) if cycle_base else None,
        "oil_rel_cycle_min": (oil_min / cycle_base - 1) if cycle_base else None,
        "oil_rel_cycle_last": (current / cycle_base - 1) if cycle_base else None,
        "oil_rel_recent_mean": (oil_mean / recent_center - 1) if recent_center else None,
        "oil_rel_recent_last": (current / recent_center - 1) if recent_center else None,
        "oil_slope_rel_cycle": (oil_slope / cycle_base) if cycle_base else None,
        "oil_decline_hours": decline_hours(series, end_index),
        "oil_drawdown_cycle": ((cycle_base - oil_min) / cycle_base) if cycle_base else None,
        "oil_range_rel_cycle": (oil_range / cycle_base) if cycle_base else None,
    }


def vibration_features(series: list[dict], end_index: int, window: list[dict]) -> dict:
    features = {
        "running_coverage_mean": mean([row["运行覆盖率"] for row in window]),
        "running_coverage_min": min(row["运行覆盖率"] for row in window),
    }
    for side in ("coupling", "chain"):
        averages = [row[f"{side}_avg"] for row in window]
        peaks = [row[f"{side}_max"] for row in window]
        avg_center, avg_scale = previous_baseline(series, end_index, f"{side}_avg", scanner.VIBRATION_BASELINE_DAYS * 24)
        peak_center, peak_scale = previous_baseline(series, end_index, f"{side}_max", scanner.VIBRATION_BASELINE_DAYS * 24)
        current_avg, current_peak = averages[-1], peaks[-1]
        features.update({
            f"{side}_avg_mean": mean(averages), f"{side}_avg_max": max(averages), f"{side}_avg_range": value_range(averages),
            f"{side}_avg_slope": slope(averages), f"{side}_peak_mean": mean(peaks), f"{side}_peak_max": max(peaks),
            f"{side}_peak_range": value_range(peaks), f"{side}_peak_slope": slope(peaks),
            f"{side}_avg_baseline_deviation": ((current_avg - avg_center) / avg_scale) if avg_center is not None and avg_scale else None,
            f"{side}_peak_baseline_deviation": ((current_peak - peak_center) / peak_scale) if peak_center is not None and peak_scale else None,
            f"{side}_avg_rel_last": (current_avg / avg_center - 1) if avg_center else None,
            f"{side}_peak_rel_last": (current_peak / peak_center - 1) if peak_center else None,
            f"{side}_avg_rel_mean": (mean(averages) / avg_center - 1) if avg_center else None,
            f"{side}_peak_rel_max": (max(peaks) / peak_center - 1) if peak_center else None,
            f"{side}_high_ratio_mean": mean([ratio(row[f"{side}_high_count"], row["运行样本数"]) for row in window]),
            f"{side}_high_high_ratio_mean": mean([ratio(row[f"{side}_high_high_count"], row["运行样本数"]) for row in window]),
        })
    features["both_sides_rising"] = int(features["coupling_avg_slope"] > 0 and features["chain_avg_slope"] > 0)
    return features


def window_status(window: list[dict], stages: list[dict]) -> tuple[int | None, str, str | None, str | None]:
    start, end = window[0]["_dt"], window[-1]["_dt"]
    overlapping = [stage for stage in stages if stage["start"] <= end and stage["end"] >= start]
    if overlapping:
        stage = min(overlapping, key=lambda item: (item["start"], item["stage_id"]))
        return 1, "rule_stage", stage["stage_id"], None
    nearby = [stage for stage in stages if stage["start"] - timedelta(hours=ISOLATION_HOURS) <= end and stage["end"] + timedelta(hours=ISOLATION_HOURS) >= start]
    if nearby:
        return None, "excluded", None, "anomaly_isolation_period"
    return 0, "stable_normal", None, None


def build_samples(rows: list[dict], model_type: str, stages: list[dict], window_hours: int) -> list[dict]:
    cycle_rows = assign_run_cycles(rows, model_type)
    by_cycle = defaultdict(list)
    for row in cycle_rows:
        by_cycle[(row["机组"], row["批次"], row["run_cycle_id"])].append(row)
    # The scanner's vibration baseline is the robust baseline over the device's
    # full valid history in the prior 7*24 hours, across restarts, excluding
    # restart transition hours.  Reproduce that here so features match the rule.
    full_by_batch: dict[tuple, list[dict]] = {}
    index_by_batch: dict[tuple, dict] = {}
    if model_type == "vibration":
        grouped = defaultdict(list)
        for row in cycle_rows:
            if not row["restart_hour"]:
                grouped[(row["机组"], row["批次"])].append(row)
        for key, series in grouped.items():
            series.sort(key=lambda row: row["_dt"])
            full_by_batch[key] = series
            index_by_batch[key] = {row["_dt"]: index for index, row in enumerate(series)}
    samples = []
    sequence = 0
    for (unit, batch, cycle_id), series in sorted(by_cycle.items()):
        series.sort(key=lambda row: row["_dt"])
        for end_index in range(window_hours - 1, len(series)):
            window = series[end_index - window_hours + 1:end_index + 1]
            if any(window[index]["_dt"] - window[index - 1]["_dt"] > timedelta(hours=RESTART_GAP_HOURS) for index in range(1, len(window))):
                continue
            if any(row["restart_hour"] for row in window):
                exclude_reason = "contains_restart_transition_hour"
                label, source, stage_id = None, "excluded", None
            else:
                scoped_stages = [stage for stage in stages if stage["unit"] == unit and stage["batch"] == batch]
                label, source, stage_id, exclude_reason = window_status(window, scoped_stages)
            sequence += 1
            sample_id = f"{model_type}_{unit}_B{batch}_{window[0]['_dt']:%Y%m%d%H}_{window[-1]['_dt']:%Y%m%d%H}_{sequence:06d}"
            event_group_id = stage_id or f"{cycle_id}_normal"
            if model_type == "pressure":
                features = pressure_features(series, end_index, window)
            else:
                key = (unit, batch)
                full_series = full_by_batch[key]
                full_index = index_by_batch[key][window[-1]["_dt"]]
                features = vibration_features(full_series, full_index, window)
            samples.append({
                "sample_id": sample_id, "model_type": model_type, "unit": unit, "batch": batch,
                "run_cycle_id": cycle_id, "event_group_id": event_group_id,
                "start_hour": window[0]["小时"], "end_hour": window[-1]["小时"],
                "valid_hour_count": window_hours,
                "natural_duration_hours": int((window[-1]["_dt"] - window[0]["_dt"]).total_seconds() / 3600) + 1,
                "label": label, "label_name": "anomaly" if label == 1 else "normal" if label == 0 else "excluded",
                "label_source": source, "label_confidence": "initial_rule" if label is not None else "not_usable",
                "rule_stage_id": stage_id, "exclude_reason": exclude_reason,
                "features": features, "source_hours": [source_hour(row) for row in window],
            })
    return samples


def write_hourly(path: Path, rows: list[dict]) -> None:
    fields = ["机组", "批次", "小时", "样本数", "运行样本数", "运行覆盖率", *scanner.NUMERIC_HOURLY_FIELDS]
    with path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            writer.writerow({field: row.get(field) for field in fields})


def write_samples(base: Path, samples: list[dict]) -> Counter:
    for folder in ("candidates", "excluded"):
        (base / folder).mkdir(parents=True, exist_ok=True)
    manifest_fields = ["sample_id", "model_type", "unit", "batch", "run_cycle_id", "event_group_id", "start_hour", "end_hour", "label", "label_name", "label_source", "label_confidence", "rule_stage_id", "exclude_reason", "split"]
    counts = Counter()
    with (base / "manifest.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=manifest_fields)
        writer.writeheader()
        for sample in samples:
            folder = "excluded" if sample["label"] is None else "candidates"
            with (base / folder / f"{sample['sample_id']}.json").open("w", encoding="utf-8") as output:
                json.dump(sample, output, ensure_ascii=False, indent=2, allow_nan=False)
            writer.writerow({field: sample.get(field, "") for field in manifest_fields})
            counts[(folder, sample["label_name"], sample.get("exclude_reason") or "")] += 1
    return counts


def main() -> None:
    parser = argparse.ArgumentParser(description="从现有PKS扫描器构建可配置小时窗口的模型辅助样本")
    parser.add_argument("--start", default=DEFAULT_START)
    parser.add_argument("--end", default=DEFAULT_END)
    parser.add_argument("--chunk-days", type=int, default=14)
    parser.add_argument("--window-hours", type=int, default=WINDOW_HOURS)
    parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parent / "datasets"))
    parser.add_argument("--evolution-dir", default=str(ROOT / "cache" / "pks_scanner" / "evolution"))
    parser.add_argument("--hourly-csv", default=None, help="复用已保存的 hourly.csv，不查询数据库")
    parser.add_argument("--model-types", nargs="+", choices=("pressure", "vibration"), default=["pressure", "vibration"])
    parser.add_argument("--units", nargs="+", default=None, help="只构建指定机组，例如 8号机")
    args = parser.parse_args()
    start, end = parse_dt(args.start), parse_dt(args.end)
    if end <= start or args.chunk_days < 1 or args.window_hours < 1:
        parser.error("时间范围或 chunk-days 无效")
    output = Path(args.output_dir)
    evolution = Path(args.evolution_dir)
    for name in ("润滑油压力小时阶段段.csv", "振动异常段.csv"):
        if not (evolution / name).is_file():
            parser.error(f"缺少扫描器结果：{evolution / name}")
    if output.exists():
        shutil.rmtree(output)
    output.mkdir(parents=True)
    if args.hourly_csv:
        hourly_path = Path(args.hourly_csv)
        if not hourly_path.is_file():
            parser.error(f"缺少 hourly.csv：{hourly_path}")
        print(f"复用小时中间表：{hourly_path}")
        rows = load_hourly_csv(hourly_path)
    else:
        print("只读模式：复用 scan_pks_fault_evolution.py 的小时聚合和数据口径")
        rows, _ = fetch_hourly(start, end, args.chunk_days)
    pressure_stages = stage_ranges(evolution / "润滑油压力小时阶段段.csv", "压力阶段")
    vibration_stages = stage_ranges(evolution / "振动异常段.csv", "振动阶段")
    selected_units = set(args.units or [])
    for model_type in args.model_types:
        stages = pressure_stages if model_type == "pressure" else vibration_stages
        scoped_rows = [row for row in rows if not selected_units or row["机组"] in selected_units]
        scoped_stages = [stage for stage in stages if not selected_units or stage["unit"] in selected_units]
        base = output / model_type
        base.mkdir()
        write_hourly(base / "hourly.csv", scoped_rows)
        counts = write_samples(base, build_samples(scoped_rows, model_type, scoped_stages, args.window_hours))
        print(f"{model_type} samples: {dict(counts)}")
    print(f"小时中间数据：{len(rows):,} 行")
    print(f"样本输出目录：{output}")


if __name__ == "__main__":
    main()