filtering.py

"""Phase 3 — Filtering the pruned DFG.

Variant of Dijkstra that maximizes the *min-edge-frequency* path from source
to every node (forward) and from every node to sink (backward). The chosen
"best" incoming/outgoing edge of each node is retained unconditionally; any
remaining edge whose frequency is above the eta-percentile threshold is also
kept. This guarantees that every node sits on at least one source-to-sink
path while keeping the edge count low for low complexity.
"""

from __future__ import annotations

import math
from collections import deque
from dataclasses import dataclass

import numpy as np

from types_ import DFG, activities


@dataclass
class FilterResult:
    edges: set[tuple[str, str]]
    source: str
    sink: str


def _find_source_sink(dfg: DFG, nodes: set[str]) -> tuple[str, str]:
    """Source = no incoming arcs; sink = no outgoing arcs."""
    has_in: set[str] = {b for (_, b) in dfg.keys()}
    has_out: set[str] = {a for (a, _) in dfg.keys()}
    sources = [n for n in nodes if n not in has_in]
    sinks = [n for n in nodes if n not in has_out]
    if len(sources) != 1 or len(sinks) != 1:
        # Fall back to the sentinel labels installed during preprocessing.
        from dfg_builder import END, START

        if START in nodes and END in nodes:
            return START, END
        raise ValueError(
            f"Filtered PDFG must have exactly one source/sink; got "
            f"sources={sources}, sinks={sinks}"
        )
    return sources[0], sinks[0]


def _best_incoming(
        dfg: DFG, source: str, nodes: set[str]
) -> tuple[dict[str, float], dict[str, tuple[str, str]]]:
    """Forward BFS. Returns capacities and best incoming edge."""
    capacity: dict[str, float] = {n: 0 for n in nodes}
    capacity[source] = math.inf
    best: dict[str, tuple[str, str]] = {}

    out_adj: dict[str, list[tuple[str, int]]] = {n: [] for n in nodes}
    for (a, b), f in dfg.items():
        out_adj[a].append((b, f))

    in_queue: set[str] = {source}
    unexplored: set[str] = set(nodes) - {source}
    queue: deque[str] = deque([source])

    while queue:
        p = queue.popleft()
        in_queue.discard(p)
        for n, f_e in out_adj[p]:
            c_max = min(capacity[p], f_e)
            updated = False
            if c_max > capacity[n]:
                capacity[n] = c_max
                best[n] = (p, n)
                updated = True
            if updated:
                # Re-queue n so its successors are revisited with the new capacity.
                if n in unexplored:
                    unexplored.discard(n)
                if n not in in_queue:
                    queue.append(n)
                    in_queue.add(n)
            elif n in unexplored:
                unexplored.discard(n)
                if n not in in_queue:
                    queue.append(n)
                    in_queue.add(n)
    return capacity, best


def _best_outgoing(
        dfg: DFG, sink: str, nodes: set[str]
) -> tuple[dict[str, float], dict[str, tuple[str, str]]]:
    """Backward BFS. Returns capacities and best outgoing edge."""
    capacity: dict[str, float] = {n: 0 for n in nodes}
    capacity[sink] = math.inf
    best: dict[str, tuple[str, str]] = {}

    in_adj: dict[str, list[tuple[str, int]]] = {n: [] for n in nodes}
    for (a, b), f in dfg.items():
        in_adj[b].append((a, f))

    in_queue: set[str] = {sink}
    unexplored: set[str] = set(nodes) - {sink}
    queue: deque[str] = deque([sink])

    while queue:
        n = queue.popleft()
        in_queue.discard(n)
        for p, f_e in in_adj[n]:
            c_max = min(capacity[n], f_e)
            updated = False
            if c_max > capacity[p]:
                capacity[p] = c_max
                best[p] = (p, n)
                updated = True
            if updated:
                if p in unexplored:
                    unexplored.discard(p)
                if p not in in_queue:
                    queue.append(p)
                    in_queue.add(p)
            elif p in unexplored:
                unexplored.discard(p)
                if p not in in_queue:
                    queue.append(p)
                    in_queue.add(p)
    return capacity, best


def filter_pdfg(pdfg: DFG, eta: float) -> FilterResult:
    """Return the set of edges retained."""
    nodes = activities(pdfg)
    source, sink = _find_source_sink(pdfg, nodes)

    # Collect per-node best incoming / outgoing frequencies for the percentile.
    fmax_in: dict[str, int] = {n: 0 for n in nodes}
    fmax_out: dict[str, int] = {n: 0 for n in nodes}
    for (a, b), f in pdfg.items():
        if f > fmax_out[a]:
            fmax_out[a] = f
        if f > fmax_in[b]:
            fmax_in[b] = f

    frequencies: list[int] = []
    for n in nodes:
        if n != source:
            frequencies.append(fmax_in[n])
        if n != sink:
            frequencies.append(fmax_out[n])

    if frequencies:
        f_th = float(np.percentile(frequencies, eta * 100.0))
    else:
        f_th = 0.0

    _, best_in = _best_incoming(pdfg, source, nodes)
    _, best_out = _best_outgoing(pdfg, sink, nodes)

    kept_best: set[tuple[str, str]] = set(best_in.values()) | set(best_out.values())

    edges_out: set[tuple[str, str]] = set()
    for (a, b), f in pdfg.items():
        if (a, b) in kept_best or f > f_th:
            edges_out.add((a, b))

    return FilterResult(edges=edges_out, source=source, sink=sink)