dag_to_operators_integration.py

#!/usr/bin/env python3
"""
Complete DAG-to-Operators Integration System

This module provides full integration between:
1. Traced neural rendering DAGs (from instrumentation)
2. /Operators framework (realistic characteristics)  
3. Scheduler.IR format (for scheduling)
"""

import sys
import pickle
import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
import networkx as nx
import os
import math
import re

# Add paths for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent / 'Operators'))

# Import operator mapping and transformation
try:
    from operator_mapping import map_function_to_operator_type, map_operator_to_hardware_type
except ImportError:
    from Instrumentation.operator_mapping import map_function_to_operator_type, map_operator_to_hardware_type

# Import Scheduler IR format
try:
    from Scheduler.IR import OperatorGraph as SchedulerOperatorGraph, OperatorNode, TensorDesc
except ImportError:
    sys.path.insert(0, '.')
    from Scheduler.IR import OperatorGraph as SchedulerOperatorGraph, OperatorNode, TensorDesc

# Import /Operators framework
try:
    from operators.sampling_operator import UniformSamplerOperator, PDFSamplerOperator, FrustrumCullingOperator
    from operators.encoding_operator import HashEncodingOperator, RFFEncodingOperator
    from operators.computation_operator import MLPOperator
    from operators.blending_operator import RGBRendererOperator, DensityRendererOperator
    from utils.operator_graph import OperatorGraph as OperatorsGraph
except ImportError:
    sys.path.insert(0, 'Operators')
    from operators.sampling_operator import UniformSamplerOperator, PDFSamplerOperator, FrustrumCullingOperator
    from operators.encoding_operator import HashEncodingOperator, RFFEncodingOperator
    from operators.computation_operator import MLPOperator
    from operators.blending_operator import RGBRendererOperator, DensityRendererOperator
    from utils.operator_graph import OperatorGraph as OperatorsGraph

class OperatorFactory:
    """Factory to create actual /Operators instances from traced data."""
    
    @staticmethod
    def create_operator(function_name: str, node_data: Dict[str, Any], dim: Tuple[int, int]) -> Optional[Any]:
        """Create an actual /Operators instance from traced node data."""
        op_type = map_function_to_operator_type(function_name)
 
        # Do not create an operator for model-level wrappers; they are orchestration only
        if op_type == 'MODEL_WRAPPER':
            return None

        try:
            if 'UNIFORM_SAMPLING' in op_type or 'RAY_SAMPLING' in op_type:
                return UniformSamplerOperator(dim, sampler_type="uniform", bitwidth=16)
            
            elif 'PDF_SAMPLING' in op_type:
                return PDFSamplerOperator(dim, bitwidth=16)
            
            elif 'FRUSTUM_SAMPLING' in op_type:
                return FrustrumCullingOperator(dim, fov=60.0, near=0.1, far=100.0)
            
            elif 'HASH_ENCODING' in op_type:
                return HashEncodingOperator(
                    dim, 
                    input_dim=3,
                    num_levels=16, 
                    features_per_level=2,
                    bitwidth=16
                )
            
            elif 'POSITIONAL_ENCODING' in op_type or 'RFF_ENCODING' in op_type:
                return RFFEncodingOperator(
                    dim,
                    input_dim=3,
                    num_features=60,  # Typical NeRF positional encoding
                    bitwidth=16
                )
            
            elif 'COMPUTATION' in op_type:
                # Prefer precise configuration from tracing metadata when available
                fn = str(function_name)
                nd = node_data or {}

                # FieldHead (single Linear) — model as 1‑layer MLP with explicit in/out dims
                if 'field_components.field_heads.' in fn:
                    in_dim = None
                    out_dim = None
                    try:
                        in_dim = int(nd.get('field_head_in_dim')) if nd.get('field_head_in_dim') is not None else None
                    except Exception:
                        in_dim = None
                    try:
                        out_dim = int(nd.get('field_head_out_dim')) if nd.get('field_head_out_dim') is not None else None
                    except Exception:
                        out_dim = None

                    # Fallback: parse from shapes (B,N,C)
                    if in_dim is None:
                        try:
                            import re
                            ins = nd.get('input_shapes') or ''
                            m = re.search(r"\((?:\d+,\s*){2}(\d+)\)", str(ins))
                            if m:
                                in_dim = int(m.group(1))
                        except Exception:
                            pass
                    if out_dim is None:
                        try:
                            import re
                            outs = nd.get('output_shapes') or ''
                            m = re.search(r"\((?:\d+,\s*){2}(\d+)\)", str(outs))
                            if m:
                                out_dim = int(m.group(1))
                        except Exception:
                            pass

                    # Reasonable fallbacks if still unknown
                    if in_dim is None:
                        in_dim = 128
                    if out_dim is None:
                        out_dim = 1 if 'DensityFieldHead' in fn else 3 if 'RGBFieldHead' in fn else 4

                    return MLPOperator(
                        dim,
                        in_dim=in_dim,
                        num_layers=1,
                        layer_width=max(in_dim, out_dim),
                        out_dim=out_dim,
                        skip_connections=(),
                        use_bias=True,
                        bitwidth=16
                    )

                # MLP.forward — use exact params captured by tracer if present
                in_dim = None
                out_dim = None
                num_layers = None
                layer_width = None
                skip_connections = None

                try:
                    if nd.get('mlp_in_dim') is not None:
                        in_dim = int(nd['mlp_in_dim'])
                except Exception:
                    pass
                try:
                    if nd.get('mlp_out_dim') is not None:
                        out_dim = int(nd['mlp_out_dim'])
                except Exception:
                    pass
                try:
                    if nd.get('mlp_num_layers') is not None:
                        num_layers = int(nd['mlp_num_layers'])
                except Exception:
                    pass
                try:
                    if nd.get('mlp_layer_width') is not None:
                        layer_width = int(nd['mlp_layer_width'])
                except Exception:
                    pass
                try:
                    sc = nd.get('mlp_skip_connections')
                    if sc is not None:
                        skip_connections = tuple(int(x) for x in sc) if isinstance(sc, (list, tuple)) else None
                except Exception:
                    pass

                # Fallbacks from shapes if any are missing
                if in_dim is None:
                    try:
                        import re
                        ins = nd.get('input_shapes') or nd.get('output_shapes') or ''
                        m = re.search(r"\((?:\d+,\s*){2}(\d+)\)", str(ins))
                        if m:
                            in_dim = int(m.group(1))
                    except Exception:
                        pass

                # Reasonable defaults if still unknown (vanilla NeRF base MLP)
                if in_dim is None:
                    in_dim = 256
                if num_layers is None:
                    num_layers = 8
                if layer_width is None:
                    layer_width = 256
                if out_dim is None:
                    out_dim = layer_width
                if skip_connections is None:
                    skip_connections = (4,)

                return MLPOperator(
                    dim,
                    in_dim=in_dim,
                    num_layers=num_layers,
                    layer_width=layer_width,
                    out_dim=out_dim,
                    skip_connections=skip_connections,
                    use_bias=True,
                    bitwidth=16
                )
            
            elif 'RGB' in op_type and 'RENDERING' in op_type:
                return RGBRendererOperator(dim, background_color="random", bitwidth=16)
            
            elif (
                'DENSITY_RENDERING' in op_type
                or 'DEPTH_RENDERING' in op_type
                or 'ALPHA_BLENDING' in op_type
            ):
                return DensityRendererOperator(dim, method="expected", bitwidth=16)
            
            else:
                # Unknown types: skip creating an operator
                return None
                
        except Exception as e:
            print(f"⚠️ Failed to create operator for {function_name}: {e}")
            # Skip this node on error to avoid bogus operators
            return None

