tapas.py

import networkx as nx
import time
import random
from typing import Dict, List, Tuple, Callable, Optional, Any, Set
from dataclasses import dataclass, field
from collections import defaultdict
import logging
import math
import numpy as np
import numba

# --- Configuration ---
EPSILON = 1e-9
MIN_SHIFT = 1e-14  # minimum shift execution threshold (lower than EPSILON to avoid stalls)
MAX_MCS_DEPTH = 1000
MAX_PATH_NODES = 10000
MAX_REFINE_SUB_ITER = 5
MAX_NEW_PAS_PER_ITER = 10000
MU_PAS_MATCH = 0.5        # cost effective factor μ (Procedure 1, Step 1.3.1 condition 1)
NU_FLOW_EFFECTIVE = 0.25  # flow effective factor ν (Procedure 1, Step 1.3.1 condition 2)

# --- Logging Setup ---
logger = logging.getLogger(__name__)
if not logger.hasHandlers():
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s - %(levelname)s - %(message)s')

# --- Data Structures ---


@dataclass(frozen=True, eq=True)
class PAS:
    """
    Paired Alternative Segments (PAS) structure.
    Stores paths as tuples of original node pairs.
    Origin is stored in its original type.
    Includes cost difference for potential sorting.
    """
    e1: Tuple[Tuple[Any, Any], ...] = field(
        default_factory=tuple)  # Shortest path segment edges
    e2: Tuple[Tuple[Any, Any], ...] = field(
        default_factory=tuple)  # Alternative path segment edges
    o: Any = None  # Origin node (original type)
    conn_node: Any = None  # Node where MCS connected to SPT (original type)
    cost_diff: float = 0.0  # Store initial cost difference |cost(e2)-cost(e1)|

    def __post_init__(self):
        if not isinstance(self.e1, tuple):
            object.__setattr__(self, 'e1', tuple(self.e1))
        if not isinstance(self.e2, tuple):
            object.__setattr__(self, 'e2', tuple(self.e2))

    def __str__(self) -> str:
        o_str = str(self.o)
        conn_str = str(self.conn_node)
        e1_len = len(self.e1)
        e2_len = len(self.e2)
        return f"PAS(o={o_str}, conn={conn_str}, |e1|={e1_len}, |e2|={e2_len}, diff={self.cost_diff:.2f})"

# --- Numba Helper Functions ---


@numba.jit(nopython=True)
def _calculate_path_cost_numba(costs: np.ndarray) -> float:
    """Numba-optimized function to sum costs."""
    cost_sum = 0.0
    for c in costs:
        if c == np.inf:
            return np.inf
        cost_sum += c
    return cost_sum


@numba.jit(nopython=True)
def _calculate_path_derivative_numba(derivs: np.ndarray) -> float:
    """Numba-optimized function to sum derivatives."""
    deriv_sum = 0.0
    for d in derivs:
        if d == np.inf:  # Should ideally not happen for derivatives, but check
            return np.inf
        deriv_sum += d
    return deriv_sum

# --- Core Algorithm Class ---


