Compressor/models/common/windowing.py

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

These helpers reproduce the original generalized model pipeline's dataset
building 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",
]