class DAGToOperatorsIntegration:
    """Complete integration system for DAG transformation."""
    
    def __init__(self):
        self.operator_factory = OperatorFactory()
    
    def _is_parameter_artifact(self, function_name: str, node_id: str) -> bool:
        name = (function_name or str(node_id)).lower()
        if name == 'args' or name == 'kwargs':
            return True
        if name.startswith('args') or name.startswith('kwargs'):
            return True
        if 'args[' in name or 'kwargs[' in name:
            return True
        if name.endswith('.self') or name.endswith('].self') or '.self' in name:
            return True
        return False
    
    def extract_tensor_dimensions(self, dag_data: Dict[str, Any]) -> Tuple[int, int]:
        """Extract realistic tensor dimensions from traced DAG data."""
        # Try to find ray-bundle operations to get actual dimensions
        for node_id, node_data in dag_data.get('nodes', {}).items():
            if 'RayBundle' in str(node_id) or 'ray_bundle' in str(node_id):
                if 'inputs' in node_data:
                    for inp in node_data['inputs']:
                        if isinstance(inp, dict) and 'shape' in inp:
                            shape = inp['shape']
                            if len(shape) >= 1:
                                total_rays = shape[0]
                                # Estimate samples per ray (typical NeRF uses 64-128)
                                samples_per_ray = 64
                                if total_rays > 10000:  # Large batch
                                    return total_rays // samples_per_ray, samples_per_ray
                                else:
                                    return total_rays, samples_per_ray
        
        # If shapes are missing, attempt to infer image resolution from dataset images
        try:
            from pathlib import Path
            possible_roots = [
                Path('nerf_synthetic'),
                Path('/tmp/nerf/nerf_synthetic'),
            ]
            image_paths = []
            for root in possible_roots:
                if root.exists():
                    image_paths += list(root.glob('**/train/*.png'))
                    image_paths += list(root.glob('**/train/*.jpg'))
            if image_paths:
                sample_image = image_paths[0]
                try:
                    from PIL import Image
                    with Image.open(sample_image) as im:
                        width, height = im.size
                        if width > 0 and height > 0:
                            # Use full-frame rays as batch size
                            return width * height, 64
                except Exception:
                    # Ignore PIL errors and fall back below
                    pass
        except Exception:
            # Ignore filesystem errors and fall back below
            pass
     
        # Default neural rendering dimensions
        return 4096, 64  # 4096 rays, 64 samples per ray
    
    def _infer_node_dim(self, node_id: str, node_data: Dict[str, Any], default_dim: Tuple[int, int]) -> Tuple[int, int]:
        """Infer per-node (rays, samples) strictly from traced shapes.
        - Try node_data['inputs'] or ['outputs'] entries with concrete shapes.
        - Try parsing shape-aware function_name (e.g., "Func[(...)->key:(B, N,...)]").
        - If unable to infer both B and N, raise ValueError to signal missing instrumentation.
        """
        # 1) Direct shapes on inputs/outputs if present (dict-dag format may include 'inputs'/'outputs')
        for key in ('inputs', 'outputs'):
            for item in (node_data.get(key) or []):
                if isinstance(item, dict) and 'shape' in item:
                    shape = item['shape']
                    if isinstance(shape, (list, tuple)) and len(shape) >= 2:
                        B = shape[0]
                        N = shape[1]
                        if isinstance(B, int) and isinstance(N, int) and B > 0 and N > 0:
                            return (B, N)
        
        # 2) Parse from shape-aware function_name or output_shapes strings
        import re
        func_str = str(node_data.get('function_name', node_id))
        candidates: list[tuple[int, int]] = []
        # Extract bracket content [...] if present
        m = re.search(r"\[(.*?)\]$", func_str)
        shape_sig = m.group(1) if m else ""
        parts_to_search = [shape_sig, str(node_data.get('output_shapes', '')), str(node_data.get('input_shapes', ''))]
        for text in parts_to_search:
            # Look for patterns like :(B, N, ...) or (B, N, ...)
            for mm in re.finditer(r":?\((\d+)\s*,\s*(\d+)\b", text):
                try:
                    B = int(mm.group(1))
                    N = int(mm.group(2))
                    if B > 0 and N > 0:
                        candidates.append((B, N))
                except Exception:
                    pass
        if candidates:
            # Choose the most frequent or first candidate
            return candidates[0]
        
        # 3) Could not infer strictly from shapes
        raise ValueError(f"Unable to infer (rays, samples) for node '{node_id}'. Instrumentation missing shapes. Ensure traced functions expose tensor shapes (e.g., RaySamples, weights, positions).")
    
    def transform_dag_to_operators(self, dag_data: Dict[str, Any]) -> Tuple[OperatorsGraph, Dict[str, Any]]:
        """Transform traced DAG to actual /Operators instances."""
        print(f"🔧 Transforming DAG to /Operators instances...")
        
        # Extract baseline dimensions to use as fallback
        B_base, N_base = self.extract_tensor_dimensions(dag_data)
        default_dim = (B_base, N_base)
        print(f"   📐 Baseline dimensions: {B_base} rays × {N_base} samples")
        
        # Create /Operators graph
        operators_graph = OperatorsGraph()
        node_mapping = {}  # traced_node_id -> operator_instance
        characteristics = {
            'total_flops': 0,
            'total_memory_bytes': 0,
            'operator_types': {},
            'realistic_operators': []
        }
        
        # Create operator instances with per-node dims (strict inference, no heuristics)
        per_node_dims: Dict[str, Tuple[int, int] | None] = {}
        fallback_nodes: set[str] = set()
        
        # Build adjacency for neighbor lookups
        succs: Dict[str, List[str]] = {}
        preds: Dict[str, List[str]] = {}
        for edge in dag_data.get('edges', []) or []:
            if len(edge) >= 2:
                src_id, dst_id = edge[0], edge[1]
                succs.setdefault(src_id, []).append(dst_id)
                preds.setdefault(dst_id, []).append(src_id)
        
        # First pass: strict per-node inference; store None if unavailable
        missing: List[str] = []
        for node_id, node_data in dag_data.get('nodes', {}).items():
            function_name = node_data.get('function_name', str(node_id))
            if self._is_parameter_artifact(function_name, node_id):
                per_node_dims[node_id] = None
                continue
            try:
                node_dim = self._infer_node_dim(node_id, node_data, default_dim)
                per_node_dims[node_id] = node_dim
            except Exception:
                per_node_dims[node_id] = None
                missing.append(node_id)
        
        # Second pass: resolve missing by neighbors (predecessors, then successors)
        if missing:
            unresolved = set(missing)
            changed = True
            max_iters = 3
            it = 0
            while changed and unresolved and it < max_iters:
                changed = False
                it += 1
                to_remove = []
                for nid in list(unresolved):
                    # Check predecessors
                    found = None
                    for p in preds.get(nid, []):
                        p_fn = dag_data['nodes'][p].get('function_name', str(p))
                        if self._is_parameter_artifact(p_fn, p):
                            continue
                        if per_node_dims.get(p) is not None:
                            found = per_node_dims[p]
                            break
                    # If not found, check successors
                    if found is None:
                        for s in succs.get(nid, []):
                            s_fn = dag_data['nodes'][s].get('function_name', str(s))
                            if self._is_parameter_artifact(s_fn, s):
                                continue
                            if per_node_dims.get(s) is not None:
                                found = per_node_dims[s]
                                break
                    if found is not None:
                        per_node_dims[nid] = found
                        to_remove.append(nid)
                        changed = True
                for nid in to_remove:
                    unresolved.discard(nid)
        
        # If still unresolved, default to baseline dims and warn (lenient mode)
        still_missing = []
        for nid, dim in per_node_dims.items():
            fn = dag_data['nodes'][nid].get('function_name', str(nid))
            if self._is_parameter_artifact(fn, nid):
                continue
            if dim is None:
                still_missing.append(nid)
        if still_missing:
            try:
                sample_list = []
                for nid in still_missing[:5]:
                    f = dag_data['nodes'][nid].get('function_name', str(nid))
                    sample_list.append(f"{nid} -> {f}")
                print(f"⚠️  {len(still_missing)} nodes missing dims; defaulting to baseline {default_dim}. Examples: " + "; ".join(sample_list))
            except Exception:
                pass
            for nid in still_missing:
                per_node_dims[nid] = default_dim
                try:
                    fallback_nodes.add(str(nid))
                except Exception:
                    pass
        
        # Create operators
        for node_id, node_data in dag_data.get('nodes', {}).items():
            function_name = node_data.get('function_name', str(node_id))
            if self._is_parameter_artifact(function_name, node_id):
                continue
            # Skip nodes that fell back to baseline dims to avoid scheduling unrealistic operators
            if str(node_id) in fallback_nodes:
                continue
            node_dim = per_node_dims[node_id]  # type: ignore[arg-type]
            
            operator = self.operator_factory.create_operator(function_name, node_data, node_dim)  # type: ignore[arg-type]
            if operator:
                # Preserve trace identity on each mapped operator so downstream
                # validators can execute/inspect in transformed-graph order.
                try:
                    operator.trace_node_id = str(node_id)
                    operator.trace_function_name = str(function_name)
                    operator.trace_stage = str(node_data.get('stage', '')) if node_data.get('stage') is not None else None
                    # node ids are typically suffixed as "...#<call_index>".
                    m = re.search(r"#(\d+)$", str(node_id))
                    operator.trace_call_index = int(m.group(1)) if m else None
                except Exception:
                    pass
                operators_graph.nodes.add(operator)
                node_mapping[node_id] = operator
                
                # Collect characteristics
                characteristics['total_flops'] += operator.get_num_ops()
                characteristics['total_memory_bytes'] += (operator.input_a + operator.output) * 4
                characteristics['operator_types'][operator.op_type] = characteristics['operator_types'].get(operator.op_type, 0) + 1
                characteristics['realistic_operators'].append({
                    'original_id': node_id,
                    'function_name': function_name,
                    'op_type': operator.op_type,
                    'input_elements': operator.input_a,
                    'output_elements': operator.output,
                    'flop_count': operator.get_num_ops(),
                    'memory_bytes': (operator.input_a + operator.output) * 4,
                    'dim': node_dim,
                })
        
        # Helpers to keep graph acyclic and ordered by taxonomy
        def _taxonomy(op_obj):
            return self._map_operator_to_taxonomy(op_obj)

        _ORDER = {
            'SAMPLING': 0,
            'ENCODING': 1,
            'FIELD_COMPUTATION': 2,
            'BLENDING': 3,
        }

        def _order_of(op_obj) -> int:
            return _ORDER.get(_taxonomy(op_obj), 99)

        def _would_create_cycle(src, dst) -> bool:
            try:
                # DFS from dst to see if src is reachable
                seen = set()
                stack = [dst]
                while stack:
                    cur = stack.pop()
                    if cur is src:
                        return True
                    for ch in getattr(cur, 'children', []) or []:
                        if ch not in seen:
                            seen.add(ch)
                            stack.append(ch)
            except Exception:
                pass
            return False

        # Allowed cross-stage backward edges for coarse→fine NeRF pipelines:
        # The coarse stage's renderers (BLENDING) feed into the fine stage's
        # importance sampler (SAMPLING), and density heads (FIELD_COMPUTATION)
        # may also connect directly.
        _ALLOWED_BACKWARD = {
            ('BLENDING', 'SAMPLING'),
            ('FIELD_COMPUTATION', 'SAMPLING'),
        }

        def _safe_connect(src, dst, allow_cross_stage=False) -> bool:
            try:
                if src is dst:
                    return False
                src_tax = _taxonomy(src)
                dst_tax = _taxonomy(dst)
                if _order_of(dst) < _order_of(src):
                    if not allow_cross_stage and (src_tax, dst_tax) not in _ALLOWED_BACKWARD:
                        return False
                # Prevent cycles
                if _would_create_cycle(src, dst):
                    return False
                src.add_child(dst)
                return True
            except Exception:
                return False

        # Wire dependencies based on traced edges
        for edge in dag_data.get('edges', []):
            if len(edge) >= 2:
                src_id, dst_id = edge[0], edge[1]
                if src_id in node_mapping and dst_id in node_mapping:
                    src_op = node_mapping[src_id]
                    dst_op = node_mapping[dst_id]
                    _safe_connect(src_op, dst_op)
        
        # Augment missing ENCODING -> FIELD_COMPUTATION edges when FIELD nodes have no predecessors
        # and more generally, connect any zero in-degree operator to the nearest mapped predecessor
        try:
            # Build operator indegree
            op_indegree = {}
            graph_nodes = list(operators_graph.nodes)
            for op in graph_nodes:
                op_indegree[op] = 0
            for op in graph_nodes:
                for ch in getattr(op, 'children', []) or []:
                    op_indegree[ch] = op_indegree.get(ch, 0) + 1
 
            def _op_taxonomy(op_obj):
                return self._map_operator_to_taxonomy(op_obj)
 
            def _bfs_find_predecessor(start_traced_id: str, prefer_encoding: bool = False):
                """BFS upstream in the original traced DAG to find a mapped predecessor.
                If prefer_encoding is True, first try to find ENCODING; otherwise accept any mapped op.
                If no preferred found, fall back to any mapped predecessor.
                """
                visited = set()
                queue = list(preds.get(start_traced_id, []))
                fallback_found = None
                while queue:
                    cur = queue.pop(0)
                    if cur in visited:
                        continue
                    visited.add(cur)
                    m = node_mapping.get(cur)
                    if m is not None:
                        if prefer_encoding and _op_taxonomy(m) == 'ENCODING':
                            return m
                        if fallback_found is None:
                            fallback_found = m
                    # continue walking upstream
                    for pp in preds.get(cur, []) or []:
                        if pp not in visited:
                            queue.append(pp)
                return fallback_found
 
            augmented_encoding_fc = 0
            augmented_general = 0
 
            for traced_id, op_obj in node_mapping.items():
                # Only consider nodes currently with zero in-degree in the operator graph
                if op_indegree.get(op_obj, 0) > 0:
                    continue
                taxonomy = _op_taxonomy(op_obj)
 
                found = None
                if taxonomy == 'FIELD_COMPUTATION':
                    # Prefer an ENCODING predecessor; fall back to any mapped predecessor
                    found = _bfs_find_predecessor(traced_id, prefer_encoding=True)
                    if found is not None and _safe_connect(found, op_obj):
                        augmented_encoding_fc += 1
                        continue
                # For all other zero in-degree nodes (or FC fallback), connect to any mapped predecessor
                found = _bfs_find_predecessor(traced_id, prefer_encoding=False)
                if found is not None and _safe_connect(found, op_obj):
                    augmented_general += 1
 
            if augmented_encoding_fc or augmented_general:
                print(
                    f"   🔗 Augmented zero-in-degree ops: ENCODING→FIELD={augmented_encoding_fc}, general preds added={augmented_general}"
                )
        except Exception:
            pass
 
        # Shape-based fallback: for any remaining zero in-degree FIELD_COMPUTATION,
        # connect to nearest ENCODING (preferred) or SAMPLING op with matching (B,N)
        try:
            graph_nodes = list(operators_graph.nodes)
            # Recompute indegree after previous augmentation
            op_indegree = {op: 0 for op in graph_nodes}
            for op in graph_nodes:
                for ch in getattr(op, 'children', []) or []:
                    op_indegree[ch] = op_indegree.get(ch, 0) + 1
 
            # Build (B,N) -> candidate lists for ENCODING and SAMPLING
            def _bn(shape):
                try:
                    if isinstance(shape, (list, tuple)) and len(shape) >= 2:
                        return (int(shape[0]), int(shape[1]))
                except Exception:
                    return None
                return None
 
            encoding_by_bn = {}
            sampling_by_bn = {}
 
            for op in graph_nodes:
                try:
                    t = self._map_operator_to_taxonomy(op)
                    # Prefer output shape; fallback to input if output unavailable
                    out_shape = None
                    in_shapes = None
                    if hasattr(op, 'get_output_tensor_shape'):
                        try:
                            out_shape = op.get_output_tensor_shape()
                        except Exception:
                            out_shape = None
                    if hasattr(op, 'get_input_tensor_shapes'):
                        try:
                            in_shapes = op.get_input_tensor_shapes() or []
                        except Exception:
                            in_shapes = []
                    key = _bn(out_shape) or (_bn(in_shapes[0]) if in_shapes else None)
                    if not key:
                        continue
                    if t == 'ENCODING':
                        encoding_by_bn.setdefault(key, []).append(op)
                    elif t == 'SAMPLING':
                        sampling_by_bn.setdefault(key, []).append(op)
                except Exception:
                    continue
 
            added_by_shape = 0
            for op in graph_nodes:
                try:
                    if op_indegree.get(op, 0) > 0:
                        continue
                    if self._map_operator_to_taxonomy(op) != 'FIELD_COMPUTATION':
                        continue
                    # FC input shape drives linkage
                    in_shapes = []
                    if hasattr(op, 'get_input_tensor_shapes'):
                        try:
                            in_shapes = op.get_input_tensor_shapes() or []
                        except Exception:
                            in_shapes = []
                    key = _bn(in_shapes[0]) if in_shapes else None
                    if not key:
                        continue
                    src = None
                    cands = encoding_by_bn.get(key) or []
                    if cands:
                        src = cands[0]
                    else:
                        cands = sampling_by_bn.get(key) or []
                        if cands:
                            src = cands[0]
                    if src is not None and _safe_connect(src, op):
                        added_by_shape += 1
                except Exception:
                    continue
            if added_by_shape:
                print(f"   🔗 Shape-based augmentation added {added_by_shape} ENCODING/SAMPLING→FIELD links by (B,N) match")
        except Exception:
            pass
 
        # Forward augmentation: for zero out-degree FIELD_COMPUTATION nodes, find nearest
        # downstream mapped successors through helper/wrapper nodes and reconnect.
        # This complements the zero in-degree predecessor augmentation above.
        try:
            graph_nodes = list(operators_graph.nodes)
            op_outdegree = {op: 0 for op in graph_nodes}
            for op in graph_nodes:
                for ch in getattr(op, 'children', []) or []:
                    op_outdegree[op] = op_outdegree.get(op, 0) + 1

            def _op_tax(op_obj):
                return self._map_operator_to_taxonomy(op_obj)

            def _bfs_find_successor(start_traced_id, preferred_types):
                visited = set()
                queue = list(succs.get(start_traced_id, []))
                fallback_found = None
                while queue:
                    cur = queue.pop(0)
                    if cur in visited:
                        continue
                    visited.add(cur)
                    m = node_mapping.get(cur)
                    if m is not None:
                        t = _op_tax(m)
                        if t in preferred_types:
                            return m
                        if fallback_found is None:
                            fallback_found = m
                    for ss in succs.get(cur, []) or []:
                        if ss not in visited:
                            queue.append(ss)
                return fallback_found

            augmented_forward = 0
            augmented_coarse_fine = 0
            for traced_id, op_obj in node_mapping.items():
                if _op_tax(op_obj) != 'FIELD_COMPUTATION':
                    continue
                if op_outdegree.get(op_obj, 0) > 0:
                    continue
                # Prefer to terminate compute chains at blending/rendering.
                found = _bfs_find_successor(traced_id, preferred_types=('BLENDING',))
                if found is not None and _safe_connect(op_obj, found):
                    augmented_forward += 1
                    continue
                # Coarse→fine: density heads feed into PDFSampler via get_weights.
                # Look for a SAMPLING successor (importance sampling).
                found = _bfs_find_successor(traced_id, preferred_types=('SAMPLING',))
                if found is not None and _safe_connect(op_obj, found, allow_cross_stage=True):
                    augmented_coarse_fine += 1
                    continue
                # Fallback to any downstream mapped node.
                found = _bfs_find_successor(traced_id, preferred_types=('FIELD_COMPUTATION', 'ENCODING'))
                if found is not None and _safe_connect(op_obj, found):
                    augmented_forward += 1

            if augmented_forward or augmented_coarse_fine:
                print(f"   🔗 Forward augmentation: {augmented_forward} FC→downstream, {augmented_coarse_fine} FC→SAMPLING (coarse→fine)")
            else:
                # Count remaining dangling FC for diagnostics
                dangling_fc = sum(1 for op in graph_nodes if _op_tax(op) == 'FIELD_COMPUTATION' and op_outdegree.get(op, 0) == 0)
                if dangling_fc:
                    print(f"   ⚠️ Forward augmentation: {dangling_fc} FIELD_COMPUTATION ops remain without children")
        except Exception as _fwd_err:
            import traceback
            print(f"   ⚠️ Forward augmentation failed: {_fwd_err}")
            traceback.print_exc()

        # Core-only edge projection: contract wrappers/aux nodes to enforce
        # SAMPLING/ENCODING -> FIELD_COMPUTATION -> BLENDING dependencies
        try:
            graph_nodes = list(operators_graph.nodes)
            # Build reverse map: operator instance -> traced id(s)
            op_to_traced = {}
            for tid, op in node_mapping.items():
                op_to_traced.setdefault(op, []).append(tid)
 
            def _taxonomy(op_obj) -> str:
                return self._map_operator_to_taxonomy(op_obj)
 
            # Helper: (B,N) key from operator shapes if available
            def _bn_of_op(op_obj):
                try:
                    if hasattr(op_obj, 'get_input_tensor_shapes'):
                        ins = op_obj.get_input_tensor_shapes() or []
                        if ins and isinstance(ins[0], (list, tuple)) and len(ins[0]) >= 2:
                            return (int(ins[0][0]), int(ins[0][1]))
                except Exception:
                    pass
                try:
                    dim = getattr(op_obj, 'dim', None)
                    if isinstance(dim, (list, tuple)) and len(dim) >= 2:
                        return (int(dim[0]), int(dim[1]))
                except Exception:
                    pass
                return None
 
            # For each core operator, add edges from nearest upstream core producers
            added_core_edges = 0
            for op in graph_nodes:
                t = _taxonomy(op)
                if t not in ('SAMPLING','ENCODING','FIELD_COMPUTATION','BLENDING'):
                    continue
                # Collect candidate traced ids for this operator
                tids = op_to_traced.get(op, [])
                if not tids:
                    continue
                # Determine preferred upstream type(s)
                if t == 'FIELD_COMPUTATION':
                    preferred = ('ENCODING','SAMPLING')
                elif t == 'BLENDING':
                    preferred = ('FIELD_COMPUTATION',)
                else:
                    preferred = tuple()
                if not preferred:
                    continue
                # For each traced id mapped to this operator, BFS upstream to nearest core of preferred types
                bn_self = _bn_of_op(op)
                found_sources = []
                for tid in tids:
                    visited = set()
                    queue = list(preds.get(tid, []))
                    local_found = []
                    while queue and len(local_found) < 2:  # collect a couple to reduce fan-in
                        cur = queue.pop(0)
                        if cur in visited:
                            continue
                        visited.add(cur)
                        src_op = node_mapping.get(cur)
                        if src_op is not None:
                            tax = _taxonomy(src_op)
                            if tax in preferred:
                                if bn_self is None or _bn_of_op(src_op) is None or _bn_of_op(src_op) == bn_self:
                                    local_found.append(src_op)
                                    continue
                        # Continue walking upstream through non-core or non-preferred
                        for pp in preds.get(cur, []) or []:
                            if pp not in visited:
                                queue.append(pp)
                    found_sources.extend(local_found)
                # Add edges from sources to current operator
                for src in found_sources:
                    try:
                        if op not in getattr(src, 'children', []) and _safe_connect(src, op):
                            added_core_edges += 1
                    except Exception:
                        pass
            if added_core_edges:
                print(f"   🔗 Core-only projection added {added_core_edges} wrapper-contracted edges")
        except Exception:
            pass

        # Summary of dims present
        try:
            unique_dims = sorted(set(dim for dim in per_node_dims.values() if dim is not None))
            for (b, n) in unique_dims:
                print(f"   🎯 Dims present: {b} rays × {n} samples")
        except Exception:
            pass

        # Optional: Transitive reduction on operator edges to remove redundant parent→grandchild links
        # Example: if A→B and B→C exist, drop A→C. Keeps graph minimal for readability.
        try:
            # Build operator graph in NetworkX
            Gop = nx.DiGraph()
            graph_nodes = list(operators_graph.nodes)
            idx = {op: i for i, op in enumerate(graph_nodes)}
            for op in graph_nodes:
                Gop.add_node(idx[op])
            for op in graph_nodes:
                for ch in getattr(op, 'children', []) or []:
                    if ch in idx:
                        Gop.add_edge(idx[op], idx[ch])
            if nx.is_directed_acyclic_graph(Gop):
                tr = nx.algorithms.dag.transitive_reduction(Gop)
                # Remove edges not in transitive reduction
                to_remove = set(Gop.edges()) - set(tr.edges())
                if to_remove:
                    removed = 0
                    # Build reverse map for quick lookup
                    rev = {i: op for op, i in idx.items()}
                    for u, v in to_remove:
                        src = rev.get(u)
                        dst = rev.get(v)
                        if src is None or dst is None:
                            continue
                        try:
                            if hasattr(src, 'children') and dst in src.children:
                                src.children = [c for c in src.children if c is not dst]
                                removed += 1
                        except Exception:
                            continue
                    if removed:
                        print(f"   ✂️  Transitive reduction removed {removed} redundant edges")
        except Exception:
            # Never fail transformation due to a visualization/cleanup step
            pass
        
        print(f"   [OK] Created {len(operators_graph)} realistic operators")
        print(f"   📊 Total FLOPs: {characteristics['total_flops']:,}")
        print(f"   💾 Total Memory: {characteristics['total_memory_bytes']:,} bytes ({characteristics['total_memory_bytes']/1024/1024:.1f} MB)")
        
        return operators_graph, characteristics
    
    def _map_operator_to_taxonomy(self, operator) -> str:
        """Map /Operators instance to 4-stage unified taxonomy."""
        operator_class = type(operator).__name__
        
        # Field Sampler (SAMPLING) - sampling operations along rays or in space
        if any(term in operator_class for term in ['Sampler', 'Sample', 'FrustumCulling', 'FrustrumCulling']):
            return 'SAMPLING'
        
        # Encoding (ENCODING) - transform spatial coordinates to feature vectors
        elif any(term in operator_class for term in ['Encoding', 'Encoder', 'Hash', 'RFF', 'Positional', 'Fourier']):
            return 'ENCODING'
        
        # Blending (BLENDING) - aggregate scene properties to final pixel color
        elif any(term in operator_class for term in ['Render', 'Blend', 'Volume', 'RGB', 'Alpha', 'Composite']):
            return 'BLENDING'

        # Field Computation (FIELD_COMPUTATION) - compute scene properties (density, color)
        # Keep this after BLENDING to avoid classifying *Renderer* operators as FIELD_COMPUTATION.
        elif any(term in operator_class for term in ['MLP', 'Network', 'Field', 'Computation', 'Density', 'Color']):
            return 'FIELD_COMPUTATION'
        
        # Default fallback
        else:
            print(f"⚠️ Unknown operator class: {operator_class}, defaulting to FIELD_COMPUTATION")
            return 'FIELD_COMPUTATION'
    
    def operators_to_scheduler_ir(self, operators_graph: OperatorsGraph, node_mapping: Dict[str, Any]) -> SchedulerOperatorGraph:
        """Convert /Operators instances to Scheduler.IR format.
        Supports optional non-collapsing rendering mode via env NERARCHSIM_SPLIT_RENDERING=1
        which splits BLENDING/VOLUME_RENDERING into multiple smaller ops by samples dimension.
        """
        print(f"🔄 Converting to Scheduler.IR format...")
        split_rendering = False
        
        scheduler_graph = SchedulerOperatorGraph()
        node_list = list(operators_graph.nodes)
        
        # Map original operator index -> list of scheduler node ids
        op_idx_to_sched_ids: Dict[int, List[str]] = {}
        
        # First pass: create scheduler nodes (with optional rendering split)
        for i, operator in enumerate(node_list):
            # Prefer shape helpers if implemented; otherwise fall back to element counts
            in_shape = None
            out_shape = None
            try:
                if hasattr(operator, "get_input_tensor_shapes") and hasattr(operator, "get_output_tensor_shape"):
                    input_shapes = operator.get_input_tensor_shapes()
                    output_shape = operator.get_output_tensor_shape()
                    if input_shapes and isinstance(input_shapes[0], (list, tuple)):
                        in_shape = list(input_shapes[0])
                    if isinstance(output_shape, (list, tuple)):
                        out_shape = list(output_shape)
            except Exception:
                in_shape = None
                out_shape = None
            # Fallbacks based on element counts when shapes are unavailable
            if in_shape is None:
                in_elems = getattr(operator, "input_a", 1) or 1
                in_shape = [int(in_elems), 1]
            if out_shape is None:
                out_elems = getattr(operator, "output", 1) or 1
                out_shape = [int(out_elems), 1]
            # Ensure minimum rank 2
            if len(in_shape) < 2:
                in_shape = list(in_shape) + [1] * (2 - len(in_shape))
            if len(out_shape) < 2:
                out_shape = list(out_shape) + [1] * (2 - len(out_shape))
            taxonomy_op_type = self._map_operator_to_taxonomy(operator)
             
            # Determine how many rendering instances to create: do not replicate; keep one-to-one with instrumentation
            render_slices = 1
            # no-op: we don't split rendering here; instrumentation should provide multiple calls if any
             
            created_ids: List[str] = []
            for r in range(render_slices):
                node_id = f"op_{i}" if render_slices == 1 else f"op_{i}_r{r}"
                # For rendering replication by chunks, set shapes to [chunk_size, N, C] -> [chunk_size, C]
                inputs = [TensorDesc(shape=in_shape, dtype='float32')]
                outputs = [TensorDesc(shape=out_shape, dtype='float32')]
                scheduler_node = OperatorNode(
                    id=node_id,
                    op_type=taxonomy_op_type,
                    inputs=inputs,
                    outputs=outputs,
                    call_count=1,
                    metadata={
                        'flop_count': operator.get_num_ops(),
                        'memory_bytes': sum(t.bytes() for t in inputs) + sum(t.bytes() for t in outputs),
                        'hardware_type': map_operator_to_hardware_type(taxonomy_op_type),
                        'realistic_characteristics': True,
                        'operator_class': type(operator).__name__,
                        'input_elements': operator.input_a,
                        'output_elements': operator.output,
                        'original_op_type': operator.op_type,
                        'render_slice_index': None,
                        'render_slices': None,
                    }
                )
                # Attach MLP details when applicable
                try:
                    from operators.computation_operator import MLPOperator as _MLP
                    if isinstance(operator, _MLP):
                        scheduler_node.metadata.update({
                            'mlp_in_dim': operator.in_dim,
                            'mlp_num_layers': operator.num_layers,
                            'mlp_layer_width': operator.layer_width,
                            'mlp_out_dim': operator.out_dim,
                            'mlp_skip_connections': list(operator.skip_connections) if operator.skip_connections else [],
                            'mlp_layer_weight_shapes': operator.get_layer_weight_shapes(),
                        })
                except Exception:
                    pass
                scheduler_graph.nodes[node_id] = scheduler_node
                created_ids.append(node_id)
            op_idx_to_sched_ids[i] = created_ids
        
        # Second pass: wire dependencies
        for i, operator in enumerate(node_list):
            src_ids = op_idx_to_sched_ids.get(i, [])
            for child in getattr(operator, 'children', []) or []:
                if child in node_list:
                    j = node_list.index(child)
                    dst_ids = op_idx_to_sched_ids.get(j, [])
                    if not src_ids or not dst_ids:
                        continue
                    # If destination has multiple slices, connect to its first slice
                    dst_id = dst_ids[0]
                    for src_id in src_ids:
                        scheduler_graph.edges.append((src_id, dst_id))
            # No replication; nothing to chain
            ids = op_idx_to_sched_ids.get(i, [])
            if ids and len(ids) > 1:
                for a, b in zip(ids[:-1], ids[1:]):
                    scheduler_graph.edges.append((a, b))
        
        print(f"   [OK] Converted to Scheduler.IR with {len(scheduler_graph.nodes)} nodes")
        return scheduler_graph
    
    def analyze_transformation_impact(self, original_dag: Dict[str, Any], characteristics: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze the transformation impact."""
        original_nodes = len(original_dag.get('nodes', {}))
        transformed_nodes = len(characteristics['realistic_operators'])
        
        total_flops = characteristics['total_flops']
        total_memory_mb = characteristics['total_memory_bytes'] / (1024 * 1024)
        
        return {
            'transformation_summary': {
                'nodes_processed': f"{transformed_nodes}/{original_nodes}",
                'total_flops': f"{total_flops:,}",
                'total_memory_mb': f"{total_memory_mb:.1f} MB",
                'operator_distribution': characteristics['operator_types']
            },
            'improvement_comparison': {
                'before_transformation': {
                    'tensor_shapes': 'Generic [1]',
                    'flop_counts': '0',
                    'memory_modeling': '4 bytes per operator',
                    'performance_analysis': 'Meaningless'
                },
                'after_transformation': {
                    'tensor_shapes': 'Realistic (786K - 8M elements)',
                    'flop_counts': f'{total_flops:,}',
                    'memory_modeling': f'{total_memory_mb:.1f} MB realistic workload',
                    'performance_analysis': 'Hardware bottleneck identification'
                }
            },
            'enabled_capabilities': [
                'Roofline analysis with actual operator characteristics',
                'Memory bandwidth bottleneck identification', 
                'Realistic accelerator design feedback',
                'Accurate neural rendering performance prediction',
                'Hardware-software co-design optimization'
            ]
        }


def _operator_validation_artifacts(
    operators_graph: OperatorsGraph,
    out_prefix: str,
    verbose: bool = True,
) -> Dict[str, str]:
    """
    Generate validation artifacts for transformed operator graph:
      - <out_prefix>_graph.png
      - <out_prefix>_summary.json
      - <out_prefix>_summary.md
    """
    artifacts: Dict[str, str] = {}
    try:
        nodes = list(operators_graph.nodes)
        node_type_counts: Dict[str, int] = {}
        indeg: Dict[Any, int] = {n: 0 for n in nodes}
        outdeg: Dict[Any, int] = {n: 0 for n in nodes}
        for n in nodes:
            t = type(n).__name__
            node_type_counts[t] = node_type_counts.get(t, 0) + 1
            children = list(getattr(n, "children", []) or [])
            outdeg[n] = len(children)
            for c in children:
                indeg[c] = indeg.get(c, 0) + 1

        def _tax(op_obj) -> str:
            cn = type(op_obj).__name__
            if any(term in cn for term in ['Sampler', 'Sample', 'FrustumCulling', 'FrustrumCulling']):
                return 'SAMPLING'
            if any(term in cn for term in ['Encoding', 'Encoder', 'Hash', 'RFF', 'Positional', 'Fourier']):
                return 'ENCODING'
            if any(term in cn for term in ['Render', 'Blend', 'Volume', 'RGB', 'Alpha', 'Composite']):
                return 'BLENDING'
            if any(term in cn for term in ['MLP', 'Network', 'Field', 'Computation', 'Density', 'Color']):
                return 'FIELD_COMPUTATION'
            return 'FIELD_COMPUTATION'

        sinks_by_tax: Dict[str, int] = {}
        roots_by_tax: Dict[str, int] = {}
        for n in nodes:
            t = _tax(n)
            if outdeg.get(n, 0) == 0:
                sinks_by_tax[t] = sinks_by_tax.get(t, 0) + 1
            if indeg.get(n, 0) == 0:
                roots_by_tax[t] = roots_by_tax.get(t, 0) + 1

        summary = {
            "total_nodes": len(nodes),
            "node_type_counts": node_type_counts,
            "roots_by_taxonomy": roots_by_tax,
            "sinks_by_taxonomy": sinks_by_tax,
            "validation_checks": {
                "has_blending_sinks": bool(sinks_by_tax.get("BLENDING", 0) > 0),
                "field_computation_sink_count": int(sinks_by_tax.get("FIELD_COMPUTATION", 0)),
            },
        }

        graph_png = f"{out_prefix}_graph.png"
        operators_graph.plot_graph(
            title="Operator Graph Validation View",
            save_path=graph_png
        )
        artifacts["validation_graph_png"] = graph_png

        summary_json = f"{out_prefix}_summary.json"
        with open(summary_json, "w") as f:
            json.dump(summary, f, indent=2)
        artifacts["validation_summary_json"] = summary_json

        md = [
            "# Operator Graph Validation",
            "",
            f"- Total operators: {summary['total_nodes']}",
            "",
            "## Node Type Counts",
        ]
        for k, v in sorted(node_type_counts.items(), key=lambda kv: kv[1], reverse=True):
            md.append(f"- {k}: {v}")
        md += [
            "",
            "## Sinks by Taxonomy",
        ]
        for k, v in sorted(sinks_by_tax.items(), key=lambda kv: kv[1], reverse=True):
            md.append(f"- {k}: {v}")
        md += [
            "",
            "## Checks",
            f"- Has BLENDING sinks: {summary['validation_checks']['has_blending_sinks']}",
            f"- FIELD_COMPUTATION sink count: {summary['validation_checks']['field_computation_sink_count']}",
        ]
        summary_md = f"{out_prefix}_summary.md"
        with open(summary_md, "w") as f:
            f.write("\n".join(md) + "\n")
        artifacts["validation_summary_md"] = summary_md

        if verbose:
            print(f"   🧪 Validation graph: {graph_png}")
            print(f"   🧪 Validation summary: {summary_json}")
    except Exception as e:
        if verbose:
            print(f"   ⚠️ Validation artifact generation failed: {e}")
    return artifacts

def load_and_transform_traced_dag(dag_path: str) -> Tuple[SchedulerOperatorGraph, Dict[str, Any]]:
    """Complete pipeline: Load traced DAG → /Operators → Scheduler.IR"""
    print(f"🚀 Complete DAG Transformation Pipeline")
    print("=" * 50)
    
    # Load traced DAG
    with open(dag_path, 'rb') as f:
        dag_data = pickle.load(f)
    
    if isinstance(dag_data, nx.DiGraph):
        # Convert NetworkX to dict format
        dict_dag = {"nodes": {}, "edges": list(dag_data.edges())}
        for node_id, node_data in dag_data.nodes(data=True):
            node_data_copy = node_data.copy()
            node_data_copy['function_name'] = str(node_id)
            dict_dag["nodes"][node_id] = node_data_copy
        dag_data = dict_dag
    
    print(f"📥 Loaded traced DAG: {len(dag_data['nodes'])} nodes, {len(dag_data.get('edges', []))} edges")
    
    # Transform to /Operators
    integration = DAGToOperatorsIntegration()
    operators_graph, characteristics = integration.transform_dag_to_operators(dag_data)

    # Optional validation artifacts to verify instrumentation/transformation flow.
    # Enabled by setting NERARCHSIM_VALIDATE_RENDERING=1.
    try:
        validate = str(os.environ.get("NERARCHSIM_VALIDATE_RENDERING", "0")).strip().lower() in ("1", "true", "yes", "on")
        if validate:
            base = Path(dag_path).resolve()
            out_prefix = str(base.parent / "operator_graph_validation")
            _operator_validation_artifacts(operators_graph, out_prefix=out_prefix, verbose=True)
    except Exception:
        pass
 
    # Persist operator graph visuals for analysis (picked up by analyze -> visuals/)
    # Disabled by default; enable via env:
    #   NERARCHSIM_PLOT_OPERATORS=1           -> plots coarse operator graph
    #   NERARCHSIM_PLOT_FINE_OPERATORS=1      -> plots fine operator graph (can be large)
    # Optional filters:
    #   NERARCHSIM_PLOT_INCLUDE="MLP,Encoding"   (substring match on op_type or class)
    #   NERARCHSIM_PLOT_MAX_NODES=200
    try:
        def _env(new_key: str, old_key: str, default: str = "0") -> str:
            v = os.environ.get(new_key)
            if v is None or str(v).strip() == "":
                v = os.environ.get(old_key, default)
            return str(v)
        def _should(new_key: str, old_key: str) -> bool:
            v = _env(new_key, old_key, "0").strip().lower()
            return v in ("1", "true", "yes", "on")

        include_filter = _env("NERARCHSIM_PLOT_INCLUDE", "RENDERSIM_PLOT_INCLUDE", "").strip()
        include_terms = [t.strip().lower() for t in include_filter.split(',') if t.strip()] if include_filter else []
        max_nodes_env = _env("NERARCHSIM_PLOT_MAX_NODES", "RENDERSIM_PLOT_MAX_NODES", "")
        try:
            max_nodes = int(max_nodes_env) if max_nodes_env else None
        except Exception:
            max_nodes = None

        def _filtered_nodes(all_nodes):
            nodes = list(all_nodes)
            if include_terms:
                def _match(op):
                    try:
                        name = getattr(op, 'op_type', '') or type(op).__name__
                        name = str(name).lower()
                        return any(term in name for term in include_terms)
                    except Exception:
                        return False
                nodes = [op for op in nodes if _match(op)]
            if max_nodes is not None and len(nodes) > max_nodes:
                nodes = nodes[:max_nodes]
            return nodes

        # Default: emit SVG alongside PNG unless user explicitly disabled
        try:
            if os.environ.get("NERARCHSIM_PLOT_SVG") in (None, "") and os.environ.get("RENDERSIM_PLOT_SVG") in (None, ""):
                os.environ["NERARCHSIM_PLOT_SVG"] = "1"
        except Exception:
            pass

        if _should("NERARCHSIM_PLOT_OPERATORS", "RENDERSIM_PLOT_OPERATORS"):
            # Coarse graph is small; no filtering needed
            operators_graph.plot_graph(title="Operator Graph", save_path="operator_graph.png")
        if _should("NERARCHSIM_PLOT_FINE_OPERATORS", "RENDERSIM_PLOT_FINE_OPERATORS"):
            from Operators.utils.operator_graph import FineOperatorGraph
            fine = FineOperatorGraph()
            fine.nodes.extend(_filtered_nodes(operators_graph.nodes))
            fine.plot_graph(title="Fine Operator Graph", save_path="operator_graph_fine.png")
    except Exception:
        pass
 
    # Convert to Scheduler.IR
    scheduler_graph = integration.operators_to_scheduler_ir(operators_graph, {})
    
    # Analyze impact
    impact = integration.analyze_transformation_impact(dag_data, characteristics)
    
    print(f"\n🎉 Transformation Complete!")
    print(f"   {impact['transformation_summary']['nodes_processed']} operators with realistic characteristics")
    print(f"   {impact['transformation_summary']['total_flops']} FLOPs vs. 0 before")
    print(f"   {impact['transformation_summary']['total_memory_mb']} vs. negligible before")
    
    return scheduler_graph, impact

if __name__ == "__main__":
    # When running this module directly, enable validation artifacts by default.
    os.environ.setdefault("NERARCHSIM_VALIDATE_RENDERING", "1")
    # Test the complete integration
    if len(sys.argv) > 1:
        dag_path = sys.argv[1]
    else:
        dag_path = "execution_dag.pkl"
    
    if Path(dag_path).exists():
        try:
            scheduler_graph, impact = load_and_transform_traced_dag(dag_path)
            
            print(f"\n🧪 Integration Test Results:")
            print(f"[OK] Scheduler.IR nodes: {len(scheduler_graph.nodes)}")
            print(f"[OK] Scheduler.IR edges: {len(scheduler_graph.edges)}")
            print(f"[OK] Ready for realistic neural rendering scheduling!")
            
        except Exception as e:
            print(f"[ERROR] Integration failed: {e}")
            import traceback
            traceback.print_exc()
    else:
        print(f"[ERROR] DAG file not found: {dag_path}")
        print(f"💡 Generate with: python nerfstudio/nerfstudio/scripts/eval.py ... --enable-trace")