Compressor/models/common/splitting.py

"""Event-group dataset splitting shared by both models.

Ported from the original generalized model pipeline so the held-out split is
identical for a fixed seed.
"""
from __future__ import annotations

import csv
import json
import random
import shutil
from collections import Counter, defaultdict
from pathlib import Path

SPLIT_NAMES = ("train", "validation", "test")


def normalize_splits(values: list[float]):
    total = sum(values)
    if total <= 0:
        raise ValueError("切分比例之和必须大于 0")
    return tuple((name, value / total) for name, value in zip(SPLIT_NAMES, values))


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


def allocate(groups: list[tuple], splits, assignments: dict, weight=None) -> None:
    if weight is None:
        weight = len
    weights = [weight(members) for _, members, _, _ in groups]
    total = sum(weights)
    targets = {name: total * ratio for name, ratio in splits}
    counts = Counter()
    for (group_id, members, _, _), group_weight in zip(groups, weights):
        split = min(
            (name for name, _ in splits),
            key=lambda name: (counts[name] - targets[name], counts[name], name),
        )
        assignments[group_id] = split
        counts[split] += group_weight


def assign_groups(rows: list[dict], splits, seed: int, positive_basis: str = "samples") -> dict[str, str]:
    grouped = defaultdict(list)
    for row in rows:
        grouped[row["event_group_id"]].append(row)
    positives = []
    normal_by_unit = defaultdict(list)
    for group_id, members in grouped.items():
        first = min(members, key=lambda row: row["start_hour"])
        item = (group_id, members, first["start_hour"], first["unit"])
        if first["label"] == "1":
            positives.append(item)
        else:
            normal_by_unit[first["unit"]].append(item)

    if len(positives) < 3:
        raise RuntimeError("独立异常事件少于3个，无法同时建立训练、验证和测试正例")

    assignments = {}
    rng = random.Random(seed)
    rng.shuffle(positives)
    positive_weight = (lambda members: 1) if positive_basis == "groups" else None
    allocate(positives, splits, assignments, weight=positive_weight)
    by_split = defaultdict(list)
    for group_id, members, _, _ in positives:
        by_split[assignments[group_id]].append((group_id, members))
    for name, _ in splits:
        if by_split[name]:
            continue
        donor = max((other for other, _ in splits), key=lambda other: len(by_split[other]))
        group_id, members = by_split[donor].pop()
        assignments[group_id] = name
        by_split[name].append((group_id, members))

    for unit in sorted(normal_by_unit):
        groups = normal_by_unit[unit]
        rng.shuffle(groups)
        allocate(groups, splits, assignments)
        if not any(assignments[group_id] == "train" for group_id, _, _, _ in groups):
            raise RuntimeError(f"{unit} 的正常事件组没有进入训练集，无法计算该机组基准")
    return assignments


def write_split(dataset: Path, splits, seed: int, positive_basis: str) -> dict:
    dataset = Path(dataset)
    manifest = load_manifest(dataset / "manifest.csv")
    candidates = [row for row in manifest if row["label"] in {"0", "1"}]
    assignments = assign_groups(candidates, splits, seed, positive_basis)
    for split, _ in splits:
        target = dataset / split
        if target.exists():
            shutil.rmtree(target)
        target.mkdir()

    for row in candidates:
        split = assignments[row["event_group_id"]]
        source = dataset / "candidates" / f"{row['sample_id']}.json"
        target = dataset / split / source.name
        target.hardlink_to(source)
        row["split"] = split

    fields = list(manifest[0])
    with (dataset / "manifest.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(manifest)

    summary = defaultdict(Counter)
    group_summary = defaultdict(Counter)
    group_first = {}
    for row in candidates:
        split = row["split"]
        summary[split][f"samples_{row['label_name']}"] += 1
        summary[split]["samples_total"] += 1
        group_first.setdefault(row["event_group_id"], row)
    for row in group_first.values():
        group_summary[row["split"]][f"groups_{row['label_name']}"] += 1
        group_summary[row["split"]]["groups_total"] += 1
    return {
        "dataset": dataset.name,
        "sample_counts": {split: dict(summary[split]) for split, _ in splits},
        "group_counts": {split: dict(group_summary[split]) for split, _ in splits},
        "excluded_samples": sum(row["label"] == "" for row in manifest),
        "target_ratios": {name: round(ratio, 4) for name, ratio in splits},
        "seed": seed,
        "positive_event_rule": "positive event groups are sampled globally with a fixed seed; every split keeps at least one independent positive event",
        "rule": "event_group_id is never split across train, validation, and test",
    }


__all__ = ["SPLIT_NAMES", "normalize_splits", "load_manifest", "allocate", "assign_groups", "write_split"]