"""Stage-2: pin core horizontal beam + agg row via rigid soft-block translates.""" from __future__ import annotations from collections import deque from typing import Any from netx_api.mcp_remote.topology.layout_metrics import count_edge_crossings from netx_api.mcp_remote.topology.layout_ops.hotspots import fix_overlaps_local from netx_api.mcp_remote.topology.layout_ops.partition import partition_soft_blocks from netx_api.mcp_remote.topology.layout_ops.score import score_state from netx_api.mcp_remote.topology.layout_ops.state import LayoutParams, LayoutState, OpResult from netx_api.mcp_remote.topology.layout_ops.transforms import normalize_origin def _hops_to_cores( src: str, cores: list[str], adj: dict[str, set[str]] ) -> tuple[int, str]: if not cores: return 99, "" core_set = set(cores) if src in core_set: return 0, src seen = {src} q: deque[tuple[str, int]] = deque([(src, 0)]) while q: u, d = q.popleft() for v in adj.get(u, ()): if v in seen: continue if v in core_set: return d + 1, v seen.add(v) q.append((v, d + 1)) return 99, cores[0] def _core_order(state: LayoutState) -> list[str]: cores = [ n for n in state.positions if state.layers.get(n) == "core" ] if len(cores) < 2: # fall back: high-degree hubs that look like CN for n, nm in state.names.items(): if "-CN" in (nm or "") and n not in cores: cores.append(n) cores = sorted( set(cores), key=lambda n: (-len(state.adj.get(n, ())), state.names.get(n, n)), ) return cores def _agg_hubs(state: LayoutState, *, min_degree: int = 4) -> list[str]: aggs = [ n for n in state.positions if state.layers.get(n) == "agg" and len(state.adj.get(n, ())) >= min_degree ] return sorted( aggs, key=lambda n: (-len(state.adj.get(n, ())), state.names.get(n, n)), ) def _apply_targets( state: LayoutState, targets: dict[str, tuple[float, float]], by_hub: dict[str, Any], ) -> LayoutState: pos = dict(state.positions) for hid, (tx, ty) in targets.items(): block = by_hub.get(hid) if not block or hid not in pos: if hid in pos: pos[hid] = (tx, ty) continue hx, hy = pos[hid] dx, dy = tx - hx, ty - hy for n in block.node_ids: if n in pos: x, y = pos[n] pos[n] = (x + dx, y + dy) st = state.copy() st.positions = pos st = fix_overlaps_local(st, LayoutParams()).state st = normalize_origin(st, LayoutParams()).state return st def pin_beam_rigid( state: LayoutState, params: LayoutParams | None = None, *, step: float | None = None, gap_scales: tuple[float, ...] = (1.0, 1.3, 1.6, 0.85), ) -> OpResult: """Rigid-translate soft blocks so cores form a beam and aggs sit on a row. Picks the trial with lowest (overlaps, crossings). Does not worsen crossings vs the input when a better trial exists; otherwise returns input unchanged. """ params = params or LayoutParams() step_px = float(step if step is not None else max(params.pitch, 180.0)) cores = _core_order(state) if len(cores) < 2: return OpResult( state=state, moved=set(), op="pin_beam", params={"cores": cores}, note="need_ge_2_cores", ) aggs = _agg_hubs(state) aggs = sorted( aggs, key=lambda a: ( _hops_to_cores(a, cores, state.adj)[0], state.names.get(a, a), ), ) blocks = partition_soft_blocks(state, mode="hub_territory") by_hub = {b.hub_id: b for b in blocks if b.hub_id} links = list(state.links) base_cross = count_edge_crossings(state.positions, links) base_score = score_state(state) base_ov = int(base_score.get("footprint_overlap_pairs") or 0) + int( base_score.get("label_overlap_pairs") or 0 ) gap0 = step_px * 2.4 beam_y = 0.0 trials: list[tuple[str, LayoutState, int, int]] = [] for gs in gap_scales: gap = gap0 * gs for petal_mul in (4.0, 4.5, 5.5): petal_dy = step_px * petal_mul targets: dict[str, tuple[float, float]] = {} for i, c in enumerate(cores): targets[c] = (i * gap, beam_y) if aggs: lo = targets[cores[0]][0] - gap * 0.5 hi = targets[cores[-1]][0] + gap * 0.5 for i, a in enumerate(aggs): _, pref = _hops_to_cores(a, cores, state.adj) prefer_x = targets.get(pref, targets[cores[0]])[0] if len(aggs) == 1: x = prefer_x else: x = lo + (hi - lo) * i / (len(aggs) - 1) x = 0.55 * x + 0.45 * prefer_x targets[a] = (x, beam_y + petal_dy) st2 = _apply_targets(state, targets, by_hub) sc = score_state(st2) ov = int(sc.get("footprint_overlap_pairs") or 0) + int( sc.get("label_overlap_pairs") or 0 ) cross = int(sc.get("edge_crossings") or 0) trials.append((f"g{gs}_dy{petal_dy:.0f}", st2, ov, cross)) # Y-align only (preserve x): often safer on already-readable canvases cy = sum(state.positions[c][1] for c in cores if c in state.positions) / len( cores ) targets_y: dict[str, tuple[float, float]] = {} for c in cores: if c in state.positions: targets_y[c] = (state.positions[c][0], cy) ay = cy + step_px * 4.0 for a in aggs: if a in state.positions: targets_y[a] = (state.positions[a][0], ay) if targets_y: st_y = _apply_targets(state, targets_y, by_hub) sc = score_state(st_y) ov = int(sc.get("footprint_overlap_pairs") or 0) + int( sc.get("label_overlap_pairs") or 0 ) cross = int(sc.get("edge_crossings") or 0) trials.append(("y_align", st_y, ov, cross)) if not trials: return OpResult( state=state, moved=set(), op="pin_beam", note="no_trials" ) best_name, best_st, best_ov, best_cross = min( trials, key=lambda t: (t[2], t[3]) ) # Refuse internal accept if worse than input on both axes we care about if best_ov > base_ov or best_cross > base_cross + max(20, int(base_cross * 0.05)): return OpResult( state=state, moved=set(), op="pin_beam", params={ "cores": cores, "aggs": aggs, "best_trial": best_name, "best_crossings": best_cross, "base_crossings": base_cross, "accepted": False, }, note="no_improvement", ) moved = { n for n, p in best_st.positions.items() if n in state.positions and p != state.positions[n] } best_st.meta = dict(best_st.meta or {}) best_st.meta["pin_beam"] = { "trial": best_name, "cores": cores, "aggs": aggs, "crossings_before": base_cross, "crossings_after": best_cross, "overlaps_after": best_ov, } return OpResult( state=best_st, moved=moved, op="pin_beam", params={ "cores": cores, "aggs": aggs, "best_trial": best_name, "accepted": True, "crossings_before": base_cross, "crossings_after": best_cross, }, note=f"pin_beam:{best_name} {base_cross}->{best_cross}", ) def pin_beam_params_from_overrides(overrides: dict[str, Any] | None) -> dict[str, Any]: out: dict[str, Any] = {} if not overrides: return out if overrides.get("step") is not None: try: out["step"] = float(overrides["step"]) except (TypeError, ValueError): pass return out