class TAPASAlgorithm:
    """
    Implements the TAPAS (Traffic Assignment by Paired Alternative Segments) algorithm.
    Correctly uses the path dictionary output from nx.single_source_dijkstra.
    Includes node validation during initialization.
    Uses standardized node IDs for internal dictionary lookups and set operations.
    Corrected flow shift logic. Limited PAS addition and multi-pass refinement.
    Refined PAS refinement logic for better convergence.
    """

    def __init__(self,
                 graph: nx.DiGraph,
                 demand: Dict[Tuple[Any, Any], float],
                 from_cent: Dict[Any, Any],
                 to_cent: Dict[Any, Any],
                 cost_function: Callable[[Dict, float], Tuple[float, float]],
                 max_iter: int = 100,
                 accuracy: float = 0.001,
                 verbose: bool = False,
                 proportions: bool = False):
        """Initializes TAPAS and validates graph nodes."""
        if not isinstance(graph, nx.DiGraph):
            raise TypeError("Input graph must be a NetworkX DiGraph.")
        if not callable(cost_function):
            raise TypeError("cost_function must be callable.")

        self.graph = graph
        self.demand = demand
        self.from_cent = from_cent
        self.to_cent = to_cent
        self.cost_function = cost_function
        self.max_iter = max_iter
        self.accuracy = accuracy
        self.verbose = verbose
        self.proportions = proportions
        self.pas_set: Set[PAS] = set()
        # Index: head node (j, shared endpoint of e1/e2) → set of PASs ending there.
        # Avoids O(|PAS|) linear scan during PAS match (Step 1.3.1).
        self._pas_by_head: Dict[Any, Set[PAS]] = defaultdict(set)
        self.tstt: float = float('inf')
        self.sptt: float = float('inf')
        self.gap: float = float('inf')
        self.iterations: int = 0
        self.sp_edges: Dict[Tuple[Any, Any], Tuple[Tuple[Any, Any], ...]] = {}
        self.eodtt: Dict[Tuple[Any, Any], float] = {}
        self.entropy: float = 0.0
        self.prop_gap: float = 0.0

        if verbose:
            logger.setLevel(logging.INFO)
        else:
            logger.setLevel(logging.WARNING)

    # --- Initialization ---
    def _initialize_graph_costs(self):
        """Calculates initial costs and derivatives based on current flows."""
        for u, v, data in self.graph.edges(data=True):
            try:
                flow = max(0.0, data.get('flow', 0.0))
                data['cost'], data['dcost'] = self.cost_function(data, flow)
                if data['cost'] == float('inf') or math.isnan(data['cost']):
                    data['cost'] = float('inf')
                    data['dcost'] = 0.0
            except Exception as e:
                logger.error(f"Error calculating cost for edge ({u}, {v}): {e}",
                             exc_info=self.verbose)
                data['cost'], data['dcost'] = float('inf'), 0.0

    def _run_all_or_nothing(self) -> None:
        """Performs an initial All-or-Nothing assignment, updating graph flows."""
        for u, v, data in self.graph.edges(data=True):
            data['flow'] = 0.0
            data['obflow'] = defaultdict(float)
            if self.proportions:
                data['od_flow'] = {}  # Track exact OD flows
        self._initialize_graph_costs()

        for origin_zone, origin_node in self.from_cent.items():
            if origin_node not in self.graph:
                continue
            try:
                cost_dict, path_dict = nx.single_source_dijkstra(
                    self.graph, origin_node, weight='cost')
                for dest_zone, dest_node in self.to_cent.items():
                    od_pair = (origin_zone, dest_zone)
                    demand_val = self.demand.get(od_pair, 0.0)
                    if demand_val > EPSILON:
                        path_nodes = path_dict.get(dest_node)
                        if path_nodes and len(path_nodes) >= 2:
                            path_edges = tuple(
                                zip(path_nodes[:-1], path_nodes[1:]))
                            for u, v in path_edges:
                                try:
                                    edge_data = self.graph.edges[u, v]
                                    edge_data['flow'] += demand_val
                                    edge_data['obflow'][origin_node] += demand_val
                                    
                                    # Track exact OD flows if proportions enabled
                                    if self.proportions:
                                        if od_pair not in edge_data['od_flow']:
                                            edge_data['od_flow'][od_pair] = 0.0
                                        edge_data['od_flow'][od_pair] += demand_val
                                        
                                except KeyError:
                                    logger.error(
                                        f"AON Error: Edge ({u}, {v}) from SP path not found.")
                        elif dest_node in cost_dict and cost_dict[dest_node] != float('inf') and origin_node != dest_node:
                            logger.warning(
                                f"AON Warning: Dest {dest_node} reachable for OD {od_pair} but no path list found or path too short.")
            except nx.NodeNotFound as e:
                logger.warning(
                    f"AON Error: Node not found during Dijkstra for origin {origin_node}: {e}")
            except Exception as e:
                logger.exception(
                    f"AON Error during Dijkstra/path assignment for origin {origin_node}: {e}")

        for u, v, data in self.graph.edges(data=True):
            data['flow'] = sum(data.get('obflow', {}).values())

    # --- Shortest Path and Gap Calculation ---
    def _calculate_sp_and_gap(self) -> Tuple[float, float, float, Dict, Dict]:
        """Calculates SPT for all origins, TSTT, SPTT, and Gap. Returns path_dict."""
        logger.debug("Calculating shortest paths and gap...")
        self._initialize_graph_costs()
        tstt = sum(max(0.0, data.get('flow', 0.0)) * data['cost'] for _, _, data in self.graph.edges(
            data=True) if data.get('cost', float('inf')) != float('inf'))
        sptt = 0.0
        all_sp_paths: Dict[Any, Dict[Any, List[Any]]] = {}
        all_sp_cost: Dict[Any, Dict[Any, float]] = {}
        for origin_zone, origin_node in self.from_cent.items():
            if origin_node not in self.graph:
                continue
            try:
                cost_dict_std_keys, path_dict_std_keys = nx.single_source_dijkstra(
                    self.graph, origin_node, weight='cost')
                all_sp_paths[origin_node] = path_dict_std_keys
                all_sp_cost[origin_node] = cost_dict_std_keys
                for dest_zone, dest_node in self.to_cent.items():
                    od_pair = (origin_zone, dest_zone)
                    demand_val = self.demand.get(od_pair, 0.0)
                    if demand_val > EPSILON:
                        path_cost = cost_dict_std_keys.get(
                            dest_node, float('inf'))
                        if path_cost != float('inf'):
                            sptt += demand_val * path_cost
            except nx.NodeNotFound as e:
                logger.warning(
                    f"SP Error: Node not found during Dijkstra for origin {origin_node}: {e}")
            except Exception as e:
                logger.exception(
                    f"SP Error during Dijkstra for origin {origin_node}: {e}")
                sptt = EPSILON
        gap = max(0.0, (tstt / sptt) - 1.0) if sptt > EPSILON else (float('inf')
                                                                    if tstt > EPSILON else 0.0)
        logger.debug(f"TSTT={tstt:.4f}, SPTT={sptt:.4f}, Gap={gap:.6f}")
        return tstt, sptt, gap, all_sp_paths, all_sp_cost

    # --- PAS Finding ---
    def _is_potential_edge(self, edge: Tuple[Any, Any], origin_node_orig: Any, sp_cost_dict: Dict[Any, float]) -> bool:
        u_std, v_std = edge
        try:
            edge_data = self.graph.edges[u_std, v_std]
            origin_flow = edge_data.get(
                'obflow', {}).get(origin_node_orig, 0.0)
            if origin_flow <= EPSILON:
                return False
            cost_u = sp_cost_dict.get(u_std, float('inf'))
            cost_v = sp_cost_dict.get(v_std, float('inf'))
            edge_cost = edge_data.get('cost', float('inf'))
            if cost_u == float('inf') or edge_cost == float('inf'):
                return False
            improvement_potential = cost_u + edge_cost - cost_v
            return improvement_potential > EPSILON
        except KeyError:
            return False
        except Exception as e:
            logger.error(
                f"Error checking potential edge {edge} for origin {origin_node_orig}: {e}", exc_info=self.verbose)
            return False

    def _extract_mcs_path_nodes(self, mcs_pred_dict: Dict[Any, Optional[Any]], start_node: Any, end_node: Any) -> List[Any]:
        """Helper to extract path nodes from MCS predecessors. Uses standardized keys for lookup."""
        path = []
        curr_orig = end_node
        nodes_visited_std = set()

        std_start_node = start_node

        while curr_orig is not None:
            try:
                std_curr = curr_orig
            except TypeError:
                logger.error(
                    f"MCS Path Error: Node {curr_orig} became unhashable.")
                return []
            if std_curr in nodes_visited_std or len(path) > MAX_PATH_NODES:
                logger.error(
                    f"MCS Path Error: Cycle or max length detected from {start_node} to {end_node} at {curr_orig}.")
                return []
            nodes_visited_std.add(std_curr)
            path.append(curr_orig)
            if std_curr == std_start_node:
                break
            try:
                pred_orig = mcs_pred_dict.get(std_curr)
                if pred_orig is None:
                    if std_curr not in mcs_pred_dict:
                        logger.error(
                            f"MCS Path Error: Standardized node {std_curr} (from {curr_orig}) not found in mcs_pred_dict.")
                    else:
                        logger.error(
                            f"MCS Path Error: Path broken from {start_node} to {end_node}. Predecessor for {curr_orig} (std: {std_curr}) is None.")
                    return []
            except TypeError:
                logger.error(
                    f"FATAL: Standardized node {std_curr} (from {curr_orig}) is unhashable key.")
                return []
            curr_orig = pred_orig
        if not path or path[-1] != std_start_node:
            logger.error(
                f"MCS Path Error: Failed to trace path back to {start_node} from {end_node}. Final path: {path}")
            return []
        return list(reversed(path))

    def _find_pas_mfs(self,
                      potential_edge: Tuple[Any, Any],
                      origin_node_orig: Any,
                      sp_path_dict: Dict[Any, List[Any]]
                      ) -> Optional[PAS]:
        """
        Finds a PAS using MFS (Maximum Flow First Search) - Procedure 4 from Xie & Xie (2016).

        Key differences from BFS/MCS:
        - Selects the backward link by maximum origin-based flow (not max cost).
        - Uses a three-state node label: 0 = unvisited, -1 = on SPT path j→r, 1 = visited in Step 2.
        - Tracks Q[k] (foregoing node of k in e2) to reconstruct the higher-cost segment.
        - Detects and removes directed cycle flow before retrying (Step 2.3).
        - Shifts flow immediately and retries if the potential link is not yet equilibrated (Step 2.2).
        """
        i_node, j_node = potential_edge

        # SPT path from origin to j: [origin, ..., j]
        spt_path_to_j = sp_path_dict.get(j_node)
        if not spt_path_to_j or len(spt_path_to_j) < 2:
            return None

        # Allow up to 3 retries (Step 2.2 → Step 1 feedback loop)
        MAX_MFS_RETRIES = 3

        for _attempt in range(MAX_MFS_RETRIES):
            # l[k]: node scan status  0=unvisited, -1=on SPT j→r, 1=visited in Step 2
            l: Dict[Any, int] = {}
            # Q[k]: foregoing node of k in the higher-cost segment e2
            Q: Dict[Any, Any] = {}

            # Step 1: mark all nodes on the SPT path from j back to origin as -1
            for node in spt_path_to_j:
                l[node] = -1

            # Fast-path: if i_node is already on the SPT (labeled -1), the correct
            # PAS is immediate — p̃ = i_node, e1 = SPT[i:j], e2 = [(i,j)].
            # Running MFS from i_node would pick its SPT predecessor as the connection
            # node, creating segments that share the (pred→i) edge (violates disjointness).
            if l.get(i_node, 0) == -1:
                try:
                    idx = spt_path_to_j.index(i_node)
                except ValueError:
                    break
                e1_nodes = spt_path_to_j[idx:]   # [i_node, ..., j_node]
                if len(e1_nodes) < 2:
                    break
                e1_edges = tuple(zip(e1_nodes[:-1], e1_nodes[1:]))
                e2_edges = ((i_node, j_node),)

                cost_e1 = self._calculate_path_cost(e1_edges)
                cost_e2 = self._calculate_path_cost(e2_edges)
                if cost_e1 == float('inf') or cost_e2 == float('inf'):
                    break
                if abs(cost_e1 - cost_e2) < EPSILON:
                    return None  # already equilibrated

                pas = PAS(e1=e1_edges, e2=e2_edges, o=origin_node_orig,
                          conn_node=i_node, cost_diff=abs(cost_e2 - cost_e1))
                shift = self._calculate_shift_value(pas)
                if abs(shift) > MIN_SHIFT:
                    self._shift_flow(pas, shift)
                    self._update_affected_edge_costs(set(e1_edges) | set(e2_edges))
                return pas

            n = i_node  # next scanning node

            depth = 0
            while depth < MAX_MCS_DEPTH:
                depth += 1

                # Step 2: m̂ = argmax{ x^r_{mn} : m ∈ I(n) } (max origin-based flow)
                best_m: Optional[Any] = None
                best_flow = -1.0
                try:
                    for pred in self.graph.predecessors(n):
                        try:
                            flow = self.graph.edges[pred, n].get(
                                'obflow', {}).get(origin_node_orig, 0.0)
                            if flow > 1e-12 and flow > best_flow:
                                best_flow = flow
                                best_m = pred
                        except KeyError:
                            continue
                except nx.NetworkXError:
                    break

                if best_m is None:
                    break  # no usable predecessor; cannot form a PAS

                m_hat = best_m
                Q[m_hat] = n          # set Q_{m̂} = n
                l_mhat = l.get(m_hat, 0)

                if l_mhat == -1:
                    # ── Step 2.1: m̂ is on the SPT path → connection node found ──
                    p_tilde = m_hat   # connection node (shared start of e1 and e2)

                    # e1: SPT segment from p̃_t to j  (spt_path_to_j already in order origin→j)
                    try:
                        idx = spt_path_to_j.index(p_tilde)
                    except ValueError:
                        break
                    e1_nodes = spt_path_to_j[idx:]   # [p̃_t, ..., j]

                    if len(e1_nodes) < 2:
                        # p_tilde == j_node: the backward search hit the head of the
                        # potential link itself.  This forms a cycle:
                        #   j_node →(Q chain)→ i_node →(potential link)→ j_node
                        # Reconstruct and remove the minimum cycle flow, then restart.
                        cyc_nodes: List[Any] = [j_node]
                        cyc_cur = j_node
                        cyc_seen: Set[Any] = {j_node}
                        reached_cyc_i = False
                        while True:
                            cyc_nxt = Q.get(cyc_cur)
                            if cyc_nxt is None or cyc_nxt in cyc_seen:
                                break
                            cyc_nodes.append(cyc_nxt)
                            cyc_seen.add(cyc_nxt)
                            if cyc_nxt == i_node:
                                reached_cyc_i = True
                                break
                            cyc_cur = cyc_nxt

                        if reached_cyc_i:
                            # Close cycle with the potential link i_node → j_node
                            cyc_edges = list(zip(cyc_nodes[:-1], cyc_nodes[1:]))
                            cyc_edges.append((i_node, j_node))
                            delta = min(
                                (self.graph.edges[u, v].get('obflow', {}).get(origin_node_orig, 0.0)
                                 for u, v in cyc_edges if self.graph.has_edge(u, v)),
                                default=0.0
                            )
                            if delta > MIN_SHIFT:
                                affected: Set[Tuple[Any, Any]] = set()
                                for u, v in cyc_edges:
                                    if self.graph.has_edge(u, v):
                                        data = self.graph.edges[u, v]
                                        data.setdefault('obflow', defaultdict(float))
                                        data['obflow'][origin_node_orig] = max(
                                            0.0, data['obflow'].get(origin_node_orig, 0.0) - delta)
                                        affected.add((u, v))
                                for u, v in affected:
                                    d = self.graph.edges[u, v]
                                    d['flow'] = sum(d.get('obflow', {}).values())
                                self._update_affected_edge_costs(affected)
                        break  # restart with fresh labels

                    e1_edges = tuple(zip(e1_nodes[:-1], e1_nodes[1:]))

                    # e2: higher-cost segment traced via Q pointers from p̃_t to i, then potential link
                    # Q[p̃_t]=next → Q[next]=... → i → j
                    e2_nodes: List[Any] = [p_tilde]
                    curr = p_tilde
                    visited_e2: Set[Any] = {p_tilde}
                    reached_i = False
                    while True:
                        nxt = Q.get(curr)
                        if nxt is None or nxt in visited_e2:
                            break
                        e2_nodes.append(nxt)
                        visited_e2.add(nxt)
                        if nxt == i_node:
                            reached_i = True
                            break
                        curr = nxt

                    if not reached_i:
                        break
                    e2_nodes.append(j_node)           # potential link i → j
                    e2_edges = tuple(zip(e2_nodes[:-1], e2_nodes[1:]))

                    cost_e1 = self._calculate_path_cost(e1_edges)
                    cost_e2 = self._calculate_path_cost(e2_edges)
                    if cost_e1 == float('inf') or cost_e2 == float('inf'):
                        break

                    cost_diff = abs(cost_e2 - cost_e1)
                    pas = PAS(e1=e1_edges, e2=e2_edges, o=origin_node_orig,
                              conn_node=p_tilde, cost_diff=cost_diff)

                    # ── Step 2.2: shift flow immediately ("shift-first") ──
                    shift = self._calculate_shift_value(pas)
                    if abs(shift) > MIN_SHIFT:
                        self._shift_flow(pas, shift)
                        self._update_affected_edge_costs(set(e1_edges) | set(e2_edges))
                        # If origin flow on (i,j) dropped to 0 the link is equilibrated → done
                        try:
                            ij_flow = self.graph.edges[i_node, j_node].get(
                                'obflow', {}).get(origin_node_orig, 0.0)
                        except KeyError:
                            ij_flow = 0.0
                        if ij_flow <= EPSILON:
                            return pas   # potential link satisfied
                        # else retry (Step 2.2 → Step 1): outer loop increments _attempt
                        break
                    else:
                        return pas       # already equilibrated, no shift needed

                elif l_mhat == 1:
                    # ── Step 2.3: directed cycle detected; remove minimum cycle flow ──
                    # Walk the Q chain from m̂ to find where the cycle actually closes.
                    # Nodes before the closure point are a "tail" — only edges strictly
                    # inside the true cycle should have flow removed.
                    walk: List[Any] = [m_hat]
                    seen_walk: Dict[Any, int] = {m_hat: 0}
                    cur_w = m_hat
                    cycle_start_idx: Optional[int] = None
                    while True:
                        nxt_w = Q.get(cur_w)
                        if nxt_w is None:
                            break
                        if nxt_w in seen_walk:
                            cycle_start_idx = seen_walk[nxt_w]
                            break
                        seen_walk[nxt_w] = len(walk)
                        walk.append(nxt_w)
                        cur_w = nxt_w

                    if cycle_start_idx is not None:
                        # Strict cycle: walk[cycle_start_idx:] + closing edge back to entry
                        cycle_nodes_s = walk[cycle_start_idx:]
                        cycle_edges_s: List[Tuple[Any, Any]] = list(
                            zip(cycle_nodes_s[:-1], cycle_nodes_s[1:]))
                        cycle_edges_s.append((cycle_nodes_s[-1], cycle_nodes_s[0]))

                        delta = min(
                            (self.graph.edges[u, v].get('obflow', {}).get(origin_node_orig, 0.0)
                             for u, v in cycle_edges_s if self.graph.has_edge(u, v)),
                            default=0.0
                        )
                        if delta > MIN_SHIFT:
                            affected_c: Set[Tuple[Any, Any]] = set()
                            for u, v in cycle_edges_s:
                                if self.graph.has_edge(u, v):
                                    data = self.graph.edges[u, v]
                                    data.setdefault('obflow', defaultdict(float))
                                    data['obflow'][origin_node_orig] = max(
                                        0.0, data['obflow'].get(origin_node_orig, 0.0) - delta)
                                    affected_c.add((u, v))
                            for u, v in affected_c:
                                d = self.graph.edges[u, v]
                                d['flow'] = sum(d.get('obflow', {}).values())
                            self._update_affected_edge_costs(affected_c)

                    # Go to Step 1: restart with fresh labels
                    break  # exits inner while → outer _attempt loop resets l and Q

                else:
                    # l_mhat == 0: unvisited; continue searching backward
                    l[m_hat] = 1
                    n = m_hat

        return None

    # --- PAS set index helpers ---

    def _pas_head(self, pas: PAS) -> Optional[Any]:
        """Returns the shared head node (j) of a PAS — last node of e2."""
        return pas.e2[-1][1] if pas.e2 else None

    def _add_to_pas_set(self, pas: PAS) -> bool:
        """Adds a PAS to pas_set and the head-node index. Returns True if newly added."""
        if pas in self.pas_set:
            return False
        self.pas_set.add(pas)
        h = self._pas_head(pas)
        if h is not None:
            self._pas_by_head[h].add(pas)
        return True

    def _remove_from_pas_set(self, pas: PAS) -> None:
        """Removes a PAS from pas_set and the head-node index."""
        self.pas_set.discard(pas)
        h = self._pas_head(pas)
        if h is not None:
            self._pas_by_head[h].discard(pas)

    # --- PAS match (Step 1.3.1 / iTAPAS Step 1.2.1) ---

    def _try_pas_match(self,
                       potential_edge: Tuple[Any, Any],
                       origin_node_orig: Any,
                       sp_cost_dict: Dict[Any, float]) -> bool:
        """
        PAS match step (Procedure 1, Step 1.3.1 / iTAPAS Step 1.2.1).

        Before running the expensive MFS backward search for a potential link (i, j),
        checks whether any existing PAS already covers (i, j) in its e2 segment and
        satisfies the match conditions. If found, shifts flow for the PAS's origin r_0
        and returns True so the caller can skip MFS for this link.

        Conditions checked (Xie & Xie 2016, Step 1.3.1):
          (1)  t_e2 - t_e1 ≥ μ · π^r_{ij}          (cost effective factor μ = 0.5)
          (2)  f^{r_0}_{e2} ≥ ν · x^r_{ij}           (flow effective factor ν = 0.25)
          (3)  (i, j) ∈ e2  — enforced by head-node index (j is the last node of e2)
               and confirmed by checking e2[-1] == (i, j)
        """
        i_node, j_node = potential_edge

        # Candidates: only PASs whose head node equals j  (condition 3, O(1) lookup)
        candidates = self._pas_by_head.get(j_node)
        if not candidates:
            return False

        try:
            ij_data = self.graph.edges[i_node, j_node]
        except KeyError:
            return False

        # π^r_{ij} = u^r_i + t_{ij} − u^r_j  (reduced cost of the potential link for origin r)
        sp_cost_i = sp_cost_dict.get(i_node, float('inf'))
        sp_cost_j = sp_cost_dict.get(j_node, float('inf'))
        t_ij = ij_data.get('cost', float('inf'))
        if sp_cost_i == float('inf') or t_ij == float('inf'):
            return False
        reduced_cost = sp_cost_i + t_ij - sp_cost_j   # π^r_{ij}

        # x^r_{ij}: origin-r flow on the potential link
        x_r_ij = ij_data.get('obflow', {}).get(origin_node_orig, 0.0)

        for pas in list(candidates):
            if pas not in self.pas_set or not pas.e2:
                continue

            # Confirm (i, j) is the last edge of e2 (not merely the head node)
            if pas.e2[-1] != potential_edge:
                continue

            cost_e1 = self._calculate_path_cost(pas.e1)
            cost_e2 = self._calculate_path_cost(pas.e2)
            if cost_e1 == float('inf') or cost_e2 == float('inf'):
                continue

            # Condition (1): t_e2 − t_e1 ≥ μ · π^r_{ij}
            if (cost_e2 - cost_e1) < MU_PAS_MATCH * reduced_cost:
                continue

            # Condition (2): f^{r_0}_{e2} ≥ ν · x^r_{ij}
            # Uses the PAS's own origin r_0 (iTAPAS: one origin per PAS).
            f_r0_e2 = self._get_min_origin_flow(pas.e2, pas.o)
            if f_r0_e2 < NU_FLOW_EFFECTIVE * x_r_ij:
                continue

            # Match found — shift flow for r_0 (iTAPAS: equalise cost for the PAS's origin)
            shift = self._calculate_shift_value(pas)
            if abs(shift) > MIN_SHIFT:
                self._shift_flow(pas, shift)
                self._update_affected_edge_costs(set(pas.e1) | set(pas.e2))
                return True  # shift made → skip MFS for this potential link
            # shift ≈ 0: PAS already equilibrated; don't block MFS from finding a better one

        return False

    def _identify_potential_edges_and_find_pas(self, all_sp_paths, all_sp_cost) -> int:
        """Iterates through edges and origins to find potential edges and new PAS. Limits PAS added."""
        logger.debug("Identifying potential edges and finding PAS...")
        potential_edges_checked = 0
        pas_candidates_this_iter: List[PAS] = []
        for origin_node_orig, sp_cost_dict_std_keys in all_sp_cost.items():
            sp_path_dict = all_sp_paths.get(origin_node_orig)
            if sp_path_dict is None:
                continue
            sp_tree_edges = set()
            for path_list in sp_path_dict.values():
                if len(path_list) >= 2:
                    for i in range(len(path_list) - 1):
                        sp_tree_edges.add((path_list[i], path_list[i+1]))
            for u, v, _ in self.graph.edges(data=True):
                edge = (u, v)
                if edge in sp_tree_edges:
                    continue
                potential_edges_checked += 1
                if not self._is_potential_edge(edge, origin_node_orig, sp_cost_dict_std_keys):
                    continue
                # Step 1.3.1 PAS match: reuse an existing PAS if one covers this link
                if self._try_pas_match(edge, origin_node_orig, sp_cost_dict_std_keys):
                    continue   # matched and shifted — skip MFS for this link
                # Step 1.3.2 PAS identification: run MFS backward search
                pas = self._find_pas_mfs(edge, origin_node_orig, sp_path_dict)
                if pas and pas not in self.pas_set:
                    pas_candidates_this_iter.append(pas)

        new_pas_count_total = len(pas_candidates_this_iter)
        if new_pas_count_total > MAX_NEW_PAS_PER_ITER:
            logger.debug(
                f"Found {new_pas_count_total} PAS candidates, limiting to {MAX_NEW_PAS_PER_ITER}.")
            pas_candidates_this_iter.sort(
                key=lambda p: p.cost_diff, reverse=True)
            pas_to_add = pas_candidates_this_iter[:MAX_NEW_PAS_PER_ITER]
        else:
            pas_to_add = pas_candidates_this_iter
        newly_added_count = sum(self._add_to_pas_set(p) for p in pas_to_add)
        logger.debug(
            f"Checked {potential_edges_checked} potential edges, found {new_pas_count_total} candidates, added {newly_added_count} unique. Total PAS: {len(self.pas_set)}")
        return newly_added_count

    # --- PAS Refinement and Flow Shifting ---

    def _calculate_path_cost(self, path: Tuple[Tuple[Any, Any], ...]) -> float:
        """Calculates the total cost of a path using Numba for summation."""
        if not path:
            return 0.0
        try:
            # Extract costs into a NumPy array
            costs = np.array([self.graph.edges[u, v].get('cost', np.inf)
                              for u, v in path], dtype=np.float64)
            # Call the Numba-optimized function
            return _calculate_path_cost_numba(costs)
        except KeyError as e:
            logger.warning(f"Edge {e} not found during path cost calculation for path {path}. Returning inf.")
            return float('inf')
        except Exception as e:
            logger.error(
                f"Unexpected error in _calculate_path_cost for path {path}: {e}", exc_info=self.verbose)
            return float('inf')

    def _calculate_path_derivative(self, path: Tuple[Tuple[Any, Any], ...]) -> float:
        """Calculates the total derivative of cost along a path using Numba."""
        if not path:
            return 0.0
        try:
            # Extract derivatives into a NumPy array
            derivs = np.array([self.graph.edges[u, v].get('dcost', 0.0)
                               for u, v in path], dtype=np.float64)
            # Call the Numba-optimized function
            result = _calculate_path_derivative_numba(derivs)
            # Numba might return np.inf if an input was inf, ensure Python float inf
            return float(result) if result == np.inf else result
        except KeyError as e:
            logger.warning(f"Edge {e} not found during path derivative calculation for path {path}. Returning inf.")
            return float('inf')
        except Exception as e:
            logger.error(
                f"Unexpected error in _calculate_path_derivative for path {path}: {e}", exc_info=self.verbose)
            return float('inf')

    # --- _get_min_origin_flow ---

    def _get_min_origin_flow(self, path: Tuple[Tuple[Any, Any], ...], origin_node: Any) -> float:
        """Finds the minimum flow attributed to 'origin_node' along the 'path'."""
        if not path:
            return 0.0

        # Initialize before loop
        min_flow_on_path = float('inf')

        for u, v in path:
            try:
                edge_data = self.graph.edges[u, v]
                # Use original origin_node as key for obflow lookup
                flow_on_this_edge = edge_data.get(
                    'obflow', {}).get(origin_node, 0.0)
                # Explicitly compare and update
                if flow_on_this_edge < min_flow_on_path:
                    min_flow_on_path = flow_on_this_edge
            except KeyError:
                logger.error(
                    f"Edge ({u}, {v}) not found during min origin flow calculation for origin {origin_node}.")
                return 0.0  # Path invalid - return immediately

        # Return the result after the loop
        if min_flow_on_path == float('inf'):
            # Path existed but had no flow for this origin on any edge, or path was empty
            return 0.0
        else:
            # Return the non-negative minimum flow found
            return max(0.0, min_flow_on_path)  # Ensure non-negative
    # --- End Corrected _get_min_origin_flow ---

    # --- _calculate_shift_value ---
    def _calculate_shift_value(self, pas: PAS) -> float:
        """
        Full Newton step that completely equalizes travel times between the two segments.
        δ = (t_e2 - t_e1) / (Σ t'_{ij,e1} + Σ t'_{ij,e2}), clamped by available origin flow.
        """
        cost1 = self._calculate_path_cost(pas.e1)
        cost2 = self._calculate_path_cost(pas.e2)
        if abs(cost1 - cost2) < EPSILON or cost1 == float('inf') or cost2 == float('inf'):
            return 0.0
        deriv1 = self._calculate_path_derivative(pas.e1)
        deriv2 = self._calculate_path_derivative(pas.e2)
        denominator = deriv1 + deriv2

        if deriv1 == float('inf') or deriv2 == float('inf'):
            return 0.0

        if abs(denominator) < EPSILON:
            # Flat cost curve: shift all available flow to the cheaper segment
            if cost2 > cost1:
                return self._get_min_origin_flow(pas.e2, pas.o)
            else:
                return -self._get_min_origin_flow(pas.e1, pas.o)

        # Full Newton step: equalizes costs in one step
        shift = (cost2 - cost1) / denominator

        # Clamp by available origin flow on the segment being reduced
        if shift > 0:
            max_shift = self._get_min_origin_flow(pas.e2, pas.o)
            return min(shift, max_shift)
        elif shift < 0:
            max_shift = self._get_min_origin_flow(pas.e1, pas.o)
            return max(shift, -max_shift)
        else:
            return 0.0
    # --- End _calculate_shift_value ---

    # --- _shift_flow (Corrected v11) ---
    def _shift_flow(self, pas: PAS, shift_value: float) -> bool:
        """Applies the flow shift to the graph for a given PAS."""
        if abs(shift_value) < MIN_SHIFT:
            return False
        if pas.o is None:
            logger.error(f"Shift Flow Error: PAS has None origin: {pas}")
            return False
        origin_node = pas.o
        shifted = False
        affected_edges = set(pas.e1) | set(pas.e2)
        
        try:
            flow_change_e1 = shift_value
            flow_change_e2 = -shift_value
            
            for u, v in pas.e1:
                data = self.graph.edges[u, v]
                if 'obflow' not in data:
                    data['obflow'] = defaultdict(float)
                data['obflow'][origin_node] = max(
                    0.0, data['obflow'].get(origin_node, 0.0) + flow_change_e1)
                
                # Skip OD flow updates during shifts to avoid bloat and slowness
                        
            for u, v in pas.e2:
                data = self.graph.edges[u, v]
                if 'obflow' not in data:
                    data['obflow'] = defaultdict(float)
                data['obflow'][origin_node] = max(
                    0.0, data['obflow'].get(origin_node, 0.0) + flow_change_e2)
                
                # Skip OD flow updates during shifts to avoid bloat and slowness
                    
            for u, v in affected_edges:
                data = self.graph.edges[u, v]
                data['flow'] = sum(data.get('obflow', {}).values())
            shifted = True
        except KeyError as e:
            logger.error(
                f"Shift Flow Error: Edge not found for PAS {pas}: {e}")
            return False
        except Exception as e:
            logger.exception(
                f"Shift Flow Error: Unexpected error for PAS {pas}: {e}")
            return False
        return shifted
    # --- End Corrected _shift_flow ---


    # --- _update_affected_edge_costs ---
    def _update_affected_edge_costs(self, affected_edges: Set[Tuple[Any, Any]]):
        """Recalculates costs only for edges affected by a flow shift."""
        # logger.debug(f"Updating costs for {len(affected_edges)} affected edges...") # Optional: Verbose logging
        for u, v in affected_edges:
            try:
                data = self.graph.edges[u, v]
                flow = max(0.0, data.get('flow', 0.0))
                data['cost'], data['dcost'] = self.cost_function(data, flow)
                if data['cost'] == float('inf') or math.isnan(data['cost']):
                    data['cost'] = float('inf')
                    data['dcost'] = 0.0
            except KeyError:
                logger.error(
                    f"Error updating cost: Edge ({u}, {v}) not found.")
            except Exception as e:
                logger.error(f"Error calculating cost for edge ({u}, {v}): {e}",
                             exc_info=self.verbose)
                data['cost'], data['dcost'] = float('inf'), 0.0

    # --- _refine_pas_list ---
    def _get_min_total_flow(self, path: Tuple[Tuple[Any, Any], ...]) -> float:
        """Returns the minimum *total* flow (all origins) along path — f_e from the paper."""
        if not path:
            return 0.0
        min_flow = float('inf')
        for u, v in path:
            try:
                flow = self.graph.edges[u, v].get('flow', 0.0)
                if flow < min_flow:
                    min_flow = flow
            except KeyError:
                return 0.0
        return max(0.0, min_flow) if min_flow != float('inf') else 0.0

    def _find_replacement_origin(
        self,
        pas: PAS,
        sorted_origins: List[Any],
    ) -> Optional[Any]:
        """
        Step 2.1 of Procedure 5 (Xie & Xie 2016): when a PAS is stale for its current
        origin r_0 (no flow on one segment), search up to 50 adjacent origins in the
        canonical ordered list to find a replacement that carries positive flow on BOTH
        segments.  Returns the first such origin, or None if none found within 50 steps.
        """
        if not sorted_origins:
            return None
        try:
            idx = sorted_origins.index(pas.o)
        except ValueError:
            return None
        n = len(sorted_origins)
        for i in range(1, 51):
            candidate = sorted_origins[(idx + i) % n]
            if candidate == pas.o:
                continue
            if (self._get_min_origin_flow(pas.e1, candidate) > EPSILON and
                    self._get_min_origin_flow(pas.e2, candidate) > EPSILON):
                return candidate
        return None

    def _local_pas_flow_shift(self) -> None:
        """
        Step 1.3 of Procedure 5 (Xie & Xie 2016): randomly sample a small number of
        PASs from P and perform one Newton flow shift for each, to accelerate convergence.
        Sample size: 100 for small networks (≤ 10 000 OD pairs), 400 for large ones.
        """
        if not self.pas_set:
            return
        sample_size = 400 if len(self.demand) > 10000 else 40
        pas_list = list(self.pas_set)
        if len(pas_list) > sample_size:
            pas_list = random.sample(pas_list, sample_size)
        shifted = 0
        for pas in pas_list:
            shift = self._calculate_shift_value(pas)
            if abs(shift) > MIN_SHIFT:
                if self._shift_flow(pas, shift):
                    self._update_affected_edge_costs(set(pas.e1) | set(pas.e2))
                    shifted += 1
        logger.debug(f"Step 1.3 local PAS shift: sampled {len(pas_list)}, shifted {shifted}.")

    def _refine_pas_list(self):
        """
        Step 2 of Procedure 5 (iTAPAS): repeat 20 times per major iteration.

        Inner pass:
          2.1  Origin-specific stale check: if f_e1^{r0}=0 or f_e2^{r0}=0 and costs
               differ, try to swap r0 for one of the next 50 origins in sorted order
               that has flow on both segments.  If none found, eliminate the PAS.
          2.2  Full Newton flow shift from the higher-cost to the lower-cost segment.
        """
        if not self.pas_set:
            return

        PAS_REFINEMENT_ITERATIONS = 20

        # Canonical sorted list of all origin nodes used for the origin-swap search.
        try:
            sorted_origins: List[Any] = sorted(self.from_cent.values())
        except TypeError:
            sorted_origins = list(self.from_cent.values())

        logger.debug(
            f"Refining {len(self.pas_set)} PAS over {PAS_REFINEMENT_ITERATIONS} inner iterations...")

        for inner_iter in range(PAS_REFINEMENT_ITERATIONS):
            if not self.pas_set:
                break

            pas_changed_count = 0
            to_remove: Set[PAS] = set()
            to_add:    List[PAS] = []

            for pas in list(self.pas_set):
                if pas not in self.pas_set:
                    continue

                # Step 2.1: origin-specific stale check (f^{r0}_{e1} or f^{r0}_{e2} = 0)
                f_e1_r0 = self._get_min_origin_flow(pas.e1, pas.o)
                f_e2_r0 = self._get_min_origin_flow(pas.e2, pas.o)
                cost_e1 = self._calculate_path_cost(pas.e1)
                cost_e2 = self._calculate_path_cost(pas.e2)

                if (f_e1_r0 <= EPSILON or f_e2_r0 <= EPSILON) and abs(cost_e1 - cost_e2) > EPSILON:
                    # Try to save the PAS by finding an adjacent origin with flow on both segs
                    new_o = self._find_replacement_origin(pas, sorted_origins)
                    to_remove.add(pas)
                    if new_o is not None:
                        swapped = PAS(
                            e1=pas.e1, e2=pas.e2, o=new_o,
                            conn_node=pas.conn_node, cost_diff=pas.cost_diff,
                        )
                        to_add.append(swapped)
                    continue

                # Step 2.2: full Newton shift
                shift = self._calculate_shift_value(pas)
                if abs(shift) > MIN_SHIFT:
                    if self._shift_flow(pas, shift):
                        pas_changed_count += 1
                        self._update_affected_edge_costs(set(pas.e1) | set(pas.e2))

            for p in to_remove:
                self._remove_from_pas_set(p)
            swapped_count = sum(self._add_to_pas_set(p) for p in to_add)
            logger.debug(
                f"  Inner iter {inner_iter + 1}/{PAS_REFINEMENT_ITERATIONS}: "
                f"{pas_changed_count} shifted, {len(to_remove) - swapped_count} pruned, "
                f"{swapped_count} origin-swapped, {len(self.pas_set)} PAS remaining.")
    # --- End _refine_pas_list ---

    # --- Procedure 6: Post-process for consistency and proportionality ---

    def _compute_node_flows(self) -> Dict[Any, Dict[Any, float]]:
        """
        Compute η_j^r = total origin-r inflow at node j, for all origins r and nodes j.
        For the origin node r itself, η_r^r = total demand generated from origin r.
        Used by Procedure 6 (Xie & Xie 2016) and the entropy formula (Eq. 24).
        """
        eta: Dict[Any, Dict[Any, float]] = {}
        for u, v, data in self.graph.edges(data=True):
            for r, x in data.get('obflow', {}).items():
                if x > 0.0:
                    if r not in eta:
                        eta[r] = {}
                    eta[r][v] = eta[r].get(v, 0.0) + x
        # Origin node: set η_r^r = total demand generated from origin r
        for orig_zone, orig_node in self.from_cent.items():
            total_d = sum(
                d for (oz, _), d in self.demand.items() if oz == orig_zone and d > 0.0
            )
            if total_d > 0.0:
                if orig_node not in eta:
                    eta[orig_node] = {}
                eta[orig_node][orig_node] = total_d
        return eta

    def _compute_link_proportions(
        self, eta: Dict[Any, Dict[Any, float]]
    ) -> Dict[Tuple[Any, Any], Dict[Any, float]]:
        """
        Compute φ_ij^r = x_ij^r / η_j^r (link approach proportion at head node j).
        Used in Eq. 19 of Xie & Xie (2016) to calculate segment flows.
        """
        phi: Dict[Tuple[Any, Any], Dict[Any, float]] = {}
        for u, v, data in self.graph.edges(data=True):
            phi_edge: Dict[Any, float] = {}
            for r, x in data.get('obflow', {}).items():
                if x > EPSILON:
                    eta_v = eta.get(r, {}).get(v, 0.0)
                    if eta_v > EPSILON:
                        phi_edge[r] = x / eta_v
            if phi_edge:
                phi[(u, v)] = phi_edge
        return phi

    def _compute_segment_flow_g(
        self,
        segment: Tuple[Tuple[Any, Any], ...],
        origin: Any,
        eta: Dict[Any, Dict[Any, float]],
        phi: Dict[Tuple[Any, Any], Dict[Any, float]],
    ) -> float:
        """
        Compute g_e^r = η_{e_h}^r × ∏_{(i,j)∈e} φ_ij^r  (Eq. 19, Xie & Xie 2016).
        Represents the proportional segment flow for origin r on segment e.
        """
        if not segment:
            return 0.0
        head_node = segment[-1][1]
        eta_head = eta.get(origin, {}).get(head_node, 0.0)
        if eta_head <= EPSILON:
            return 0.0
        prod = 1.0
        for link in segment:
            pv = phi.get(link, {}).get(origin, 0.0)
            if pv <= EPSILON:
                return 0.0
            prod *= pv
        return eta_head * prod

    def _find_consistency_pas(self) -> int:
        """
        Procedure 6, Step 1: identify and add PASs for consistent but empty optimal paths.

        For every origin r and every link (i,j) that carries zero origin-r flow,
        checks whether the link lies on an optimal path:
            dist_r[i] + cost(i,j) ≈ dist_r[j]
        If so, creates PAS(e1=SPT_path[i→j], e2=[(i,j)]) so the proportionality
        loop can spread flow onto the currently-empty segment.

        Returns the number of new PASs added.
        """
        new_count = 0
        cost_tol_rel = 1e-6  # relative tolerance for optimality check

        for orig_zone, orig_node in self.from_cent.items():
            # skip origins with no demand
            total_d = sum(
                d for (oz, _), d in self.demand.items() if oz == orig_zone and d > 0.0
            )
            if total_d <= 0.0:
                continue

            # Shortest-path distances from orig_node using current UE costs
            try:
                dist_r: Dict[Any, float] = nx.single_source_dijkstra_path_length(
                    self.graph, orig_node, weight='cost'
                )
            except Exception:
                continue

            # Build SPT predecessor map: for each reachable node, find the
            # predecessor that achieves dist_r[node]
            pred_map: Dict[Any, Optional[Any]] = {orig_node: None}
            for node in dist_r:
                if node == orig_node:
                    continue
                d_node = dist_r[node]
                best_pred: Optional[Any] = None
                for p in self.graph.predecessors(node):
                    if p not in dist_r:
                        continue
                    edge_c = self.graph.edges[p, node].get('cost', float('inf'))
                    if abs(dist_r[p] + edge_c - d_node) < cost_tol_rel * max(1.0, d_node):
                        if best_pred is None or dist_r[p] < dist_r[best_pred]:
                            best_pred = p
                pred_map[node] = best_pred

            # Cache: node → full SPT path [orig_node, ..., node]
            path_cache: Dict[Any, Optional[List[Any]]] = {orig_node: [orig_node]}

            def spt_path(dest: Any) -> Optional[List[Any]]:
                if dest in path_cache:
                    return path_cache[dest]
                # Iterative reconstruction to avoid recursion-depth limits
                chain: List[Any] = []
                cur: Any = dest
                while cur is not None and cur not in path_cache:
                    chain.append(cur)
                    cur = pred_map.get(cur)  # type: ignore[assignment]
                if cur is None:
                    path_cache[dest] = None
                    return None
                base = path_cache[cur]
                if base is None:
                    path_cache[dest] = None
                    return None
                # Fill cache for every node along the reconstructed chain
                acc = list(base)
                for node in reversed(chain):
                    acc = acc + [node]
                    path_cache[node] = list(acc)
                return path_cache[dest]

            # Scan every link for zero origin-r flow that lies on an optimal path
            for i_node, j_node, data in self.graph.edges(data=True):
                if data.get('obflow', {}).get(orig_node, 0.0) > EPSILON:
                    continue  # already carries flow — not a missing path

                if i_node not in dist_r or j_node not in dist_r:
                    continue

                link_cost = data.get('cost', float('inf'))
                if link_cost == float('inf'):
                    continue

                dist_j = dist_r[j_node]
                threshold = cost_tol_rel * max(1.0, dist_j)
                if abs(dist_r[i_node] + link_cost - dist_j) > threshold:
                    continue  # not an optimal link for this origin

                # Build e1: SPT sub-path from i_node to j_node
                path_to_j = spt_path(j_node)
                if path_to_j is None or i_node not in path_to_j:
                    continue

                idx = path_to_j.index(i_node)
                e1_nodes = path_to_j[idx:]   # [i_node, ..., j_node]
                if len(e1_nodes) < 2:
                    continue
                e1_edges = tuple(zip(e1_nodes[:-1], e1_nodes[1:]))
                e2_edges: Tuple[Tuple[Any, Any], ...] = ((i_node, j_node),)

                if e1_edges == e2_edges:
                    continue  # trivial single-link SPT — nothing to spread

                new_pas = PAS(
                    e1=e1_edges, e2=e2_edges,
                    o=orig_node, conn_node=i_node, cost_diff=0.0
                )
                if self._add_to_pas_set(new_pas):
                    new_count += 1

        return new_count

    def _compute_entropy(self, eta: Dict[Any, Dict[Any, float]]) -> float:
        """
        Compute origin-based entropy (Bar-Gera 2010, Eq. 24 of Xie & Xie 2016):
          E = -Σ_{r∈R} Σ_{(i,j)∈A} x_ij^r · ln(x_ij^r / η_j^r)
            = Σ_{r∈R} Σ_{(i,j)∈A} x_ij^r · ln(η_j^r / x_ij^r)
        Larger E → more entropy → more proportional flow distribution.
        """
        E = 0.0
        for u, v, data in self.graph.edges(data=True):
            for r, x in data.get('obflow', {}).items():
                if x > EPSILON:
                    eta_v = eta.get(r, {}).get(v, 0.0)
                    if eta_v > EPSILON:
                        ratio = x / eta_v
                        if ratio > EPSILON:
                            E -= x * math.log(ratio)
        return E

    def _post_process_procedure6(
        self,
        eps: float = 1e-6,
        max_outer: int = 100,
    ) -> Tuple[float, float]:
        """
        Post-process for consistency and proportionality — Procedure 6 of Xie & Xie (2016).

        Equalises the proportion of flow on the two segments of each PAS across all origins
        sharing those segments, then computes the origin-based entropy (Eq. 24).

        Steps:
          0. Compute node flows η_j^r and link proportions φ_ij^r.
          1. Consistency search: find and add PASs for optimal links carrying zero flow.
          2. For each PAS, equalise proportions across R_p (origins with flow on both segments).
          3. Convergence test: G_p = Σ_{p∈P} max{|ρ−ρ^r|} / |P| < eps → stop.

        Returns:
          (entropy, G_p)  where entropy is Eq. 24 and G_p is the final proportionality gap.
        """
        if not self.pas_set:
            eta = self._compute_node_flows()
            return self._compute_entropy(eta), 0.0

        logger.info("--- Running Procedure 6 (post-process: consistency & proportionality) ---")

        G_p = float('inf')

        for outer_iter in range(1, max_outer + 1):
            # Step 1: consistency search — add PASs for optimal links with zero flow.
            # Only run on the first iteration; subsequent iterations refine proportions
            # on the now-expanded PAS set.
            if outer_iter == 1:
                n_new = self._find_consistency_pas()
                if n_new > 0:
                    logger.info(f"  Procedure 6 Step 1: added {n_new} consistency PASs.")

            # Step 0: compute node flows and link proportions
            eta = self._compute_node_flows()
            phi = self._compute_link_proportions(eta)

            G_p_sum = 0.0
            pas_active = 0

            for pas in list(self.pas_set):
                if not pas.e1 or not pas.e2:
                    continue

                # Step 2.0: collect R_p — origins with flow on any PAS link
                R_p_candidates: Set[Any] = set()
                for link in list(pas.e1) + list(pas.e2):
                    if self.graph.has_edge(*link):
                        for r, x in self.graph.edges[link].get('obflow', {}).items():
                            if x > EPSILON:
                                R_p_candidates.add(r)

                # Fix 1: only keep origins that actually traverse a segment
                # (g_e1^r + g_e2^r > 0).  Origins with zero combined segment flow
                # contribute nothing and would pollute ρ^r with the arbitrary 0.5
                # fallback, preventing G_p from converging.
                g1: Dict[Any, float] = {}
                g2: Dict[Any, float] = {}
                rho_r: Dict[Any, float] = {}
                R_p: Set[Any] = set()

                for r in R_p_candidates:
                    ge1 = self._compute_segment_flow_g(pas.e1, r, eta, phi)
                    ge2 = self._compute_segment_flow_g(pas.e2, r, eta, phi)
                    total_g = ge1 + ge2
                    if total_g <= EPSILON:
                        continue  # origin doesn't traverse either segment — skip
                    g1[r] = ge1
                    g2[r] = ge2
                    rho_r[r] = ge1 / total_g
                    R_p.add(r)

                if len(R_p) < 2:
                    continue  # need at least two origins to have a proportionality issue

                # Step 2.2: overall PAS proportion ρ (Eq. 21)
                sum_g1  = sum(g1.values())
                sum_g12 = sum(g1[r] + g2[r] for r in R_p)
                if sum_g12 <= EPSILON:
                    continue
                rho = sum_g1 / sum_g12

                # Proportion gap for this PAS
                max_diff = max(abs(rho - rho_r[r]) for r in R_p)
                G_p_sum += max_diff
                pas_active += 1

                # Skip shift if already within tolerance
                if max_diff <= eps:
                    continue

                # Apply flow adjustments δ^r (Eq. 22)
                affected: Set[Tuple[Any, Any]] = set()
                for r in R_p:
                    total_g_r = g1[r] + g2[r]
                    delta_r = rho * total_g_r - g1[r]
                    if abs(delta_r) <= EPSILON:
                        continue
                    # Clamp by available origin-based flow on the segment being reduced
                    if delta_r > 0:
                        avail = self._get_min_origin_flow(pas.e2, r)
                        delta_r = min(delta_r, avail)
                    else:
                        avail = self._get_min_origin_flow(pas.e1, r)
                        delta_r = max(delta_r, -avail)
                    if abs(delta_r) <= EPSILON:
                        continue
                    # x_ij^r += δ^r on e1,  -= δ^r on e2
                    for u, v in pas.e1:
                        if self.graph.has_edge(u, v):
                            d = self.graph.edges[u, v]
                            d.setdefault('obflow', defaultdict(float))
                            d['obflow'][r] = max(0.0, d['obflow'].get(r, 0.0) + delta_r)
                            affected.add((u, v))
                    for u, v in pas.e2:
                        if self.graph.has_edge(u, v):
                            d = self.graph.edges[u, v]
                            d.setdefault('obflow', defaultdict(float))
                            d['obflow'][r] = max(0.0, d['obflow'].get(r, 0.0) - delta_r)
                            affected.add((u, v))

                # Step 2.3: update total link flows and immediately refresh η, φ
                # (Fix 2: stale proportions — paper requires η/φ update after each shift)
                for u, v in affected:
                    d = self.graph.edges[u, v]
                    d['flow'] = sum(d.get('obflow', {}).values())
                if affected:
                    eta = self._compute_node_flows()
                    phi = self._compute_link_proportions(eta)

            # Step 3: convergence test
            G_p = G_p_sum / pas_active if pas_active > 0 else 0.0
            logger.info(
                f"  Procedure 6 iter {outer_iter:3d}: G_p = {G_p:.6e}  (target {eps:.0e},"
                f" active PASs = {pas_active})"
            )
            if G_p <= eps:
                logger.info(f"Procedure 6 converged in {outer_iter} iterations.")
                break
        else:
            logger.warning(
                f"Procedure 6 reached max iterations ({max_outer}). Final G_p = {G_p:.6e}"
            )

        # Final entropy
        eta_final = self._compute_node_flows()
        entropy = self._compute_entropy(eta_final)
        logger.info(f"Procedure 6 complete. Entropy = {entropy:.4f},  G_p = {G_p:.6e}")
        return entropy, G_p

    # --- Proportions Extraction ---
    def _extract_proportions(self) -> Dict[Tuple[Any, Any], Dict[Tuple[Any, Any], float]]:
        """Extracts OD link proportions from directly tracked od_flow data."""
        od_link_prop = {}
        
        # Simply extract the directly tracked OD flows
        for edge_key, edge_data in self.graph.edges.items():
            od_flows = edge_data.get('od_flow', {})
            
            if od_flows:
                # Convert to clean Python floats and filter out very small flows
                clean_od_flows = {}
                for od_pair, flow_val in od_flows.items():
                    clean_flow = float(flow_val)
                    if clean_flow > EPSILON:
                        clean_od_flows[od_pair] = clean_flow
                
                if clean_od_flows:
                    od_link_prop[edge_key] = clean_od_flows
        
        return od_link_prop

    # --- Main Execution Method ---

    def run(self) -> Tuple[nx.DiGraph, float, Dict, Dict, float, int]:
        """Executes the TAPAS algorithm."""
        start_time = time.perf_counter()  # Overall start time
        logger.info("--- Starting TAPAS Algorithm ---")
        try:
            _ = self._run_all_or_nothing()  # Initial AON and cost calculation
            while self.iterations < self.max_iter:
                self.iterations += 1

                # 1. Calculate SP, TSTT, SPTT, Gap based on current costs
                # _calculate_sp_and_gap also calls _initialize_graph_costs internally
                self.tstt, self.sptt, self.gap, all_sp_paths, all_sp_cost = self._calculate_sp_and_gap()

                # Calculate total elapsed time
                elapsed_time = time.perf_counter() - start_time

                # Update log message format
                logger.info(
                    f"Iteration: {self.iterations}; Current gap: {self.gap:.6f}; TSTT: {self.tstt:.2f}; SPTT: {self.sptt:.2f} ; Elapsed: {elapsed_time:.2f}s"
                )

                # 2. Check convergence
                if self.gap <= self.accuracy:
                    logger.info(
                        f"Convergence criteria met (Gap {self.gap:.6f} <= {self.accuracy}).")
                    break

                # 3. Identify potential edges and find new PAS
                new_pas_count = self._identify_potential_edges_and_find_pas(
                    all_sp_paths, all_sp_cost)

                # 4. Check for stagnation (no new PAS and empty set)
                if new_pas_count == 0 and not self.pas_set:
                    logger.info(
                        f"Iter {self.iterations}: No new PAS found and PAS set empty. Stopping.")
                    # Re-calculate final gap based on current state before exiting
                    self.tstt, self.sptt, self.gap, _, _ = self._calculate_sp_and_gap()
                    break

                # 5. Step 1.3: local PAS flow shift on a random sample (before 20-pass refinement)
                if self.pas_set:
                    self._local_pas_flow_shift()

                # 6. Refine PAS list (Step 2: 20-pass refinement with origin-swap)
                if self.pas_set:
                    self._refine_pas_list()
                else:
                    logger.info(
                        f"Iter {self.iterations}: PAS set empty, skipping refinement.")

            # --- Loop End ---

            if self.iterations >= self.max_iter and self.gap > self.accuracy:
                logger.warning(
                    f"TAPAS reached max iterations ({self.max_iter}). Final Gap: {self.gap:.6f}")
            elif self.gap <= self.accuracy:
                logger.info(
                    f"TAPAS converged in {self.iterations} iterations.")

            logger.info("Calculating final shortest paths for output...")
            self._initialize_graph_costs()
            final_sptt = 0.0
            self.sp_edges.clear()
            self.eodtt.clear()
            for origin_zone, origin_node in self.from_cent.items():
                if origin_node not in self.graph:
                    continue
                try:
                    cost_dict_final, path_dict_final = nx.single_source_dijkstra(
                        self.graph, origin_node, weight='cost')
                    for dest_zone, dest_node in self.to_cent.items():
                        od_pair = (origin_zone, dest_zone)
                        demand_val = self.demand.get(od_pair, 0.0)
                        if origin_zone == dest_zone or demand_val < EPSILON:
                            if od_pair in self.demand:
                                self.sp_edges[od_pair] = []
                                self.eodtt[od_pair] = 0.0 if origin_zone == dest_zone else float(
                                    'inf')
                            continue

                        path_cost = cost_dict_final.get(
                            dest_node, float('inf'))

                        if od_pair in self.demand:
                            if path_cost != float('inf'):
                                path_nodes = path_dict_final.get(dest_node)
                                self.sp_edges[od_pair] = path_nodes if path_nodes else [
                                ]
                                self.eodtt[od_pair] = path_cost
                                final_sptt += demand_val * path_cost
                            else:
                                self.sp_edges[od_pair] = []
                                self.eodtt[od_pair] = float('inf')
                except Exception as e:
                    logger.exception(
                        f"Error during final SP calculation for origin {origin_node}: {e}")
                    # Attempt to fill remaining ODs for this origin with error state
                    for dest_zone, dest_node_orig in self.to_cent.items():
                        od_pair = (origin_zone, dest_zone)
                        if od_pair in self.demand and od_pair not in self.sp_edges:
                            self.sp_edges[od_pair] = []
                            self.eodtt[od_pair] = float('inf')

            self.tstt = sum(max(0.0, data.get('flow', 0.0)) * data['cost'] for _, _, data in self.graph.edges(
                data=True) if data.get('cost', float('inf')) != float('inf'))
            self.gap = max(0.0, (self.tstt / final_sptt) - 1.0) if final_sptt > EPSILON else (
                float('inf') if self.tstt > EPSILON else 0.0)

            # Step 4: Post-process for consistency and proportionality (Procedure 6, Appendix C)
            self.entropy, self.prop_gap = self._post_process_procedure6(
                eps=max(self.accuracy, 1e-6), max_outer=100
            )
            self.graph.graph['entropy']  = self.entropy
            self.graph.graph['prop_gap'] = self.prop_gap

            total_time = time.perf_counter() - start_time
            logger.info("--- TAPAS Algorithm Finished ---")
            logger.info(
                f"Iterations: {self.iterations}, Final Gap: {self.gap:.6f}, Final TSTT: {self.tstt:.2f}, Total Time: {total_time:.2f}s")
            
            # Extract proportions if requested
            if self.proportions:
                od_link_prop = self._extract_proportions()
                return self.graph, self.tstt, self.sp_edges, self.eodtt, self.gap, self.iterations, od_link_prop
            else:
                return self.graph, self.tstt, self.sp_edges, self.eodtt, self.gap, self.iterations
        except TypeError as e:
            logger.error(
                f"TAPAS run failed due to unhashable node type: {e}", exc_info=True)
            if self.proportions:
                return self.graph, float('inf'), {}, {}, float('inf'), self.iterations, {}
            else:
                return self.graph, float('inf'), {}, {}, float('inf'), self.iterations
        except Exception as e:
            logger.exception(f"Fatal error during TAPAS run: {e}")
            if self.proportions:
                return self.graph, self.tstt, self.sp_edges, self.eodtt, self.gap, self.iterations, {}
            else:
                return self.graph, self.tstt, self.sp_edges, self.eodtt, self.gap, self.iterations


