extra_materials/graph_visualizer_blip.py

import torch
import torch as t
import html
import json
import math
import re
import io
import base64
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from typing import Any, Dict, List, Tuple, OrderedDict, Optional
from torch import Tensor
import torch.nn.functional as F
from torchvision.utils import make_grid
from torch.utils.data import Dataset

# ---------------------------------------------------------
# 1. HTML / Viz Helpers
# ---------------------------------------------------------

def _render_token_html(tokens: List[str], acts: List[float], target_idx: int, max_global: float) -> str:
    html_str = ""
    for i, (tok, act) in enumerate(zip(tokens, acts)):
        ratio = act / max_global if max_global > 0 else 0
        alpha = max(0.0, min(1.0, ratio))
        bg_color = f"rgba(0, 200, 83, {alpha})" if act > 0 else "transparent"
        # Handle special chars
        safe_tok = html.escape(tok).replace('\n', '↵')
        is_target = (i == target_idx)
        target_class = "target-token" if is_target else ""
        html_str += f'<span class="token {target_class}" style="background-color: {bg_color};">{safe_tok}<span class="tooltip">{act:.4f}</span></span>'
    return html_str

def generate_interactive_html_embeddable(data: Dict[str, Any]) -> str:
    css_content = """<style>
        .interval-block { border-bottom: 1px solid #eee; margin-bottom: 8px; padding-bottom: 3px; }
        .interval-header { background: #f1f3f5; padding: 2px 6px; font-size: 0.65rem; font-weight: bold; color: #555; display: flex; justify-content: space-between; border-radius: 2px; }
        .example-row { display: flex; align-items: flex-start; padding: 4px 0; border-bottom: 1px solid #f0f0f0; margin-bottom: 4px; background: #fff; }
        .meta-col { width: 45px; flex-shrink: 0; display: flex; flex-direction: column; align-items: center; margin-right: 6px; border-right: 1px solid #f0f0f0; }
        .act-val { color: #28a745; font-weight: bold; font-size: 0.7rem; }
        .seq-col { flex-grow: 1; font-family: 'Consolas', monospace; font-size: 10px; line-height: 1.1; white-space: pre-wrap; cursor: pointer; padding-left: 4px; }
        .token { display: inline; padding: 0 1px; border-radius: 2px; position: relative; }
        .token:hover { outline: 1px solid #444; z-index: 10; }
        .target-token { border-bottom: 2px solid #333; font-weight: bold; }
        .token .tooltip { visibility: hidden; background-color: #333; color: #fff; text-align: center; padding: 4px; border-radius: 4px; position: absolute; z-index: 100; bottom: 120%; left: 50%; transform: translateX(-50%); font-size: 10px; white-space: nowrap; pointer-events: none; opacity: 0; transition: opacity 0.2s; }
        .token:hover .tooltip { visibility: visible; opacity: 1; }
    </style>"""
    html_content = f"""<div class="activation-viz">{css_content}
        <div class="act-header" style="padding-bottom:2px; margin-bottom:4px; border-bottom:1px solid #eee;"><strong>Max Act:</strong> {data['max_act']:.4f}</div>"""
    max_global = data['max_act']
    for interval in data['intervals']:
        density_pct = interval['density'] * 100
        html_content += f"""<div class="interval-block"><div class="interval-header"><span>{interval['min']:.2f}-{interval['max']:.2f}</span><span>{density_pct:.5f}%</span></div>"""
        for ex in interval['examples']:
            short_html = _render_token_html(ex['short']['tokens'], ex['short']['activations'], ex['short']['target_idx'], max_global)
            full_html = _render_token_html(ex['full']['tokens'], ex['full']['activations'], ex['full']['target_idx'], max_global)
            html_content += f"""<div class="example-row"><div class="meta-col"><span class="act-val">{ex['target_act']:.2f}</span></div><div class="seq-col" onclick="this.querySelector('.short-view').style.display=this.querySelector('.short-view').style.display==='none'?'block':'none';this.querySelector('.full-view').style.display=this.querySelector('.full-view').style.display==='none'?'block':'none';"><div class="short-view">{short_html}</div><div class="full-view" style="display:none;">{full_html}</div></div></div>"""
        html_content += "</div>"
    html_content += "</div>"
    return html_content

# ---------------------------------------------------------
# 2. Vision Processing Helpers
# ---------------------------------------------------------

def check_vision_sae(hook_name: str) -> bool:
    """Detects if a hook belongs to the vision component of BLIP."""
    return "visual" in hook_name or "vision" in hook_name

def denorm(img_tensor):
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
    # Ensure inputs are on the same device
    mean = mean.to(img_tensor.device)
    std = std.to(img_tensor.device)
    return (img_tensor * std + mean).clamp(0, 1)

