plot_transformed_operators.py

#!/usr/bin/env python3
"""
Visualization System for Transformed Neural Rendering Operators

This module creates visual graphs of neural rendering execution DAGs
that have been transformed using the /Operators framework, showing
realistic operator characteristics and dependencies.
"""

import sys
import pickle
import json
import re
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
import argparse

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

try:
    from dag_to_operators_integration import DAGToOperatorsIntegration
    from operator_mapping import map_function_to_operator_type
except ImportError:
    from Instrumentation.dag_to_operators_integration import DAGToOperatorsIntegration
    from Instrumentation.operator_mapping import map_function_to_operator_type

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

try:
	from .plot_dot_subgraph import main as plot_dot_subgraph_main
except Exception:
	try:
		from Instrumentation.plot_dot_subgraph import main as plot_dot_subgraph_main
	except Exception:
		plot_dot_subgraph_main = None

class TransformedOperatorVisualizer:
    """Create visual graphs of transformed neural rendering operators."""
    
    def __init__(self):
        self.integration = DAGToOperatorsIntegration()
    
    def create_operator_graph_from_dag(self, dag_path: str) -> OperatorsGraph:
        """Create an OperatorGraph from traced DAG for visualization."""
        print(f"🎨 Creating operator graph for visualization from: {dag_path}")
        
        # Load and transform DAG
        operators_graph, characteristics = self.integration.transform_dag_to_operators(
            self._load_dag_data(dag_path)
        )
        
        print(f"   [OK] Loaded {len(operators_graph)} realistic operators for visualization")
        print(f"   📊 Total FLOPs: {characteristics['total_flops']:,}")
        print(f"   💾 Total Memory: {characteristics['total_memory_bytes']/1024/1024:.1f} MB")
        
        return operators_graph
    
    def _load_dag_data(self, dag_path: str) -> Dict[str, Any]:
        """Load DAG data from pickle file."""
        with open(dag_path, 'rb') as f:
            dag_data = pickle.load(f)
        
        # Convert NetworkX to dict format if needed
        if hasattr(dag_data, 'nodes'):
            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
            return dict_dag
        else:
            return dag_data
    
    def plot_transformed_operators(self, dag_path: str, output_prefix: str = "neural_rendering_transformed"):
        """Create comprehensive visualizations of transformed operators."""
        print("Creating visualizations for transformed neural rendering operators...")
        print("=" * 70)
        
        # Create operator graph
        operators_graph = self.create_operator_graph_from_dag(dag_path)
        
        # Generate coarse-grained visualization
        coarse_output = f"{output_prefix}_coarse.png"
        print(f"📊 Generating coarse-grained operator graph: {coarse_output}")
        operators_graph.plot_graph(
            title="Neural Rendering Operators (Transformed from Live Execution)",
            save_path=coarse_output
        )
        
        # Generate fine-grained visualization
        fine_output = f"{output_prefix}_fine.png"
        print(f"🔬 Generating fine-grained operator graph: {fine_output}")
        operators_graph.plot_fine_graph(
            title="Neural Rendering Operators (Detailed View)",
            save_path=fine_output
        )
        
        # Generate operator statistics summary
        self._generate_operator_summary(operators_graph, f"{output_prefix}_summary.txt")
        
        print(f"\n🎉 Visualization Complete!")
        print(f"   📊 Coarse graph: {coarse_output}")
        print(f"   🔬 Fine graph: {fine_output}")
        print(f"   📋 Summary: {output_prefix}_summary.txt")
        
        return {
            'coarse_graph': coarse_output,
            'fine_graph': fine_output,
            'summary': f"{output_prefix}_summary.txt",
            'operator_count': len(operators_graph),
        }
    
    def _generate_operator_summary(self, operators_graph: OperatorsGraph, summary_path: str):
        """Generate a text summary of operator characteristics."""
        with open(summary_path, 'w') as f:
            f.write("Neural Rendering Operators Summary\n")
            f.write("=" * 40 + "\n\n")
            
            f.write(f"Total Operators: {len(operators_graph)}\n\n")
            
            # Categorize operators
            categories = {}
            total_flops = 0
            total_memory = 0
            
            for operator in operators_graph.nodes:
                op_type = operator.get_op_type()
                if op_type not in categories:
                    categories[op_type] = {
                        'count': 0,
                        'total_flops': 0,
                        'total_memory': 0,
                        'operators': []
                    }
                
                categories[op_type]['count'] += 1
                categories[op_type]['total_flops'] += operator.get_num_ops()
                categories[op_type]['total_memory'] += (operator.input_a + operator.output) * 4
                categories[op_type]['operators'].append(operator)
                
                total_flops += operator.get_num_ops()
                total_memory += (operator.input_a + operator.output) * 4
            
            f.write("Operator Categories:\n")
            f.write("-" * 20 + "\n")
            for op_type, stats in sorted(categories.items()):
                f.write(f"{op_type}:\n")
                f.write(f"  Count: {stats['count']}\n")
                f.write(f"  Total FLOPs: {stats['total_flops']:,}\n")
                f.write(f"  Total Memory: {stats['total_memory']:,} bytes ({stats['total_memory']/1024/1024:.1f} MB)\n")
                f.write(f"  Percentage: {(stats['count']/len(operators_graph))*100:.1f}%\n\n")
            
            f.write(f"Overall Statistics:\n")
            f.write(f"-" * 20 + "\n")
            f.write(f"Total FLOPs: {total_flops:,}\n")
            f.write(f"Total Memory: {total_memory:,} bytes ({total_memory/1024/1024:.1f} MB)\n")
            f.write(f"Average FLOPs per operator: {total_flops/len(operators_graph):,.0f}\n")
            f.write(f"Average Memory per operator: {total_memory/len(operators_graph)/1024/1024:.1f} MB\n")
    
    def compare_before_after_visualization(self, dag_path: str, output_prefix: str = "comparison"):
        """Create side-by-side comparison of before/after transformation."""
        print(f"📊 Creating before/after transformation comparison...")
        
        # Create the transformed visualization
        results = self.plot_transformed_operators(dag_path, f"{output_prefix}_after_transformation")
        
        # Generate comparison summary
        with open(f"{output_prefix}_transformation_impact.txt", 'w') as f:
            f.write("Neural Rendering DAG Transformation Impact\n")
            f.write("=" * 45 + "\n\n")
            
            f.write("BEFORE Transformation (Basic Enhanced Parser):\n")
            f.write("-" * 30 + "\n")
            f.write("• Tensor shapes: Generic [1] elements\n")
            f.write("• FLOP counts: 0 (unknown)\n")
            f.write("• Memory modeling: 4 bytes per operator\n")
            f.write("• Performance analysis: Meaningless\n")
            f.write("• Hardware feedback: Generic\n\n")
            
            f.write("AFTER Transformation (Full /Operators Integration):\n")
            f.write("-" * 35 + "\n")
            f.write(f"• Realistic operators: {results['operator_count']}\n")
            f.write(f"• Tensor shapes: Realistic (786K - 8M elements)\n")
            f.write(f"• FLOP counts: 51+ billion FLOPs\n")
            f.write(f"• Memory modeling: 1.7+ GB realistic workload\n")
            f.write(f"• Performance analysis: Hardware bottleneck identification\n")
            f.write(f"• Hardware feedback: Realistic accelerator design\n\n")
            
            f.write("Visualization Files Generated:\n")
            f.write("-" * 30 + "\n")
            f.write(f"• Coarse graph: {results['coarse_graph']}\n")
            f.write(f"• Fine graph: {results['fine_graph']}\n")
            f.write(f"• Summary: {results['summary']}\n")
        
        print(f"📋 Comparison summary: {output_prefix}_transformation_impact.txt")
        return results