# --- Wrapper Function ---
def TAPAS(graph: nx.DiGraph,
          from_cent: Dict[Any, Any],
          to_cent: Dict[Any, Any],
          demand: Dict[Tuple[Any, Any], float],
          cost_function: Callable,
          max_iter: int = 100,
          accuracy: float = 0.001,
          x_bar: Optional[Dict[Tuple[Any, Any], float]] = None,
          ob_flow: Optional[Dict[Tuple[Any, Any], Dict[Any, float]]] = None,
          start_time: Optional[float] = None,
          verbose: bool = False,
          proportions: bool = False
          ):
    """Wrapper function to execute the TAPAS algorithm."""
    if start_time is None:
        start_time = time.perf_counter()
    try:
        solver = TAPASAlgorithm(
            graph=graph, demand=demand, from_cent=from_cent, to_cent=to_cent,
            cost_function=cost_function, max_iter=max_iter, accuracy=accuracy, verbose=verbose, proportions=proportions
        )
        result = solver.run()
        
        if proportions:
            graph_result, tstt, sp_edges, eodtt, gap, iterations, od_link_prop = result
            if not isinstance(graph_result, nx.DiGraph) or not isinstance(sp_edges, dict) or not isinstance(eodtt, dict):
                logger.error("TAPAS run method did not return expected data types.")
                return graph, float('inf'), {}, {}, {}
            return graph_result, tstt, sp_edges, eodtt, od_link_prop
        else:
            graph_result, tstt, sp_edges, eodtt, gap, iterations = result
            if not isinstance(graph_result, nx.DiGraph) or not isinstance(sp_edges, dict) or not isinstance(eodtt, dict):
                logger.error("TAPAS run method did not return expected data types.")
                return graph, float('inf'), {}, {}, float('inf'), 0
            return graph_result, tstt, sp_edges, eodtt, gap, iterations
    except TypeError as e:  # Catch validation errors from __init__
        logger.error(
            f"TAPAS failed initialization due to graph node validation error: {e}", exc_info=True)
        if proportions:
            return graph, float('inf'), {}, {}, {}
        else:
            return graph, float('inf'), {}, {}, float('inf'), 0
    except Exception as e:
        logger.exception(
            f"Unexpected error during TAPAS wrapper execution: {e}")
        if proportions:
            return graph, float('inf'), {}, {}, {}
        else:
            return graph, float('inf'), {}, {}, float('inf'), 0