def generate_base64_image_html(feat_idx: int, values: Tensor, indices: Tensor, dataset: Dataset, top_k: int = 10) -> str:
    """
    Generates a Base64 string of the combined heatmap visualization for a vision feature.
    Returns an HTML <img> tag.
    """
    if dataset is None:
        return "<div style='color:red'>No dataset provided for vision viz.</div>"

    mask = (values > 1e-3) * (indices == feat_idx)
    if mask.sum() == 0:
        return "<div style='color:#777'>No significant vision activation found.</div>"
    
    filtered_values = values * mask.to(values.dtype) # (b, activated_dim)
    
    sumed_filtered_values = filtered_values.sum(dim=1) # (b)
    top_vals, top_indices = sumed_filtered_values.topk(k=top_k) # (topk)
    top_acts = filtered_values[top_indices, :] # (topk, activated_dim)

    crops = []
    try:
        for each in range(top_indices.shape[0]):
            item = dataset[top_indices[each].item()]
            # Handle different dataset return formats, assuming 'pixel_values' is a tensor
            if isinstance(item, dict) and 'pixel_values' in item:
                crops.append(item['pixel_values'])
            else:
                # Fallback if dataset returns tuple
                crops.append(item[0]) 
    except Exception as e:
        return f"<div>Error accessing dataset: {str(e)}</div>"

    if not crops:
        return "<div>No crops extracted</div>"

    crops_tensor = torch.stack(crops, dim=0) # [top_k, 3, 384, 384]
    
    # 1. Standard Grid
    grid = make_grid(denorm(crops_tensor), nrow=5, padding=12, pad_value=1.0)
    grid = (grid.clamp(0,1) * 255).byte().cpu().permute(1,2,0).numpy()
    grid_pil = Image.fromarray(grid)

    # # 2. Heatmap Overlay
    # overlaid_crops = []
    # for each in range(crops_tensor.shape[0]):
    #     img = denorm(crops_tensor[each]).permute(1,2,0).cpu().numpy() # [384,384,3]
    #     act_flat = top_acts[each] # [576]
        
    #     # Assume square feature map. sqrt(576) = 24
    #     side = int(math.sqrt(act_flat.shape[0]))
    #     act_map = act_flat.view(side, side)
        
    #     # Normalize per image
    #     act_map = (act_map - act_map.min()) / (act_map.max() - act_map.min() + 1e-8)
        
    #     # Upsample to image size
    #     act_up = F.interpolate(
    #         act_map.unsqueeze(0).unsqueeze(0), 
    #         size=(img.shape[0], img.shape[1]), 
    #         mode='bilinear', 
    #         align_corners=False
    #     )[0][0]
    #     act_np = act_up.numpy()
        
    #     # Apply colormap
    #     heat_rgba = plt.cm.jet(act_np)
    #     heat_rgb = heat_rgba[:,:,:3]
    #     alpha = 0.4
    #     overlaid = img * (1 - alpha) + heat_rgb * alpha
    #     overlaid = overlaid.clip(0, 1)
    #     overlaid_tensor = torch.from_numpy(overlaid).float().permute(2,0,1)
    #     overlaid_crops.append(overlaid_tensor)
        
    # overlaid_tensor_stack = torch.stack(overlaid_crops, dim=0)
    # overlay_grid = make_grid(overlaid_tensor_stack, nrow=5, padding=12, pad_value=1.0)
    # overlay_grid = (overlay_grid.clamp(0,1) * 255).byte().cpu().permute(1,2,0).numpy()
    # overlay_grid_pil = Image.fromarray(overlay_grid)

    # # 3. Combine
    # width = grid_pil.width + overlay_grid_pil.width
    # height = max(grid_pil.height, overlay_grid_pil.height)
    # combined = Image.new('RGB', (width, height), (255, 255, 255))
    # combined.paste(grid_pil, (0, 0))
    # combined.paste(overlay_grid_pil, (grid_pil.width, 0))
    
    # 4. To Base64
    buffered = io.BytesIO()
    grid_pil.save(buffered, format="JPEG", quality=85)
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    
    return f'<img src="data:image/jpeg;base64,{img_str}" style="width:100%; border:1px solid #ddd; border-radius:4px;">'

# ---------------------------------------------------------
# 3. Main Logic
# ---------------------------------------------------------

