Compressor/models/common/windowing.py

"""Shared window construction for the pressure and vibration models.

These helpers reproduce ``backend/model_assistant_generalized/build_datasets.py``
and only expose the parts that are identical for both model types.  The model
specific feature extraction is injected via the ``feature_fn`` argument.
"""
from __future__ import annotations

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

from .constants import (
    BASELINE_DAYS,
    MIN_HOUR_COVERAGE,
    RESET_GAP_HOURS,
    VIBRATION_BASELINE_DAYS,
    VIBRATION_RESTART_GAP_HOURS,
)
from .queries import NUMERIC_HOURLY_FIELDS, parse_dt, unit_sort_key
from .stats import median_value, robust_baseline

WINDOW_HOURS = 6
ISOLATION_HOURS = 12
RESTART_GAP_HOURS = RESET_GAP_HOURS


def finite(value):
    from .stats import finite as _finite

    return _finite(value)


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


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


def slope(values):
    from .stats import linear_slope

    return linear_slope(values)


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


def csv_rows(path: Path) -> list[dict]:
    with Path(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 assign_run_cycles(rows: list[dict], model_type: str) -> list[dict]:
    grouped = defaultdict(list)
    for row in rows:
        if row["运行覆盖率"] < 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: (unit_sort_key(row["机组"]), row["批次"], row["_dt"]))


def previous_baseline(series: list[dict], index: int, field: str, hours: int):
    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 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 = BASELINE_DAYS):
    """Median of daily oil medians over the first ``days`` calendar days of a cycle."""
    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 = [median_value(values) for _, values in sorted(by_day.items())[:days]]
    daily = [value for value in daily if value is not None]
    return 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 window_status(window: list[dict], stages: list[dict]):
    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, feature_fn):
    """Build window samples; ``feature_fn(series, end_index, window) -> 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):
                label, source, stage_id, exclude_reason = None, "excluded", None, "contains_restart_transition_hour"
            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 = feature_fn(series, end_index, window)
            else:
                key = (unit, batch)
                full_series = full_by_batch[key]
                full_index = index_by_batch[key][window[-1]["_dt"]]
                features = feature_fn(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_samples(base: Path, samples: list[dict]) -> Counter:
    base = Path(base)
    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 write_hourly(path: Path, rows: list[dict]) -> None:
    fields = ["机组", "批次", "小时", "样本数", "运行样本数", "运行覆盖率", *NUMERIC_HOURLY_FIELDS]
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    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})


__all__ = [
    "WINDOW_HOURS", "ISOLATION_HOURS", "RESTART_GAP_HOURS",
    "stage_ranges", "assign_run_cycles", "previous_baseline", "source_hour",
    "cycle_baseline", "decline_hours", "window_status", "build_samples",
    "write_samples", "write_hourly", "mean", "value_range", "slope", "ratio", "finite",
]