node_mapping.py

import json
import re

_NODE_PATTERN = re.compile(r"^([mce])(\d)([a-z]+)(\d*)$")


def classify_node(node_id):
    match = _NODE_PATTERN.match(node_id)
    if match:
        building, floor, kind, num = match.groups()
        if floor == "1":
            return "floor1_outdoor"
        elif floor == "2":
            return "floor2"
        elif floor == "3":
            return "floor3"
        elif floor == "4":
            return "floor4"
        else:
            raise ValueError(f"Unrecognized floor digit in node_id: {node_id}")

    if re.match(r"^o\d+$", node_id) or node_id in ("ch", "sh", "Q1", "Q2"):
        return "floor1_outdoor"

    raise ValueError(f"Unrecognized node_id format: {node_id}")

def build_node_mapping(topology_map_path):
    with open(topology_map_path) as f:
        graph = json.load(f)

    nodes_by_head = {}
    for node_id in graph:
        head_name = classify_node(node_id)
        nodes_by_head.setdefault(head_name, []).append(node_id)

    mapping = {}
    for head_name, node_ids in nodes_by_head.items():
        for class_idx, node_id in enumerate(sorted(node_ids)):
            mapping[node_id] = (head_name, class_idx)

    return mapping

def head_class_counts(node_mapping):
    counts = {}
    for head_name, class_idx in node_mapping.values():
        counts[head_name] = max(counts.get(head_name, -1), class_idx) + 1
    return counts