class SAEGraphVisualizer:
    def __init__(self, graph_obj, model, processor, sae_dict, tokens, full_data_dict=None, data_toks=None, dataset=None):
        """
        Args:
            graph_obj: The sparse feature graph.
            model: The Blip model object
            processor: BlipProcessor (or similar) containing .tokenizer.
            sae_dict: Dictionary containing SAE-related utilities.
            tokens: List of token IDs or Tensor of token IDs.
            full_data_dict: Dictionary for retrieving activation data (keys: hook_name).
            data_toks: Tokens corresponding to the full_data_dict dataset (for text intervals).
            dataset: The PyTorch dataset object used for retrieving images (required for vision nodes).
        """
        self.graph = graph_obj
        self.model = model
        self.processor = processor
        self.sae_dict = sae_dict
        self.full_data_dict = full_data_dict
        self.data_toks = data_toks
        self.dataset = dataset
        
        # Handle tensor or list input for tokens
        if hasattr(tokens, 'tolist'):
            token_ids = tokens.tolist()
        else:
            token_ids = tokens
            
        self.tokens = [processor.tokenizer.decode(t) for t in token_ids]
        self.prompt_str = "".join(self.tokens)

    def get_layer_fallback(self, hook_name: str) -> int:
        match = re.search(r'\.(\d+)\.', hook_name)
        return int(match.group(1)) if match else 0

    def get_component_order(self, hook_name: str, layer: int) -> int:
        # Standardize component positioning for visual layout
        if "resid" in hook_name: return 0
        if "attn" in hook_name: return 1
        if "mlp" in hook_name: return 2
        if "vision" in hook_name: return 3 # Vision components usually processed before text or parallel
        return 0
    
    def get_logits_html(self, hook_name, feat_idx):
        if self.model is None: return "Model not provided"
        try:
            sae = self.sae_dict.get(hook_name)
            if not sae: return "N/A"
            
            # Extract feature vector
            if hasattr(sae, 'W_dec'):
                vec = sae.W_dec[feat_idx]
            elif hasattr(sae, 'decoder'):
                vec = sae.decoder.weight.data.T[feat_idx]
            else:
                return "N/A"
            
            # Ensure vec is on model device
            vec = vec.to(self.model.device)
            
            # Projection (BLIP specific modification)
            # Logits = model.text_decoder.cls.forward(vec)
            logits = self.model.text_decoder.cls(vec)
            
            k = 5
            pos_vals, pos_ids = logits.topk(k)
            neg_vals, neg_ids = logits.topk(k, largest=False)
            
            def pill(t_str, v, top):
                c = "background:#d4edda;color:#155724" if top else "background:#f8d7da;color:#721c24"
                return f'<span style="display:inline-block;padding:2px 6px;margin:2px;border-radius:4px;font-size:10px;{c}" title="{v:.4f}">{html.escape(t_str)}</span>'
            
            h = "<div><div style='font-size:10px;color:#666'>Top</div>"
            for idx, v in zip(pos_ids, pos_vals):
                # Use processor tokenizer for decoding
                t_str = self.processor.tokenizer.decode([idx]).replace('\n', '↵')
                h += pill(t_str, v.item(), True)
            
            h += "<div style='font-size:10px;color:#666;margin-top:4px'>Bottom</div>"
            for idx, v in zip(neg_ids, neg_vals):
                t_str = self.processor.tokenizer.decode([idx]).replace('\n', '↵')
                h += pill(t_str, v.item(), False)
            h += "</div>"
            return h
        except Exception as e:
            return f"Error: {str(e)}"

    def get_activation_html(self, hook_name, feat_idx, seq_pos):
        # Only use this for TEXT nodes
        if check_vision_sae(hook_name):
            return "<div>Vision node activation</div>"
            
        if self.data_toks is not None and self.full_data_dict and hook_name in self.full_data_dict:
            values, indices = self.full_data_dict[hook_name]
            try:
                # Call the Blip-specific interval fetcher
                # Assuming 'fetch_feature_activation_intervals_blip' is imported or available in context
                # If not, we fall back to a generic one or assume it exists in `sae.SAE_Blip_Explaining_Utils`
                from sae.SAE_Blip_Explaining_Utils import fetch_feature_activation_intervals_blip
                data = fetch_feature_activation_intervals_blip(self.processor, feat_idx, values, indices, self.data_toks, 5, 5, 10)
                return generate_interactive_html_embeddable(data)
            except ImportError:
                return "<div style='color:red'>fetch_feature_activation_intervals_blip not found</div>"
            except Exception as e:
                return f"<div style='color:red'>Error: {str(e)}</div>"
        return "<div style='color:#999;font-style:italic'>No activation data.</div>"

    def get_vision_node_html(self, hook_name, feat_idx):
        """
        Generates the visual inspection HTML (crops + heatmaps) for a vision node.
        """
        if self.full_data_dict and hook_name in self.full_data_dict:
            values, indices = self.full_data_dict[hook_name]
            return generate_base64_image_html(feat_idx, values.float(), indices.float(), self.dataset, top_k=5) # type: ignore
        return "<div>No vision data available in full_data_dict</div>"

    def build_graph_json(self):
        nodes_list = []
        raw_links_list = []
        
        # OPTIMIZATION 1: Pre-compute layer mapping to avoid O(N*L) lookups
        hook_to_layer = {}
        for cur_layer, dict_connection in self.graph.connection.items():
            for (k, _) in dict_connection.keys():
                hook_to_layer[k.name] = cur_layer

        # 1. Collect valid nodes and Vision Requests first
        valid_text_ids = set()
        valid_vision_ids = set()
        vision_feat_set = set() # (hook, feat_idx) for HTML generation
        
        # OPTIMIZATION 2: Cache for text feature HTML to avoid redundant dataset fetches
        # Map: (hook_name, feat_idx) -> html_string
        text_act_html_cache = {}
        text_logits_cache = {}

        def get_id(hook, seq, feat, is_err): 
            prefix = "vis" if check_vision_sae(hook) else "txt"
            return f"{prefix}::{hook}::{seq}::{'err' if is_err else 'feat'}::{feat}"

        # Loop Nodes
        for (node_key, _), sparse_act in self.graph.nodes.items():
            hook = node_key.name
            is_vision = check_vision_sae(hook)
            
            # Layer determination (O(1) lookup)
            layer = hook_to_layer.get(hook, -1)
            if layer == -1: layer = self.get_layer_fallback(hook)
            
            comp = self.get_component_order(hook, layer)
            scores = self.graph.node_scores.get((node_key, _))

            def process_nodes(is_resc):
                tensor_data = sparse_act.resc if is_resc else sparse_act.act
                if tensor_data is None: return
                
                nz = tensor_data.nonzero(as_tuple=True)
                
                # UPDATED LOGIC: Handle 1D tensors (errors) vs 2D tensors (features)
                if len(nz) == 1: 
                    s_indices = nz[0]
                    f_indices = -t.ones_like(s_indices)
                else:
                    s_indices = nz[0]
                    f_indices = nz[1]

                for s, f in zip(s_indices, f_indices):
                    s_item, f_item = s.item(), f.item()
                    score_val = 0.0
                    if scores:
                        # Handle score indexing safely for 1D vs 2D
                        if len(scores.act.shape) == 1: 
                            # If scores are 1D, they usually align with the 1D tensor data (e.g. error seq)
                            idx = s_item if len(nz) == 1 else f_item
                            score_val = scores.resc[idx].item() if is_resc else scores.act[idx].item()
                        else: 
                            score_val = scores.resc[s_item].item() if is_resc else scores.act[s_item, f_item].item()

                    node_id = get_id(hook, s_item, f_item if not is_resc else 0, is_resc)
                    
                    if is_vision:
                        # For vision, we only track valid features
                        if not is_resc and f_item != -1: 
                            valid_vision_ids.add(node_id)
                            vision_feat_set.add((hook, f_item))
                    else:
                        # For text, add to graph nodes immediately
                        
                        # OPTIMIZATION 2 Implementation: Check cache
                        logits_html = ""
                        act_html = ""
                        if not is_resc:
                            cache_key = (hook, f_item)
                            if cache_key in text_act_html_cache:
                                act_html = text_act_html_cache[cache_key]
                            else:
                                act_html = self.get_activation_html(hook, f_item, s_item)
                                text_act_html_cache[cache_key] = act_html
                                
                            if cache_key in text_logits_cache:
                                logits_html = text_logits_cache[cache_key]
                            else:
                                logits_html = self.get_logits_html(hook, f_item)
                                text_logits_cache[cache_key] = logits_html
                        
                        valid_text_ids.add(node_id)
                        nodes_list.append({
                            "id": node_id, 
                            "label": "Err" if is_resc else f"{f_item}",
                            "type": "error" if is_resc else "feature", 
                            "layer": layer, 
                            "comp": comp, 
                            "seq": s_item,
                            "score": score_val,
                            "hook": hook, 
                            "feat_idx": -1 if is_resc else f_item,
                            "act_html": act_html,
                            "logits_html": logits_html,
                            "vision_links": []
                        })
                    
            process_nodes(False) # Features
            process_nodes(True)  # Errors

        # 2. Generate Vision HTML (Batch Process unique valid vision nodes)
        vision_html_map = {}
        for (v_hook, v_feat) in vision_feat_set:
            vision_html_map[(v_hook, v_feat)] = self.get_vision_node_html(v_hook, v_feat)

        # 3. Process Edges (Only considering VALID nodes)
        # Create map for fast access to text node objects to append vision links
        text_node_map = {n['id']: n for n in nodes_list}

        for (end_key, end_idx), start_dict in self.graph.edges.items():
            end_hook = end_key.name
            end_is_vision = check_vision_sae(end_hook)
            
            for (start_key, start_idx), edge_tensor in start_dict.items():
                start_hook = start_key.name
                start_is_vision = check_vision_sae(start_hook)
                
                indices = edge_tensor.coalesce().indices()
                values = self.graph.edge_scores[(end_key, end_idx)][(start_key, start_idx)].coalesce().values()
                d_start = values.shape[2] - 1
                
                for i in range(indices.shape[1]):
                    end_seq, end_feat = indices[0, i].item(), indices[1, i].item()
                    
                    src_nz = values[i].nonzero(as_tuple=True)
                    for src_seq, src_feat in zip(*src_nz):
                        src_seq, src_feat = src_seq.item(), src_feat.item()
                        weight = values[i][src_seq, src_feat].item()
                        
                        is_err_src = (src_feat == d_start)
                        src_real_feat = src_feat if not is_err_src else 0
                        
                        src_id = get_id(start_hook, src_seq, src_real_feat, is_err_src)
                        tgt_id = get_id(end_hook, end_seq, end_feat, False)
                        
                        # FILTER: Check if both nodes are valid
                        src_valid = (src_id in valid_vision_ids) if start_is_vision else (src_id in valid_text_ids)
                        tgt_valid = (tgt_id in valid_vision_ids) if end_is_vision else (tgt_id in valid_text_ids)
                        
                        if not src_valid or not tgt_valid:
                            continue

                        # Add Links
                        # Case A: Text -> Text
                        if not start_is_vision and not end_is_vision:
                            raw_links_list.append({"source": src_id, "target": tgt_id, "value": weight})
                        
                        # Case B: Vision -> Text (Input)
                        elif start_is_vision and not end_is_vision:
                            if tgt_id in text_node_map:
                                html_content = vision_html_map.get((start_hook, src_real_feat), "<div>Error</div>")
                                text_node_map[tgt_id]['vision_links'].append({
                                    "id": src_id,
                                    "hook": start_hook,
                                    "feat": src_real_feat,
                                    "score": weight,
                                    "type": "feature",
                                    "image_html": html_content,
                                    "direction": "input"
                                })

                        # Case C: Text -> Vision (Output)
                        elif not start_is_vision and end_is_vision:
                            if src_id in text_node_map:
                                html_content = vision_html_map.get((end_hook, end_feat), "<div>Error</div>")
                                text_node_map[src_id]['vision_links'].append({
                                    "id": tgt_id,
                                    "hook": end_hook,
                                    "feat": end_feat,
                                    "score": weight,
                                    "type": "feature",
                                    "image_html": html_content,
                                    "direction": "output"
                                })

        # 4. Clean up and Merge
        # Sort vision links by score for better UI presentation
        for node in nodes_list:
            if node['vision_links']:
                node['vision_links'].sort(key=lambda x: abs(x['score']), reverse=True)

        max_comps = max((n['comp'] for n in nodes_list), default=0) + 1
        max_node_score = max((abs(n['score']) for n in nodes_list), default=1.0)

        return {
            "nodes": nodes_list, 
            "links": raw_links_list, # raw_links_list only contains valid connections now
            "prompt": self.prompt_str,
            "layers": sorted(list(set(n["layer"] for n in nodes_list))),
            "tokens": self.tokens,
            "max_comps": max(3, max_comps),
            "max_node_score": max_node_score
        }

    def generate_html(self, output_path="sae_graph_blip.html"):
        data = self.build_graph_json()
        
        class NaNEncoder(json.JSONEncoder):
            def default(self, o): return 0 if isinstance(o, float) and math.isnan(o) else super().default(o)
        json_data = json.dumps(data, cls=NaNEncoder)
        
        html_template = f"""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8"><title>BLIP SAE Circuit</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/d3/7.8.5/d3.min.js"></script>
    <style>
        body {{ margin:0; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif; display: flex; height: 100vh; overflow: hidden; font-size: 13px; }}
        #sidebar {{ width: 550px; border-left: 1px solid #ddd; background: #fff; display: flex; flex-direction: column; z-index:10; box-shadow: -2px 0 5px rgba(0,0,0,0.05); }}
        
        /* Tabs */
        .tab-header {{ display: flex; border-bottom: 1px solid #ddd; background: #f8f9fa; }}
        .tab-btn {{ flex: 1; padding: 10px; border: none; background: none; cursor: pointer; font-weight: 600; color: #666; border-bottom: 3px solid transparent; }}
        .tab-btn:hover {{ background: #eee; }}
        .tab-btn.active {{ color: #28a745; border-bottom-color: #28a745; background: #fff; }}
        
        .tab-content {{ flex: 1; overflow-y: auto; display: none; padding: 0; }}
        .tab-content.active {{ display: block; }}

        .panel {{ padding: 15px; border-bottom: 1px solid #eee; }}
        .panel h3 {{ margin: 0 0 10px 0; font-size: 12px; text-transform: uppercase; color: #555; font-weight: bold; letter-spacing: 0.5px; }}
        
        .panel-row {{ display: flex; border-bottom: 1px solid #eee; }}
        .panel-col {{ flex: 1; padding: 10px; min-width: 0; }}
        .panel-col:first-child {{ border-right: 1px solid #eee; }}

        #viz {{ flex-grow: 1; position: relative; background: #fdfdfd; }}
        .prompt-box {{ position: absolute; top:10px; left:10px; background: rgba(255,255,255,0.95); padding:8px 12px; border:1px solid #ccc; border-radius: 4px; font-family: monospace; font-size: 11px; max-width: 600px; z-index: 5; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }}
        
        /* D3 Elements */
        .node circle {{ fill: #fff; stroke: #888; stroke-width: 1.5px; transition: all 0.2s; }}
        .node.feature circle {{ stroke: #28a745; fill: #eafbe7; }}
        .node.error rect {{ stroke: #fd7e14; fill: #fff3e0; }}
        .node:hover circle, .node:hover rect {{ stroke: #000; stroke-width: 2px; }}
        .node.selected circle, .node.selected rect {{ stroke: #d00; stroke-width: 3px; }}
        .link {{ fill: none; stroke: #999; stroke-opacity: 0.2; }}
        .link.highlight {{ stroke: #333; stroke-opacity: 0.8; stroke-width: 1.5px; }}
        
        /* Lists */
        .feat-item {{ display: flex; justify-content: space-between; align-items: center; padding: 6px 10px; border-bottom: 1px solid #f0f0f0; cursor: pointer; }}
        .feat-item:hover {{ background: #f8f9fa; }}
        .feat-item.vision {{ background: #e3f2fd; }}
        .feat-item.vision:hover {{ background: #bbdefb; }}
        
        .f-info {{ display: flex; flex-direction: column; overflow: hidden; flex: 1; margin-right: 5px; }}
        .f-main {{ font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; display: flex; align-items: center; }}
        .f-meta {{ font-size: 10px; color: #888; margin-top: 2px; }}
        .f-bar-con {{ display: flex; align-items: center; width: 60px; justify-content: flex-end; }}
        .f-bar-bg {{ width: 35px; height: 6px; background: #eee; margin-left: 5px; border-radius: 2px; overflow: hidden; }}
        .f-bar {{ height: 100%; }}
        
        /* Vision Preview */
        #vision-preview {{ margin-top: 10px; cursor: zoom-in; position: relative; }}
        #vision-preview img {{ display: block; width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px; }}
        #vision-preview:hover::after {{
            content: "Click to Expand";
            position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%);
            background: rgba(0,0,0,0.7); color: white; padding: 5px 10px; border-radius: 4px; font-size: 12px; pointer-events: none;
        }}

        /* Vision Index Tab */
        .vi-group {{ border-bottom: 1px solid #eee; }}
        .vi-header {{ padding: 10px; background: #f8f9fa; cursor: pointer; display: flex; justify-content: space-between; align-items: center; }}
        .vi-header:hover {{ background: #e9ecef; }}
        .vi-body {{ display: none; padding: 0; background: #fff; }}
        .vi-body.open {{ display: block; }}
        .vi-link {{ padding: 6px 10px 6px 25px; border-bottom: 1px solid #f8f8f8; font-size: 11px; color: #444; cursor: pointer; display: flex; justify-content: space-between; }}
        .vi-link:hover {{ background: #fafafa; color: #000; }}

        /* Modal */
        .modal {{ display: none; position: fixed; z-index: 2000; left: 0; top: 0; width: 100%; height: 100%; overflow: auto; background-color: rgba(0,0,0,0.85); backdrop-filter: blur(2px); }}
        .modal-content {{ margin: auto; display: block; max-width: 90%; max-height: 90vh; position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); border-radius: 4px; box-shadow: 0 0 20px rgba(0,0,0,0.5); }}
        .modal-close {{ position: absolute; top: 15px; right: 35px; color: #f1f1f1; font-size: 40px; font-weight: bold; cursor: pointer; transition: 0.3s; }}
        .modal-close:hover {{ color: #bbb; }}
        
        .axis text {{ font-family: monospace; font-size: 10px; fill: #666; }}
        .axis path, .axis line {{ stroke: #ddd; }}
    </style>
</head>
<body>

<!-- Image Modal -->
<div id="img-modal" class="modal" onclick="closeModal()">
  <span class="modal-close">&times;</span>
  <img class="modal-content" id="modal-img">
</div>

<div id="viz">
    <div class="prompt-box"><strong>Prompt:</strong> {self.prompt_str}</div>
</div>

<div id="sidebar">
    <div class="tab-header">
        <button class="tab-btn active" onclick="switchTab('inspector')">Node Inspector</button>
        <button class="tab-btn" onclick="switchTab('vision')">Vision Index</button>
    </div>

    <!-- TAB 1: INSPECTOR -->
    <div id="tab-inspector" class="tab-content active">
        <div class="panel">
            <h2 id="sb-title" style="margin:0; font-size:16px;">Select a Node</h2>
            <div id="sb-meta" style="font-size:12px; color:#666; margin-top:5px;"></div>
        </div>
        <div id="sb-content" style="display:none">
            <div class="panel-row">
                <div class="panel-col"><h3>Inputs</h3><div id="list-in"></div></div>
                <div class="panel-col"><h3>Outputs</h3><div id="list-out"></div></div>
            </div>
            
            <div class="panel">
                <h3>Vision Connections</h3>
                <div id="list-vision" style="max-height: 150px; overflow-y:auto; margin-bottom:10px;"></div>
                <div id="vision-preview" onclick="expandImage(this)">
                    <div style="text-align:center; padding:20px; color:#999; font-size:11px; background:#fafafa; border:1px dashed #ccc;">
                        Select a vision node above to see activations.<br>(Click image to expand)
                    </div>
                </div>
            </div>
            
            <div class="panel"><h3>Logits</h3><div id="logits"></div></div>

            <div class="panel"><h3>Activations (Text)</h3><div id="activations"></div></div>
        </div>
    </div>

    <!-- TAB 2: VISION INDEX -->
    <div id="tab-vision" class="tab-content">
        <div id="vision-index-list"></div>
    </div>
</div>

<script>
    const data = {json_data};
    const width = document.getElementById('viz').clientWidth;
    const height = document.getElementById('viz').clientHeight;
    let selectedNodeId = null;

    // --- TABS & MODAL ---
    function switchTab(tabId) {{
        document.querySelectorAll('.tab-content').forEach(el => el.classList.remove('active'));
        document.querySelectorAll('.tab-btn').forEach(el => el.classList.remove('active'));
        document.getElementById('tab-' + tabId).classList.add('active');
        document.querySelector(`.tab-btn[onclick="switchTab('${{tabId}}')"]`).classList.add('active');
    }}

    function expandImage(el) {{
        const img = el.querySelector('img');
        if(img) {{
            document.getElementById('modal-img').src = img.src;
            document.getElementById('img-modal').style.display = "block";
        }}
    }}
    function closeModal() {{ document.getElementById('img-modal').style.display = "none"; }}

    // --- BUILD VISION INDEX ---
    // Aggregates all vision nodes and their connections
    const visionIndex = {{}};
    data.nodes.forEach(node => {{
        if(node.vision_links) {{
            node.vision_links.forEach(vl => {{
                if(!visionIndex[vl.id]) {{
                    visionIndex[vl.id] = {{
                        id: vl.id, hook: vl.hook, feat: vl.feat, img: vl.image_html, 
                        conns: []
                    }};
                }}
                visionIndex[vl.id].conns.push({{
                    textNode: node,
                    score: vl.score,
                    direction: vl.direction
                }});
            }});
        }}
    }});
    
    const vList = document.getElementById('vision-index-list');
    const sortedVisKeys = Object.keys(visionIndex).sort();
    if(sortedVisKeys.length === 0) {{ vList.innerHTML = "<div style='padding:15px; color:#999'>No vision nodes found.</div>"; }}
    
    sortedVisKeys.forEach(key => {{
        const v = visionIndex[key];
        const group = document.createElement('div');
        group.className = "vi-group";
        
        // Header
        const header = document.createElement('div');
        header.className = "vi-header";
        header.innerHTML = `<div><strong>${{v.feat}}</strong> <span style="font-size:10px; color:#888">${{v.hook.split('.').slice(-2).join('.')}}</span></div> <div style="font-size:10px; background:#e3f2fd; color:#0d47a1; padding:2px 6px; border-radius:10px;">${{v.conns.length}} Links</div>`;
        
        // Body (Connections)
        const body = document.createElement('div');
        body.className = "vi-body";
        
        // Preview Image in Header (Toggle)
        const imgDiv = document.createElement('div');
        imgDiv.style.padding = "10px";
        imgDiv.style.textAlign = "center";
        imgDiv.style.borderBottom = "1px solid #f0f0f0";
        imgDiv.innerHTML = v.img;
        // Make image in list clickable too
        imgDiv.onclick = (e) => {{ e.stopPropagation(); expandImage(imgDiv); }}; 
        body.appendChild(imgDiv);

        v.conns.sort((a,b) => Math.abs(b.score) - Math.abs(a.score));
        v.conns.forEach(c => {{
            const link = document.createElement('div');
            link.className = "vi-link";
            const dirArrow = c.direction === 'input' ? '→' : '←';
            const col = c.score > 0 ? "#28a745" : "#dc3545";
            link.innerHTML = `<span>${{dirArrow}} Text Node <strong>${{c.textNode.feat_idx}}</strong> (Seq ${{c.textNode.seq}})</span> <span style="color:${{col}}">${{c.score.toFixed(3)}}</span>`;
            link.onclick = () => {{
                switchTab('inspector');
                // Trigger D3 click
                const n = d3.selectAll(".node").filter(dn => dn.id === c.textNode.id);
                if(!n.empty()) n.dispatch("click");
            }};
            body.appendChild(link);
        }});

        header.onclick = () => {{
            body.classList.toggle('open');
        }};
        
        group.appendChild(header);
        group.appendChild(body);
        vList.appendChild(group);
    }});


    // --- D3 VISUALIZATION ---
    if(!data.nodes.length) {{ document.getElementById('viz').innerHTML = "No nodes found."; }}
    else {{
        const svg = d3.select("#viz").append("svg").attr("width","100%").attr("height","100%")
            .call(d3.zoom().on("zoom", e => g.attr("transform", e.transform))).append("g");
        const g = svg;
        
        const margin = {{top: 50, right: 100, bottom: 80, left: 60}};
        const layers = data.layers.length ? data.layers : [0];
        const maxNodeScore = data.max_node_score || 1.0;
        
        const xScale = d3.scaleLinear()
            .domain([0, data.tokens.length - 1]) 
            .range([margin.left, width - margin.right]);
            
        const minL = Math.min(...layers), maxL = Math.max(...layers);
        const yScale = d3.scaleLinear()
            .domain([minL - 0.5, maxL + 1.5]) 
            .range([height - margin.bottom, margin.top]);

        const xAxis = d3.axisBottom(xScale)
            .tickValues(d3.range(0, data.tokens.length)) 
            .tickFormat(i => data.tokens[i] || i);
            
        g.append("g").attr("class", "axis")
            .attr("transform", `translate(0, ${{height - margin.bottom}})`)
            .call(xAxis)
            .selectAll("text")
            .style("text-anchor", "start") 
            .attr("dx", "10px") 
            .attr("dy", "-2px")
            .attr("transform", "rotate(45)");

        g.append("text").attr("x", width/2).attr("y", height-10).attr("text-anchor", "middle")
         .attr("fill", "#999").attr("font-size","10px").text("Sequence Position");

        const simulation = d3.forceSimulation(data.nodes)
            .force("link", d3.forceLink(data.links).id(d => d.id).strength(0.01)) 
            .force("charge", d3.forceManyBody().strength(-100))
            .force("collide", d3.forceCollide(8).strength(1))
            .force("x", d3.forceX(d => xScale(d.seq)).strength(2.0))
            .force("y", d3.forceY(d => {{
                const compOffset = (d.comp || 0) / (data.max_comps || 3) * 0.8;
                return yScale(d.layer + compOffset);
            }}).strength(2.0));

        const link = g.append("g").selectAll("path").data(data.links).join("path")
            .attr("class", "link")
            .attr("d", d => {{
                const src = data.nodes.find(n=>n.id===d.source);
                const tgt = data.nodes.find(n=>n.id===d.target);
                if(!src || !tgt) return null;
                const dx = tgt.x - src.x;
                return `M${{src.x}},${{src.y}} C${{src.x + dx/2}},${{src.y}} ${{tgt.x - dx/2}},${{tgt.y}} ${{tgt.x}},${{tgt.y}}`;
            }});

        const node = g.append("g").selectAll("g").data(data.nodes).join("g")
            .attr("class", d => "node " + d.type)
            .call(d3.drag()
                .on("start", (e,d) => {{ if(!e.active) simulation.alphaTarget(0.3).restart(); d.fx=d.x; d.fy=d.y; }})
                .on("drag", (e,d) => {{ d.fx=e.x; d.fy=e.y; }})
                .on("end", (e,d) => {{ if(!e.active) simulation.alphaTarget(0); d.fx=null; d.fy=null; }})
            )
            .on("click", (e,d) => selectNode(d));

        node.append("title")
            .text(d => `Feature: ${{d.type==='error'?'Error':d.feat_idx}}\\nHook: ${{d.hook}}\\nSeq: ${{d.seq}}\\nScore: ${{d.score.toFixed(5)}}`);

        node.filter(d=>d.type==='feature').append("circle").attr("r", 5);
        node.filter(d=>d.type==='error').append("rect").attr("x",-4).attr("y",-4).attr("width",8).attr("height",8).attr("transform","rotate(45)");

        simulation.on("tick", () => {{
            link.attr("d", d => {{
                if (isNaN(d.source.x) || isNaN(d.target.x)) return "";
                const dx = d.target.x - d.source.x, dy = d.target.y - d.source.y;
                return `M${{d.source.x}},${{d.source.y}} C${{d.source.x + dx/2}},${{d.source.y}} ${{d.target.x - dx*0.5}},${{d.target.y}} ${{d.target.x}},${{d.target.y}}`;
            }});
            node.attr("transform", d => `translate(${{d.x}},${{d.y}})`);
        }});

        function selectNode(d) {{
            selectedNodeId = d.id;
            d3.selectAll(".node").classed("selected", false);
            d3.select(event.currentTarget).classed("selected", true);
            link.classed("highlight", l => l.source.id === d.id || l.target.id === d.id);
            
            const scW = Math.min(Math.abs(d.score) / maxNodeScore * 50, 50);
            const scoreBarHtml = `<div style="display:inline-block; width:50px; height:6px; background:#eee; margin-left:5px; vertical-align:middle; border-radius:1px;"><div style="width:${{scW}}px; height:100%; background:#555"></div></div>`;
            
            document.getElementById("sb-title").innerHTML = d.type==='error' ? "Error Node" : `Feature ${{d.feat_idx}}`;
            document.getElementById("sb-meta").innerHTML = `Hook: ${{d.hook}}<br>Seq: ${{d.seq}}<br>Score: ${{d.score.toFixed(5)}} ${{scoreBarHtml}}`;
            document.getElementById("sb-content").style.display = "block";
            document.getElementById("activations").innerHTML = d.act_html;
            document.getElementById("logits").innerHTML = d.logits_html || "N/A";
            
            updateList("list-in", l => l.target.id === d.id, l => l.source);
            updateList("list-out", l => l.source.id === d.id, l => l.target);
            updateVisionList(d);
            
            // Switch to inspector tab if not active
            switchTab('inspector');
        }}

        function updateVisionList(node) {{
            const div = document.getElementById("list-vision");
            const prev = document.getElementById("vision-preview");
            div.innerHTML = "";
            prev.innerHTML = `<div style="text-align:center; padding:20px; color:#999; font-size:11px; background:#fafafa; border:1px dashed #ccc;">
                        Select a vision node above to see activations.<br>(Click image to expand)
                    </div>`;
            
            if (!node.vision_links || node.vision_links.length === 0) {{
                div.innerHTML = "<div style='color:#999;font-size:11px; padding:10px;'>None</div>";
                return;
            }}
            
            node.vision_links.forEach(v => {{
                const row = document.createElement("div");
                row.className = "feat-item vision";
                
                const col = v.score > 0 ? "#28a745" : "#dc3545";
                const dirStr = v.direction === 'input' ? '(In)' : '(Out)';
                
                row.innerHTML = `
                    <div class="f-info">
                        <span class="f-main">VIS:${{v.feat}} ${{dirStr}}</span>
                        <span class="f-meta">${{v.hook.split('.').slice(-2).join('.')}}</span>
                    </div>
                    <div class="f-bar-con">
                        <span style="font-family:monospace;color:${{col}}">${{v.score.toFixed(3)}}</span>
                    </div>
                `;
                
                row.onclick = () => {{
                     prev.innerHTML = v.image_html;
                }};
                div.appendChild(row);
            }});
        }}

        function updateList(id, filter, getOther) {{
            const div = document.getElementById(id);
            div.innerHTML = "";
            const items = data.links.filter(filter);
            const sumEdgeVal = items.reduce((sum, item) => sum + Math.abs(item.value), 0) || 1e-9;
            items.sort((a,b) => Math.abs(b.value) - Math.abs(a.value));
            
            if(!items.length) div.innerHTML = "<div style='color:#999;font-size:11px; padding:10px;'>None</div>";
            
            items.forEach(l => {{
                const otherNode = getOther(l); 
                if(!otherNode) return;
                const row = document.createElement("div");
                row.className = "feat-item";
                
                const pct = (Math.abs(l.value) / sumEdgeVal) * 100;
                const w = Math.min(pct, 100); 
                const col = l.value > 0 ? "#28a745" : "#dc3545";
                const nScW = Math.min((Math.abs(otherNode.score) / maxNodeScore) * 20, 20);
                
                row.innerHTML = `
                    <div class="f-info">
                        <span class="f-main" title="${{otherNode.hook}}">
                            ${{otherNode.type==='error'?'Err':otherNode.feat_idx}} 
                            <div style="display:inline-block; width:20px; height:4px; background:#eee; vertical-align:middle; margin-left:4px;"><div style="width:${{nScW}}px; height:100%; background:#888;"></div></div>
                        </span>
                        <span class="f-meta">${{otherNode.hook.split('.').slice(-2).join('.')}} • Seq ${{otherNode.seq}}</span>
                    </div>
                    <div class="f-bar-con">
                        <span style="font-family:monospace;color:${{col}}">${{l.value.toFixed(3)}}</span>
                        <div class="f-bar-bg"><div class="f-bar" style="width:${{w}}%; background:${{col}}"></div></div>
                    </div>
                `;
                row.onclick = () => {{
                    const n = d3.selectAll(".node").filter(dn => dn.id === otherNode.id);
                    if(!n.empty()) {{ n.dispatch("click"); }}
                }};
                div.appendChild(row);
            }});
        }}
    }}
</script></body></html>
"""
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_template)
        print(f"Visualization saved to {output_path}")