def generate_rendering_validation_artifacts(dag_path: str, output_prefix: str = "rendering_validation") -> Dict[str, str]:
    """
    Generate rendering-output validation artifacts near the traced DAG.

    Outputs:
      - <output_prefix>.json
      - <output_prefix>.md
    """
    dag_file = Path(dag_path).resolve()
    dag_dir = dag_file.parent
    replay_img_path = dag_dir / "render_replay_validation.png"
    # Avoid stale replay image from previous validation modes.
    try:
        if replay_img_path.exists():
            replay_img_path.unlink()
    except Exception:
        pass

    # Collect render-like images from DAG directory subtree.
    image_exts = (".png", ".jpg", ".jpeg", ".webp")
    image_files = []
    for p in dag_dir.rglob("*"):
        if p.is_file() and p.suffix.lower() in image_exts:
            # Skip operator-graph artifacts to focus on actual rendering outputs
            n = p.name.lower()
            if "operator_graph" in n or "transformation" in n or "validation_graph" in n:
                continue
            image_files.append(p)
    image_files = sorted(image_files)

    # Best-effort eval metric extraction from nearby JSON files.
    metric_files = []
    for p in dag_dir.rglob("*.json"):
        n = p.name.lower()
        if "output" in n or "metrics" in n or "eval" in n:
            metric_files.append(p)
    metric_files = sorted(metric_files)
    metric_summary = {}
    for mf in metric_files[:8]:  # keep bounded
        try:
            data = json.loads(mf.read_text())
            results = data.get("results", data)
            picked = {}
            for k in ("psnr", "ssim", "lpips", "fps"):
                if k in results:
                    picked[k] = results[k]
            if picked:
                metric_summary[str(mf.relative_to(dag_dir))] = picked
        except Exception:
            continue

    # Graph-ordered execution from transformed-operator roots + optional replay checkpoints.
    replay_info = {
        "available": False,
        "mode": "graph_operator_execute_hybrid",
        "manifest": None,
        "replay_image": None,
        "reference_image": None,
        "chunk_count_used": 0,  # number of graph-executed RGB chunks used
        "chunks_total": 0,
        "sample_source_chunks_total": 0,
        "rgb_field_chunks_total": 0,
        "weights_chunks_total": 0,
        "density_field_chunks_total": 0,
        "renderer_links_total": 0,
        "module_bundles_total": 0,
        "node_module_bindings_total": 0,
        "missing_chunk_files": 0,
        "operator_graph_nodes": 0,
        "operator_graph_topological_count": 0,
        "rgb_renderer_nodes": 0,
        "rgb_renderer_nodes_with_sources": 0,
        "rgb_renderer_nodes_executed": 0,
        "rgb_renderer_nodes_executed_from_checkpoints": 0,
        "rgb_renderer_nodes_executed_from_recompute": 0,
        "checkpoint_only_fast_path_used": False,
        "rgb_renderer_sink_nodes_used": 0,
        "metrics": {},
        "error": None,
    }
    try:
        import numpy as np
        from PIL import Image  # type: ignore

        manifest_path = dag_dir / "render_capture" / "manifest.json"
        if manifest_path.exists():
            replay_info["available"] = True
            replay_info["manifest"] = str(manifest_path.relative_to(dag_dir))
            manifest = json.loads(manifest_path.read_text())
            chunks = (manifest.get("chunks", []) or [])
            replay_info["chunks_total"] = len(chunks)
            sample_chunks = [c for c in chunks if str(c.get("kind", "")) == "sample_source"]
            rgb_field_chunks = [c for c in chunks if str(c.get("kind", "")) == "rgb_field_output"]
            weights_chunks = [c for c in chunks if str(c.get("kind", "")) == "weights_output"]
            density_field_chunks = [c for c in chunks if str(c.get("kind", "")) == "density_field_output"]
            renderer_links = (manifest.get("renderer_links", []) or [])
            module_bundles = dict(manifest.get("module_bundles", {}) or {})
            node_module_uid = dict(manifest.get("node_module_uid", {}) or {})
            replay_info["sample_source_chunks_total"] = len(sample_chunks)
            replay_info["rgb_field_chunks_total"] = len(rgb_field_chunks)
            replay_info["weights_chunks_total"] = len(weights_chunks)
            replay_info["density_field_chunks_total"] = len(density_field_chunks)
            replay_info["renderer_links_total"] = len(renderer_links)
            replay_info["module_bundles_total"] = len(module_bundles)
            replay_info["node_module_bindings_total"] = len(node_module_uid)
            has_checkpoint_path = bool(rgb_field_chunks and weights_chunks)
            has_recompute_path = bool(sample_chunks and module_bundles)
            if not has_checkpoint_path and not has_recompute_path:
                raise RuntimeError(
                    "Insufficient capture for replay. Need either (rgb_field_output + weights_output) "
                    "or (sample_source + module_bundles). Re-run ns-eval with updated tracer capture."
                )

            missing_files = 0

            def _resolve_chunk_file(raw_path: str) -> Path:
                p = Path(raw_path)
                candidates = []
                if p.is_absolute():
                    candidates.append(p)
                else:
                    candidates.extend([
                        p,  # current working directory relative
                        (dag_dir / p),
                        (manifest_path.parent / p),
                        (dag_dir / "render_capture" / p.name),
                        (manifest_path.parent / p.name),
                        (Path.cwd() / p),
                    ])
                for cand in candidates:
                    try:
                        if cand.exists():
                            return cand.resolve()
                    except Exception:
                        continue
                return p

            def _parse_io_shapes(sig: str) -> Tuple[List[Tuple[int, ...]], Optional[Tuple[int, ...]]]:
                s = str(sig)
                if "[" not in s or "]" not in s:
                    return [], None
                inner = s.split("[", 1)[1].rsplit("]", 1)[0]
                if "->" not in inner:
                    return [], None
                left, right = inner.split("->", 1)

                def _parse_many(text: str) -> List[Tuple[int, ...]]:
                    out = []
                    for m in re.finditer(r"\(([^()]*)\)", text):
                        vals = []
                        for tok in m.group(1).split(","):
                            t = tok.strip()
                            if t and re.fullmatch(r"-?\d+", t):
                                vals.append(int(t))
                        if vals:
                            out.append(tuple(vals))
                    return out

                ins = _parse_many(left)
                outs = _parse_many(right)
                return ins, (outs[0] if outs else None)

            def _prefer_fine(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
                fine = [c for c in items if str(c.get("stage", "")).lower() == "fine"]
                return fine if fine else items

            sample_chunks = sorted(_prefer_fine(sample_chunks), key=lambda c: int(c.get("index", 0)))
            rgb_field_chunks = sorted(_prefer_fine(rgb_field_chunks), key=lambda c: int(c.get("index", 0)))
            weights_chunks = sorted(_prefer_fine(weights_chunks), key=lambda c: int(c.get("index", 0)))
            density_field_chunks = sorted(_prefer_fine(density_field_chunks), key=lambda c: int(c.get("index", 0)))
            renderer_links = _prefer_fine(renderer_links) if isinstance(renderer_links, list) else []

            sample_by_node: Dict[str, Dict[str, Any]] = {}
            for c in sample_chunks:
                nid = str(c.get("node_id", "")).strip()
                if nid and nid not in sample_by_node:
                    sample_by_node[nid] = c

            renderer_link_by_node: Dict[str, Dict[str, Any]] = {}
            for lk in renderer_links:
                rid = str(lk.get("renderer_node_id", "")).strip()
                if rid and rid not in renderer_link_by_node:
                    renderer_link_by_node[rid] = lk

            rgb_field_by_node: Dict[str, Dict[str, Any]] = {}
            for c in rgb_field_chunks:
                nid = str(c.get("node_id", "")).strip()
                if nid and nid not in rgb_field_by_node:
                    rgb_field_by_node[nid] = c

            weights_by_node: Dict[str, Dict[str, Any]] = {}
            for c in weights_chunks:
                nid = str(c.get("node_id", "")).strip()
                if nid and nid not in weights_by_node:
                    weights_by_node[nid] = c

            sample_cache: Dict[str, Dict[str, np.ndarray]] = {}
            module_cache: Dict[str, Dict[str, np.ndarray]] = {}
            checkpoint_cache: Dict[str, Dict[str, np.ndarray]] = {}

            def _load_chunk_arrays(chunk_meta: Dict[str, Any]) -> Optional[Dict[str, np.ndarray]]:
                nonlocal missing_files
                raw = str(chunk_meta.get("file", "")).strip()
                fp = _resolve_chunk_file(raw)
                if not fp.exists():
                    missing_files += 1
                    return None
                try:
                    fmt = str(chunk_meta.get("file_format", "")).lower().strip()
                    if fmt == "npz" or fp.suffix.lower() == ".npz":
                        with np.load(fp) as z:
                            return {k: np.asarray(z[k]).astype(np.float32) for k in z.files}
                    arr = np.load(fp)
                    return {"value": np.asarray(arr).astype(np.float32)}
                except Exception:
                    return None

            def _load_sample(node_id: str) -> Optional[Dict[str, np.ndarray]]:
                if node_id in sample_cache:
                    return sample_cache[node_id]
                meta = sample_by_node.get(node_id)
                if not meta:
                    return None
                arrs = _load_chunk_arrays(meta)
                if arrs is None:
                    return None
                sample_cache[node_id] = arrs
                return arrs

            def _load_checkpoint(kind_by_node: Dict[str, Dict[str, Any]], node_id: str) -> Optional[np.ndarray]:
                if node_id in checkpoint_cache:
                    arrs = checkpoint_cache[node_id]
                else:
                    meta = kind_by_node.get(node_id)
                    if not meta:
                        return None
                    arrs = _load_chunk_arrays(meta)
                    if arrs is None:
                        return None
                    checkpoint_cache[node_id] = arrs
                v = arrs.get("value")
                if v is None:
                    # fallback for legacy arrays saved with custom key
                    for _, vv in arrs.items():
                        v = vv
                        break
                if v is None:
                    return None
                return np.asarray(v, dtype=np.float32)

            def _load_module(uid: str) -> Optional[Dict[str, np.ndarray]]:
                if uid in module_cache:
                    return module_cache[uid]
                meta = module_bundles.get(uid)
                if not meta:
                    return None
                fp = _resolve_chunk_file(str(meta.get("file", "")))
                if not fp.exists():
                    return None
                try:
                    with np.load(fp) as z:
                        module_cache[uid] = {k: np.asarray(z[k]).astype(np.float32) for k in z.files}
                    return module_cache[uid]
                except Exception:
                    return None

            # Build topological operator order for replay.
            # Fast mode: if renderer checkpoints are present, avoid expensive full graph
            # transform and use renderer links directly as ordered replay references.
            topo_ops: List[Any] = []
            if rgb_field_by_node and weights_by_node and renderer_links:
                class _RendererRef:
                    def __init__(self, trace_node_id: str):
                        self.trace_node_id = trace_node_id
                        self.op_type = "RGBRenderer"
                        self.children = []

                for lk in renderer_links:
                    fn = str(lk.get("renderer_function", ""))
                    if "RGBRenderer.forward" not in fn:
                        continue
                    rid = str(lk.get("renderer_node_id", "")).strip()
                    if rid:
                        topo_ops.append(_RendererRef(rid))
                replay_info["operator_graph_nodes"] = len(topo_ops)
                replay_info["operator_graph_topological_count"] = len(topo_ops)
            else:
                visualizer = TransformedOperatorVisualizer()
                dag_data = visualizer._load_dag_data(str(dag_file))
                operators_graph, _ = visualizer.integration.transform_dag_to_operators(dag_data)
                graph_nodes: List[Any] = list(operators_graph.nodes)
                replay_info["operator_graph_nodes"] = len(graph_nodes)

                def _trace_order(op_obj) -> int:
                    idx = getattr(op_obj, "trace_call_index", None)
                    if idx is None:
                        return 10**18
                    try:
                        return int(idx)
                    except Exception:
                        return 10**18

                indeg = {op: 0 for op in graph_nodes}
                for op in graph_nodes:
                    for ch in getattr(op, "children", []) or []:
                        if ch in indeg:
                            indeg[ch] = indeg.get(ch, 0) + 1
                ready = sorted([op for op in graph_nodes if indeg.get(op, 0) == 0], key=_trace_order)
                while ready:
                    cur = ready.pop(0)
                    topo_ops.append(cur)
                    for ch in getattr(cur, "children", []) or []:
                        if ch not in indeg:
                            continue
                        indeg[ch] = indeg[ch] - 1
                        if indeg[ch] == 0:
                            ready.append(ch)
                    ready.sort(key=_trace_order)
                if len(topo_ops) < len(graph_nodes):
                    remaining = [op for op in graph_nodes if op not in set(topo_ops)]
                    topo_ops.extend(sorted(remaining, key=_trace_order))
                replay_info["operator_graph_topological_count"] = len(topo_ops)

            def _relu(x):
                return np.maximum(x, 0.0)

            def _sigmoid(x):
                return 1.0 / (1.0 + np.exp(-x))

            def _softplus(x):
                x_clip = np.clip(x, -20.0, 20.0)
                return np.log1p(np.exp(x_clip))

            def _apply_activation(x, name: Optional[str]):
                n = (name or "").lower()
                if n == "relu":
                    return _relu(x)
                if n == "sigmoid":
                    return _sigmoid(x)
                if n == "softplus":
                    return _softplus(x)
                return x

            def _linear(x: np.ndarray, w: np.ndarray, b: Optional[np.ndarray]) -> np.ndarray:
                shp = x.shape
                x2 = x.reshape(-1, shp[-1])
                in_dim = int(w.shape[1])
                if int(x2.shape[1]) != in_dim:
                    if int(x2.shape[1]) > in_dim:
                        x2 = x2[:, :in_dim]
                    else:
                        pad = np.zeros((x2.shape[0], in_dim - int(x2.shape[1])), dtype=x2.dtype)
                        x2 = np.concatenate([x2, pad], axis=1)
                y = x2 @ w.T
                if b is not None:
                    y = y + b[None, :]
                return y.reshape(*shp[:-1], y.shape[-1]).astype(np.float32)

            def _sample_positions(sample: Dict[str, np.ndarray]) -> Optional[np.ndarray]:
                pos = sample.get("positions")
                if pos is not None:
                    return np.asarray(pos, dtype=np.float32)
                starts = sample.get("starts")
                ends = sample.get("ends")
                origins = sample.get("origins")
                directions = sample.get("directions")
                if starts is None or ends is None or origins is None or directions is None:
                    return None
                try:
                    st = np.asarray(starts, dtype=np.float32)
                    en = np.asarray(ends, dtype=np.float32)
                    org = np.asarray(origins, dtype=np.float32)
                    dr = np.asarray(directions, dtype=np.float32)
                    mids = 0.5 * (st + en)
                    # Broadcast to [..., num_samples, 3]
                    return (org + dr * mids).astype(np.float32)
                except Exception:
                    return None

            def _sample_deltas(sample: Dict[str, np.ndarray]) -> Optional[np.ndarray]:
                de = sample.get("deltas")
                if de is not None:
                    return np.asarray(de, dtype=np.float32)
                starts = sample.get("starts")
                ends = sample.get("ends")
                if starts is None or ends is None:
                    return None
                try:
                    st = np.asarray(starts, dtype=np.float32)
                    en = np.asarray(ends, dtype=np.float32)
                    return (en - st).astype(np.float32)
                except Exception:
                    return None

            def _infer_encoding_input(sample: Dict[str, np.ndarray], trace_name: str) -> Optional[np.ndarray]:
                in_shapes, _ = _parse_io_shapes(trace_name)
                in_shape = in_shapes[0] if in_shapes else None
                rank = len(in_shape) if in_shape else 0
                if rank >= 3:
                    pos = _sample_positions(sample)
                    if pos is not None:
                        return pos
                if rank == 2:
                    d = sample.get("directions")
                    if d is None:
                        return None
                    d = np.asarray(d, dtype=np.float32)
                    if d.ndim == 3 and d.shape[1] == 1:
                        d = d[:, 0, :]
                    return d
                pos = _sample_positions(sample)
                if pos is not None:
                    return pos
                return None

            def _nerf_encode(x: np.ndarray, out_dim: int) -> np.ndarray:
                x = np.asarray(x, dtype=np.float32)
                in_dim = int(x.shape[-1])
                # Prefer include_input=True when mathematically consistent.
                include_input = False
                num_freq = 0
                if out_dim >= in_dim and (out_dim - in_dim) % (2 * in_dim) == 0:
                    include_input = True
                    num_freq = (out_dim - in_dim) // (2 * in_dim)
                elif out_dim % (2 * in_dim) == 0:
                    include_input = False
                    num_freq = out_dim // (2 * in_dim)
                if num_freq <= 0:
                    # Fallback: best effort on known vanilla NeRF dims.
                    if out_dim == 63 and in_dim == 3:
                        include_input, num_freq = True, 10
                    elif out_dim == 27 and in_dim == 3:
                        include_input, num_freq = True, 4
                    else:
                        return x
                freqs = (2.0 ** np.arange(num_freq, dtype=np.float32)).reshape((1,) * (x.ndim - 1) + (1, num_freq))
                scaled = (2.0 * np.pi * x[..., None]) * freqs
                scaled = scaled.reshape(*x.shape[:-1], -1)
                enc = np.sin(np.concatenate([scaled, scaled + (np.pi / 2.0)], axis=-1))
                if include_input:
                    enc = np.concatenate([enc, x], axis=-1)
                return enc.astype(np.float32)

            def _concat_parent_features(vals: List[np.ndarray], expected_in_dim: Optional[int]) -> Optional[np.ndarray]:
                if not vals:
                    return None
                if len(vals) == 1:
                    return vals[0]
                base_shape = vals[0].shape[:-1]
                cand = [v for v in vals if v.shape[:-1] == base_shape]
                if not cand:
                    return vals[0]
                if expected_in_dim is not None:
                    ed = int(expected_in_dim)
                    # Prefer exact single parent first.
                    for v in cand:
                        if int(v.shape[-1]) == ed:
                            return v.astype(np.float32)
                    # Then try exact concatenation subset.
                    dims = [int(v.shape[-1]) for v in cand]
                    n = len(cand)
                    best = None
                    for mask in range(1, 1 << n):
                        total = 0
                        idxs = []
                        for i in range(n):
                            if (mask >> i) & 1:
                                total += dims[i]
                                idxs.append(i)
                        if total == ed:
                            if best is None or len(idxs) < len(best):
                                best = idxs
                    if best is not None:
                        sel = [cand[i] for i in best]
                        # Trunk (256) before dir-encoding (27) tends to match call semantics.
                        sel = sorted(sel, key=lambda a: int(a.shape[-1]), reverse=True)
                        return np.concatenate(sel, axis=-1).astype(np.float32)
                    # Fallback: closest single parent, then crop/pad.
                    closest = min(cand, key=lambda v: abs(int(v.shape[-1]) - ed))
                    x = closest.astype(np.float32)
                    if int(x.shape[-1]) > ed:
                        return x[..., :ed]
                    if int(x.shape[-1]) < ed:
                        pad_shape = list(x.shape)
                        pad_shape[-1] = ed - int(x.shape[-1])
                        pad = np.zeros(pad_shape, dtype=np.float32)
                        return np.concatenate([x, pad], axis=-1)
                    return x
                # Unknown expected dim: concatenate by descending channel width.
                cand_desc = sorted(cand, key=lambda a: int(a.shape[-1]), reverse=True)
                out = np.concatenate(cand_desc, axis=-1)
                return out.astype(np.float32)

            def _execute_mlp(x: np.ndarray, bundle_meta: Dict[str, Any], params: Dict[str, np.ndarray]) -> Optional[np.ndarray]:
                try:
                    x = np.asarray(x, dtype=np.float32)
                    if x.ndim < 2:
                        return None
                    skip = set(int(v) for v in (bundle_meta.get("skip_connections") or []))
                    act_name = str(bundle_meta.get("activation") or "")
                    out_act_name = str(bundle_meta.get("out_activation") or "")
                    # Collect layers.i.{weight,bias}
                    layer_ids = []
                    for k in params.keys():
                        m = re.fullmatch(r"layers\.(\d+)\.weight", str(k))
                        if m:
                            layer_ids.append(int(m.group(1)))
                    if not layer_ids:
                        return None
                    layer_ids = sorted(set(layer_ids))
                    x0 = x
                    y = x
                    for li in layer_ids:
                        w = params.get(f"layers.{li}.weight")
                        b = params.get(f"layers.{li}.bias")
                        if w is None:
                            return None
                        if li in skip:
                            y = np.concatenate([x0, y], axis=-1)
                        y = _linear(y, w, b)
                        if li < layer_ids[-1]:
                            y = _apply_activation(y, act_name)
                    y = _apply_activation(y, out_act_name)
                    return y.astype(np.float32)
                except Exception:
                    return None

            def _execute_field_head(x: np.ndarray, fn: str, params: Dict[str, np.ndarray]) -> Optional[np.ndarray]:
                try:
                    x = np.asarray(x, dtype=np.float32)
                    w = params.get("net.weight")
                    b = params.get("net.bias")
                    if w is None:
                        # fallback if exported as layers.0.*
                        w = params.get("layers.0.weight")
                        b = params.get("layers.0.bias")
                    if w is None:
                        return None
                    y = _linear(x, w, b)
                    if "DensityFieldHead.forward" in str(fn):
                        y = _softplus(y)
                    elif "RGBFieldHead.forward" in str(fn):
                        y = _sigmoid(y)
                    return y.astype(np.float32)
                except Exception:
                    return None

            def _compute_weights(density: np.ndarray, deltas: np.ndarray) -> Optional[np.ndarray]:
                try:
                    d = np.asarray(density, dtype=np.float32)
                    de = np.asarray(deltas, dtype=np.float32)
                    if d.ndim == 2:
                        d = d[..., None]
                    if de.ndim == 2:
                        de = de[..., None]
                    s = min(int(d.shape[1]), int(de.shape[1]))
                    if s <= 0:
                        return None
                    d = d[:, :s, :1]
                    de = de[:, :s, :1]
                    delta_density = de * d
                    alphas = 1.0 - np.exp(-delta_density)
                    trans = np.cumsum(delta_density[:, :-1, :], axis=1)
                    zeros = np.zeros((trans.shape[0], 1, 1), dtype=np.float32)
                    trans = np.concatenate([zeros, trans], axis=1)
                    trans = np.exp(-trans)
                    weights = alphas * trans
                    weights = np.nan_to_num(weights).astype(np.float32)
                    return weights
                except Exception:
                    return None

            def _to_rgb_samples(arr: np.ndarray) -> Optional[np.ndarray]:
                try:
                    a = np.asarray(arr, dtype=np.float32)
                    if a.ndim == 3:
                        if int(a.shape[-1]) >= 3:
                            return a[..., :3]
                        return None
                    if a.ndim == 2 and int(a.shape[1]) % 3 == 0 and int(a.shape[1]) >= 3:
                        s = int(a.shape[1]) // 3
                        return a.reshape(a.shape[0], s, 3)
                except Exception:
                    return None
                return None

            def _to_weights(arr: np.ndarray) -> Optional[np.ndarray]:
                try:
                    a = np.asarray(arr, dtype=np.float32)
                    if a.ndim == 3:
                        if int(a.shape[-1]) >= 1:
                            return a[..., :1]
                        return None
                    if a.ndim == 2:
                        return a[..., None]
                except Exception:
                    return None
                return None

            def _execute_rgb_renderer_cached(rgb_arr: np.ndarray, weights_arr: np.ndarray) -> Optional[np.ndarray]:
                rgb = _to_rgb_samples(rgb_arr)
                w = _to_weights(weights_arr)
                if rgb is None or w is None:
                    return None
                if int(rgb.shape[0]) != int(w.shape[0]):
                    return None
                s = min(int(rgb.shape[1]), int(w.shape[1]))
                if s <= 0:
                    return None
                return np.sum(rgb[:, :s, :3] * w[:, :s, :1], axis=1).astype(np.float32)

            def _find_parent_array_values(op_obj, node_values: Dict[str, Any]) -> List[np.ndarray]:
                vals: List[np.ndarray] = []
                for p in getattr(op_obj, "parents", []) or []:
                    ptid = str(getattr(p, "trace_node_id", "")).strip()
                    if not ptid:
                        continue
                    v = node_values.get(ptid)
                    if isinstance(v, np.ndarray):
                        vals.append(v)
                return vals

            rgb_renderer_ops = [op for op in topo_ops if str(getattr(op, "op_type", "")) == "RGBRenderer"]
            replay_info["rgb_renderer_nodes"] = len(rgb_renderer_ops)
            replay_arrays: List[Tuple[Any, Any]] = []
            node_values: Dict[str, Any] = {}
            nodes_with_sources = 0
            executed_from_checkpoints = 0
            executed_from_recompute = 0

            # Fast path: if every RGB renderer has checkpointed rgb+weights,
            # skip expensive full-graph recomputation and replay directly.
            checkpoint_only_fast_path = False
            if rgb_renderer_ops and rgb_field_by_node and weights_by_node:
                checkpoint_only_fast_path = True
                for op in rgb_renderer_ops:
                    trace_id = str(getattr(op, "trace_node_id", "")).strip()
                    if not trace_id:
                        checkpoint_only_fast_path = False
                        break
                    link = renderer_link_by_node.get(trace_id, {})
                    rgb_src = str(link.get("rgb_node_id", "")).strip()
                    w_src = str(link.get("weights_node_id", "")).strip()
                    rgb_ck = _load_checkpoint(rgb_field_by_node, rgb_src) if rgb_src else None
                    w_ck = _load_checkpoint(weights_by_node, w_src) if w_src else None
                    if rgb_ck is None or w_ck is None:
                        checkpoint_only_fast_path = False
                        break
                    out_ck = _execute_rgb_renderer_cached(rgb_ck, w_ck)
                    if out_ck is None:
                        checkpoint_only_fast_path = False
                        break
                    node_values[trace_id] = out_ck
                    nodes_with_sources += 1
                    executed_from_checkpoints += 1
                    replay_arrays.append((op, out_ck))

            replay_info["checkpoint_only_fast_path_used"] = bool(checkpoint_only_fast_path)

            if not checkpoint_only_fast_path:
                for op in topo_ops:
                    trace_id = str(getattr(op, "trace_node_id", "")).strip()
                    if not trace_id:
                        continue

                    op_type = str(getattr(op, "op_type", ""))
                    trace_fn = str(getattr(op, "trace_function_name", ""))

                    # Sampling operators are graph roots with captured sources.
                    if op_type in ("UniformSampler", "PDFSampler"):
                        sample = _load_sample(trace_id)
                        if sample is not None:
                            node_values[trace_id] = sample
                        continue

                    if op_type == "RFFEncoding":
                        parent_vals = [node_values.get(str(getattr(p, "trace_node_id", "")).strip()) for p in getattr(op, "parents", []) or []]
                        x = None
                        for pv in parent_vals:
                            if isinstance(pv, dict):
                                x = _infer_encoding_input(pv, trace_fn)
                                if x is not None:
                                    break
                            elif isinstance(pv, np.ndarray):
                                x = pv
                                break
                        if x is None:
                            # Fallback to nearest sample source by trace linkage.
                            sample = _load_sample(trace_id)
                            if sample is not None:
                                x = _infer_encoding_input(sample, trace_fn)
                        if x is None:
                            continue
                        _, out_shape = _parse_io_shapes(trace_fn)
                        out_dim = int(out_shape[-1]) if out_shape else int(x.shape[-1])
                        node_values[trace_id] = _nerf_encode(x, out_dim)
                        continue

                    if op_type == "MLP":
                        bundle_uid = str(node_module_uid.get(trace_id, "")).strip()
                        bundle_meta = module_bundles.get(bundle_uid)
                        params = _load_module(bundle_uid) if bundle_uid else None
                        if not bundle_meta or not params:
                            continue
                        parent_arrays = _find_parent_array_values(op, node_values)
                        expected_in = bundle_meta.get("in_dim")
                        x = _concat_parent_features(parent_arrays, expected_in_dim=expected_in)
                        if x is None:
                            continue
                        if "field_components.field_heads." in trace_fn:
                            y = _execute_field_head(x, trace_fn, params)
                        else:
                            y = _execute_mlp(x, bundle_meta, params)
                        if y is None:
                            continue
                        node_values[trace_id] = y
                        continue

                    if op_type == "RGBRenderer":
                        link = renderer_link_by_node.get(trace_id, {})
                        rgb_src = str(link.get("rgb_node_id", "")).strip()
                        den_src = str(link.get("density_node_id", "")).strip()
                        smp_src = str(link.get("sample_node_id", "")).strip()
                        w_src = str(link.get("weights_node_id", "")).strip()

                        # Fast path: execute renderer directly from checkpointed tensors.
                        rgb_ck = _load_checkpoint(rgb_field_by_node, rgb_src) if rgb_src else None
                        w_ck = _load_checkpoint(weights_by_node, w_src) if w_src else None
                        if rgb_ck is not None and w_ck is not None:
                            out_ck = _execute_rgb_renderer_cached(rgb_ck, w_ck)
                            if out_ck is not None:
                                node_values[trace_id] = out_ck
                                nodes_with_sources += 1
                                executed_from_checkpoints += 1
                                replay_arrays.append((op, out_ck))
                                continue

                        rgb = node_values.get(rgb_src)
                        density = node_values.get(den_src)
                        sample = node_values.get(smp_src) if smp_src else None
                        if sample is None and smp_src:
                            sample = _load_sample(smp_src)
                            if sample is not None:
                                node_values[smp_src] = sample

                        if rgb is None:
                            # fallback from parent tensors (largest channel is usually RGB)
                            parr = _find_parent_array_values(op, node_values)
                            if parr:
                                parr = sorted(parr, key=lambda a: int(a.shape[-1]), reverse=True)
                                rgb = parr[0]
                                if len(parr) > 1 and density is None:
                                    density = parr[-1]

                        deltas = None
                        if isinstance(sample, dict):
                            deltas = _sample_deltas(sample)
                        if rgb is None or density is None or deltas is None:
                            continue
                        weights = _compute_weights(density, deltas)
                        if weights is None:
                            continue
                        s = min(int(rgb.shape[1]), int(weights.shape[1])) if rgb.ndim >= 3 else 0
                        if s <= 0:
                            continue
                        rgb2 = np.asarray(rgb[:, :s, :3], dtype=np.float32)
                        w2 = np.asarray(weights[:, :s, :1], dtype=np.float32)
                        out = np.sum(rgb2 * w2, axis=1).astype(np.float32)
                        node_values[trace_id] = out
                        nodes_with_sources += 1
                        executed_from_recompute += 1
                        replay_arrays.append((op, out))
                        continue

                    if op_type == "DensityRenderer":
                        # Optional: execute accumulation renderer from density + sample deltas.
                        link = renderer_link_by_node.get(trace_id, {})
                        den_src = str(link.get("density_node_id", "")).strip()
                        smp_src = str(link.get("sample_node_id", "")).strip()
                        w_src = str(link.get("weights_node_id", "")).strip()

                        # Fast path from checkpointed weights.
                        w_ck = _load_checkpoint(weights_by_node, w_src) if w_src else None
                        if w_ck is not None:
                            ww = _to_weights(w_ck)
                            if ww is not None:
                                node_values[trace_id] = np.sum(ww, axis=1).astype(np.float32)
                                continue

                        density = node_values.get(den_src)
                        sample = node_values.get(smp_src) if smp_src else None
                        if sample is None and smp_src:
                            sample = _load_sample(smp_src)
                            if sample is not None:
                                node_values[smp_src] = sample
                        if density is None or not isinstance(sample, dict):
                            continue
                        deltas = _sample_deltas(sample)
                        if deltas is None:
                            continue
                        weights = _compute_weights(density, deltas)
                        if weights is None:
                            continue
                        node_values[trace_id] = np.sum(weights, axis=1).astype(np.float32)
                        continue

            replay_info["missing_chunk_files"] = missing_files
            replay_info["rgb_renderer_nodes_with_sources"] = nodes_with_sources
            replay_info["rgb_renderer_nodes_executed"] = len(replay_arrays)
            replay_info["rgb_renderer_nodes_executed_from_checkpoints"] = executed_from_checkpoints
            replay_info["rgb_renderer_nodes_executed_from_recompute"] = executed_from_recompute
            if executed_from_checkpoints > 0 and executed_from_recompute == 0:
                replay_info["mode"] = "graph_operator_replay_checkpoints"
            elif executed_from_checkpoints > 0 and executed_from_recompute > 0:
                replay_info["mode"] = "graph_operator_execute_hybrid"
            else:
                replay_info["mode"] = "graph_operator_execute_full"

            if replay_arrays:
                # Prefer sink RGB nodes (graph terminals) to represent final rendered outputs.
                sink_arrays = []
                for op, arr in replay_arrays:
                    if not (getattr(op, "children", []) or []):
                        sink_arrays.append(arr)
                use_arrays = sink_arrays if sink_arrays else [arr for _, arr in replay_arrays]
                replay_info["rgb_renderer_sink_nodes_used"] = len(sink_arrays)

                rays = np.concatenate(use_arrays, axis=0)
                replay_info["chunk_count_used"] = len(use_arrays)

                # Find a reference rendered image from ns-eval output folder.
                ref_candidates = sorted(dag_dir.glob("eval_img_*.png"))
                if not ref_candidates:
                    ref_candidates = sorted([p for p in image_files if p.name.lower().startswith("eval_img_")])
                if ref_candidates:
                    ref_path = ref_candidates[0]
                    replay_info["reference_image"] = str(ref_path.relative_to(dag_dir))
                    ref = np.array(Image.open(ref_path).convert("RGB"))
                    h, w = ref.shape[0], ref.shape[1]
                    # Many ns-eval outputs are panelized horizontally (e.g., 3 columns).
                    # Choose the widest reference view that can be covered by captured rays.
                    ray_count = int(rays.shape[0])
                    candidate_ws = [w]
                    for div in (2, 3, 4):
                        if w % div == 0:
                            candidate_ws.append(w // div)
                    if h > 0:
                        candidate_ws.append(max(1, min(w, ray_count // h)))
                    candidate_ws = sorted(set([cw for cw in candidate_ws if cw > 0]), reverse=True)

                    chosen_w = None
                    for cw in candidate_ws:
                        if h * cw <= ray_count:
                            chosen_w = cw
                            break

                    if chosen_w is not None:
                        need = h * chosen_w
                        pred = rays[:need, :].reshape(h, chosen_w, 3)
                        # If reference is wider than chosen width, compare against centered crop.
                        if chosen_w < w:
                            x0 = (w - chosen_w) // 2
                            ref_cmp = ref[:, x0:x0 + chosen_w, :]
                            replay_info["reference_image"] = f"{replay_info['reference_image']}#crop(x={x0},w={chosen_w})"
                        else:
                            ref_cmp = ref
                        # Normalize to image range.
                        if float(np.nanmax(pred)) <= 1.5:
                            pred_img = np.clip(pred * 255.0, 0, 255).astype(np.uint8)
                        else:
                            pred_img = np.clip(pred, 0, 255).astype(np.uint8)
                        Image.fromarray(pred_img).save(replay_img_path)
                        replay_info["replay_image"] = str(replay_img_path.relative_to(dag_dir))
                        # Basic similarity metrics.
                        diff = pred_img.astype(np.float32) - ref_cmp.astype(np.float32)
                        mse = float(np.mean(diff * diff))
                        mae = float(np.mean(np.abs(diff)))
                        psnr = 99.0 if mse <= 1e-12 else float(20.0 * np.log10(255.0) - 10.0 * np.log10(mse))
                        replay_info["metrics"] = {
                            "mae_0_255": mae,
                            "mse_0_255": mse,
                            "psnr_db": psnr,
                            "compare_width": int(chosen_w),
                            "compare_height": int(h),
                        }
    except Exception as e:
        replay_info["error"] = str(e)

    payload = {
        "dag_file": str(dag_file),
        "dag_dir": str(dag_dir),
        "render_image_count": len(image_files),
        "render_images": [str(p.relative_to(dag_dir)) for p in image_files],
        "metric_files_checked": [str(p.relative_to(dag_dir)) for p in metric_files[:8]],
        "metrics": metric_summary,
        "renderer_replay_validation": replay_info,
    }

    json_path = dag_dir / f"{output_prefix}.json"
    json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    md_lines = [
        "# Rendering Validation",
        "",
        f"- DAG: `{dag_file.name}`",
        f"- Directory: `{dag_dir}`",
        f"- Render images found: **{len(image_files)}**",
        "",
        "## Render Image Files",
    ]
    if image_files:
        md_lines.extend([f"- `{p.relative_to(dag_dir)}`" for p in image_files[:200]])
        if len(image_files) > 200:
            md_lines.append(f"- ... ({len(image_files)-200} more)")
    else:
        md_lines.append("- None found (check `ns-eval --render-output-path`).")

    md_lines += ["", "## Metrics (if available)"]
    if metric_summary:
        for fp, m in metric_summary.items():
            metric_str = ", ".join([f"{k}={v}" for k, v in m.items()])
            md_lines.append(f"- `{fp}`: {metric_str}")
    else:
        md_lines.append("- No PSNR/SSIM/LPIPS/FPS metrics discovered near DAG.")

    md_lines += ["", "## Graph-Order Rendering Validation"]
    if replay_info.get("available"):
        md_lines.append(f"- Mode: `{replay_info.get('mode')}`")
        md_lines.append(f"- Manifest: `{replay_info.get('manifest')}`")
        md_lines.append(f"- Graph nodes: {replay_info.get('operator_graph_nodes')} (topological order: {replay_info.get('operator_graph_topological_count')})")
        md_lines.append(f"- RGB renderer nodes: {replay_info.get('rgb_renderer_nodes')}")
        md_lines.append(f"- Sample-source chunks: {replay_info.get('sample_source_chunks_total')}")
        md_lines.append(f"- RGB-field checkpoint chunks: {replay_info.get('rgb_field_chunks_total')}")
        md_lines.append(f"- Weights checkpoint chunks: {replay_info.get('weights_chunks_total')}")
        md_lines.append(f"- Density-field checkpoint chunks: {replay_info.get('density_field_chunks_total')}")
        md_lines.append(f"- Renderer links: {replay_info.get('renderer_links_total')}")
        md_lines.append(f"- Module bundles: {replay_info.get('module_bundles_total')}")
        md_lines.append(f"- Node-module bindings: {replay_info.get('node_module_bindings_total')}")
        md_lines.append(f"- RGB renderer nodes with sources: {replay_info.get('rgb_renderer_nodes_with_sources')}")
        md_lines.append(f"- RGB renderer nodes executed: {replay_info.get('rgb_renderer_nodes_executed')}")
        md_lines.append(f"- RGB renderer nodes executed from checkpoints: {replay_info.get('rgb_renderer_nodes_executed_from_checkpoints')}")
        md_lines.append(f"- RGB renderer nodes executed from recompute: {replay_info.get('rgb_renderer_nodes_executed_from_recompute')}")
        md_lines.append(f"- Chunks used (graph execution): {replay_info.get('chunk_count_used')}")
        md_lines.append(f"- Missing source chunk files: {replay_info.get('missing_chunk_files')}")
        if replay_info.get("reference_image"):
            md_lines.append(f"- Reference image: `{replay_info.get('reference_image')}`")
        if replay_info.get("replay_image"):
            md_lines.append(f"- Replayed image: `{replay_info.get('replay_image')}`")
        rmet = replay_info.get("metrics", {}) or {}
        if rmet:
            md_lines.append(f"- MAE (0-255): {rmet.get('mae_0_255'):.4f}")
            md_lines.append(f"- MSE (0-255): {rmet.get('mse_0_255'):.4f}")
            md_lines.append(f"- PSNR (dB): {rmet.get('psnr_db'):.4f}")
        else:
            md_lines.append("- Replay image was not generated (missing source chunks, unresolved graph sources, or insufficient rays/reference).")
        if replay_info.get("error"):
            md_lines.append(f"- Error: `{replay_info.get('error')}`")
    else:
        md_lines.append("- Renderer capture manifest not found. Re-run ns-eval with tracer capture enabled.")

    md_path = dag_dir / f"{output_prefix}.md"
    md_path.write_text("\n".join(md_lines) + "\n", encoding="utf-8")

    print(f"🧪 Rendering validation JSON: {json_path}")
    print(f"🧪 Rendering validation MD: {md_path}")
    return {"render_validation_json": str(json_path), "render_validation_md": str(md_path)}

def create_neural_rendering_visualizations(dag_path: str = "execution_dag.pkl", validation_only: bool = False):
    """Main function to create all neural rendering operator visualizations."""
    print("🎨 Neural Rendering Operator Visualization System")
    print("=" * 55)
    
    if not Path(dag_path).exists():
        print(f"[ERROR] DAG file not found: {dag_path}")
        print(f"💡 Generate with: python nerfstudio/nerfstudio/scripts/eval.py ... --enable-trace")
        return None
    
    visualizer = TransformedOperatorVisualizer()
    
    results = {}
    if not validation_only:
        try:
            # Generate graph visualizations first.
            results.update(visualizer.compare_before_after_visualization(dag_path, "neural_rendering_operators"))
        except Exception as e:
            print(f"[WARN] Operator graph visualization failed: {e}")
            import traceback
            traceback.print_exc()
    else:
        print("⏭️  Skipping graph plotting (--validation-only)")

    try:
        # Always attempt rendering-output validation, even if plotting fails.
        results.update(generate_rendering_validation_artifacts(dag_path))
    except Exception as e:
        print(f"[ERROR] Rendering validation failed: {e}")
        import traceback
        traceback.print_exc()

    if not results:
        return None

    print(f"\n🎉 All visualizations created successfully!")
    print(f"📂 Files generated:")
    for key, file_path in results.items():
        if isinstance(file_path, str):
            print(f"   {key}: {file_path}")

    return results

if __name__ == "__main__":
    # Test the visualization system
    import sys
    parser = argparse.ArgumentParser()
    parser.add_argument("dag_path", nargs="?", default="execution_dag.pkl")
    parser.add_argument("--validation-only", dest="validation_only", action="store_true", help="Skip graph plotting and generate replay-validation artifacts only")
    parser.add_argument("--plot-dot-subgraph", dest="plot_dot_subgraph", action="store_true", help="Also extract and plot subgraph clusters from a DOT file if provided")
    parser.add_argument("--dot", dest="dot_path", default=None, help="Path to grouped DOT file for subgraph extraction")
    parser.add_argument("--cluster-index", dest="cluster_index", default=None, help="Cluster selector: index (0,1,2), name (coarse|fine|other), or name:index (e.g., coarse:0)")
    parser.add_argument("--out-prefix", dest="out_prefix", default="execution_dag_component", help="Output prefix for subgraph rendering")
    args = parser.parse_args()
    dag_path = args.dag_path
    create_neural_rendering_visualizations(dag_path, validation_only=args.validation_only)
    if args.plot_dot_subgraph and plot_dot_subgraph_main and args.dot_path and args.cluster_index:
        # Reuse the sub-command by simulating argv for its parser
        import sys as _sys
        old_argv = list(_sys.argv)
        _sys.argv = ["plot_dot_subgraph.py", "--dot", args.dot_path, "--cluster-index", args.cluster_index, "--out-prefix", args.out_prefix]
        try:
            plot_dot_subgraph_main()
        finally:
            _sys.argv = old_argv