import torch as t
from torch import Tensor
from .Graph_Template import Graph, GraphName, Node, Index
from typing import List, Tuple, Dict, Union, Callable, Any
from tqdm import tqdm
from transformer_lens.hook_points import HookPoint
from transformer_lens import (
    utils,
    HookedTransformer,
    ActivationCache,
)
from copy import deepcopy
from collections import OrderedDict, defaultdict
from warnings import warn
from contextlib import contextmanager
from sae_lens import (
    HookedSAETransformer,
)
from functools import partial
from itertools import product
from collections import OrderedDict
from .Graph_utils import nested_dict_to_string
from .Sparse_act import SparseAct
from torch import sparse_coo_tensor
import einops
import torch.autograd as atg

def rademacher_sample(size) -> Tensor:
    return t.randint(0, 2, size, dtype=t.float32) * 2 - 1

def diag_H(g: Tensor, x: Tensor, niter: int = 100) -> Tensor:
    diag_H = t.zeros_like(x)
    for _ in range(niter):
        v = rademacher_sample(x.shape).to(x.device)
        
        Hv = atg.grad(
            g, x, grad_outputs=v, retain_graph=True, allow_unused=True
        )[0]
        
        diag_H.add_(v * Hv)
        
    return diag_H

def diag_H_mask(g: Tensor, x: Tensor, mask: Tensor, niter: int = 100) -> Tensor:
    x_mask = x[mask]
    g_mask = g[mask]
    diag_H = t.zeros_like(x_mask)
    for _ in range(niter):
        v = rademacher_sample(x_mask.shape).to(x.device)
        
        Hv = atg.grad(
            g_mask, x, grad_outputs=v, retain_graph=True, allow_unused=True
        )[0][mask]
        
        diag_H.add_(v * Hv)
        
    return diag_H

def cache_to_sparseact(
    cache: Dict[str, Tensor] | ActivationCache, 
    act_name: str, 
    res_name: str | None = None, 
    resc_name: str | None = None,
) -> SparseAct:
    return SparseAct(
        act=cache[act_name],
        res=cache[res_name] if res_name is not None else None,
        resc=cache[resc_name] if resc_name is not None else None,
    )
    
def interpolate(start: Tensor, end: Tensor, frac: float,) -> Tensor:
    assert 0 <= frac <= 1, "frac must be in [0, 1]"
    return start * frac + end * (1 - frac)

def create_list(dict: Dict[str, Tensor], name: str) -> Dict[str, Tensor]:
    if name not in dict:
        dict[name] = 0 # type: ignore
    return dict
    
class ConnectionNode(Node):
    def __init__(self, name: str):
        self._name = name
        
    @property
    def name(self) -> str:
        return self._name
    
    def __eq__(self, other) -> bool:
        if isinstance(other, ConnectionNode):
            return self._name == other.name
        elif isinstance(other, str):
            return self._name == other
        else:
            raise NotImplementedError("other is not an instance of ConnectionNode or str")
    
    def __repr__(self) -> str:
        return self.name
    
    def __hash__(self) -> int:
        return hash(self.name)
    
class ConnectionIndex(Index):
    def __init__(
        self,
        list_index: tuple[int|None, ...] | None = None,
    ):
        if list_index is None:
            self.list_index = (None,)
        else:
            for index in list_index:
                assert type(index) == int or index == None, "index is not an instance of int or None"
            self.list_index = list_index
        
    @property
    def as_index(self) -> Tuple[int | slice, ...]:
        return tuple(slice(None) if x is None else x for x in self.list_index) # for indexing
    
    def __repr__(self) -> str:
        ret = "["
        for idx, x in enumerate(self.list_index):
            if idx > 0:
                ret += ", "
            if x is None:
                ret += ":"
            elif type(x) == int:
                ret += str(x)
            else:
                raise NotImplementedError(x)
        ret += "]"
        return ret
    
    def __eq__(self, other) -> bool:
        if isinstance(other, ConnectionIndex):
            return self.list_index == other.list_index
        elif isinstance(other, tuple):
            return self.list_index == other
        else:
            raise NotImplementedError("other is not an instance of ConnectionIndex or tuple")
    
    def __hash__(self) -> int:
        return hash(self.list_index)
    
class FeatureIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int], 
        length = 2,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))
        
class FeatureErrorIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int], 
        length = 2,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))
    
class ErrorIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int],
        length = 1,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))

def sae_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_acts_post' # e.g. ...hook_resid_pre.hook_sae_acts_post

def error_term_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_error'

def output_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_output'

def input_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_input'

def recons_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_recons'

def revert_hook_name(sae_name: str) -> str:
    return ".".join(sae_name.split(".")[:-1]) # e.g. ...hook_resid_pre.hook_sae_acts_post -> ...hook_resid_pre

def return_idx(idx: Tuple | List, length: int = 2) -> Tuple:
    assert len(idx) <= length, "The length of idx must be less than or equal to length"
    return tuple([None for _ in range(length - len(idx))] + list(idx))
    

class Feature_Graph(Graph):
    def __init__(
        self,
        model: HookedSAETransformer,
        saes: Dict[int, List[Tuple[str, Any]]], # {layer: list[{hook_position: HookedSAE}]}, can define granularity here
        use_error_term: bool = False,
    ):
        '''
        '''
        self.model = model
        self.cfg = model.cfg
        self.use_error_term = use_error_term
        self.device = self.cfg.device
        
        self.n_layers = self.cfg.n_layers
        # self.n_heads = self.cfg.n_heads
        
        assert len(saes) == self.n_layers, "please provide SAEs"
        self.saes = saes
        self.dict_saes: Dict[str, Any] = self.extract_saes(saes)
        
        self.reset_graph()
        
    def graph_type(self) -> str:
        return GraphName.feature_graph
        
    def extract_saes(self, saes: Dict[int, List[Tuple[str, Any]]]) -> Dict[str, Any]:
        dict_saes = {}
        for layer in range(self.n_layers):
            for hook_position, sae in saes[layer]:
                dict_saes[hook_position] = sae
        return dict_saes
        
    def build_default_connection(
        self,
        seq_length: int,
        token_wise: bool = False,
        inter: bool = True,
        intra: bool = False,
    ) -> Dict:
        '''
        Sample tokens to get the shape of the activations, necessary to build the graph of features
        If token_wise, the each node is a feature of a token, else, the graph is built feature-wise
        '''
        self.reset_graph()
        for layer in range(self.n_layers):
            # Setup
            self.connection[layer] = OrderedDict()
            for hook_position, _ in self.saes[layer]:
                self.connection[layer][(ConnectionNode(hook_position), ConnectionIndex())] = []
                
            # Inter-layer connections
            if inter:
                for i, (hook_position_end, _) in reversed(list(enumerate(self.saes[layer]))):
                    for j, (hook_position_start, _) in enumerate(self.saes[layer]):
                        if i > j:
                            self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
                            
            # Intra-layer connections
            if intra:
                for prev_layer in range(layer):
                    for hook_position_end, _ in self.saes[layer]:
                        for hook_position_start, _ in self.saes[prev_layer]:
                            self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
            # Sequential connection layer-1 --> layer
            else:
                if layer == 0:
                    continue
                for hook_position_start, _ in self.saes[layer-1]:
                    # connect to the first node of the next layer
                    hook_position_end = self.saes[layer][0][0] 
                    self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                        (ConnectionNode(hook_position_start), ConnectionIndex())
                    )
                
        self._build_nodes(seq_length, token_wise)
        self._build_graphs()
        return self.connection
    
    def _build_nodes(self, seq_length: int, token_wise: bool) -> None:
        self.seq_length = seq_length
        self.token_wise = token_wise
        
        for sae_name, sae in self.dict_saes.items():
            node_shape = (self.seq_length, sae.cfg.d_sae)
            self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.ones(node_shape, requires_grad=False).to(self.device),
                None,
                t.ones(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.zeros(node_shape, requires_grad=False).to(self.device),
                None,
                t.zeros(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            
    def _build_graphs(self) -> None:
        assert self.nodes, "Please build the nodes first"
        for _, connections in self.connection.items():
            for hook_position_end, list_hook_positions_start in connections.items():
                self.edges[hook_position_end] = OrderedDict()
                self.edge_scores[hook_position_end] = OrderedDict()
                
                for hook_position_start in list_hook_positions_start:
                    self.edges[hook_position_end][hook_position_start] = t.zeros([0]) # place holder
                    self.edge_scores[hook_position_end][hook_position_start] = t.zeros([0]) # place holder
        
    def reset_graph(self) -> None:
        self.connection: OrderedDict[int, OrderedDict[Tuple[Node, Index], List[Tuple[Node, Index]]]] = OrderedDict() # {layer: {hook_position_end: [hook_position_start]}}
        self.nodes: Dict[Tuple[Node, Index], SparseAct] = {} # {hook_position: SparseAct[act: Tensor['seq d_sae'], resc: Tensor['seq']]} 
        self.node_scores: Dict[Tuple[Node, Index], SparseAct] = {} 
        self.edges: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict() # {hook_position_end: {hook_position_start: Tensor}}
        self.edge_scores: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict()
        self.seq_length: int | None = None
        self.token_wise: bool | None = None
        
    def _check_graph(self) -> None:
        assert len(self.connection) > 0, "Graph is empty"
        assert len(self.nodes) > 0, "Nodes is empty"
        assert len(self.node_scores) > 0, "Node scores is empty"
        assert len(self.edges) > 0, "Edges is empty"
        assert len(self.edge_scores) > 0, "Edge scores is empty"
        assert self.seq_length is not None
        assert self.token_wise is not None
        
    def __repr__(self) -> str:
        self._check_graph()
        return nested_dict_to_string(self.connection, indent=4)
    
    def _check_feat_idx(self, feat_idx: Tuple | List) -> None:
        if self.token_wise:
            assert len(feat_idx) == 2, "Please provide the correct index"
        else:
            assert len(feat_idx) == 1, "Please provide the correct index"
            
    def _check_error_idx(self, error_idx: Tuple | List) -> None:
        if self.token_wise:
            assert len(error_idx) == 1, "Please provide the correct index"
        else:
            assert len(error_idx) == 0, "Please provide the correct index"
        
    def _active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> Tuple[List, ...]:
        self._check_graph()
        if self.token_wise:
            node = self.nodes[(sae_name, node_idx)] # act: (seq, d_sae), resc: (seq) 
        else:
            node = self.nodes[(sae_name, node_idx)].mean(dim=0) # act: (d_sae), resc: () - COLAPSED TENSOR
        
        if reverse:
            active_node = (node.act == 0).nonzero().tolist()
            active_error = (node.resc == 0).nonzero().tolist() if self.use_error_term else [] # type: ignore
        else:
            active_node = node.act.nonzero().tolist()
            active_error = node.resc.nonzero().tolist() if self.use_error_term else [] # type: ignore
            
        return active_node, active_error
            
    def active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> List[Tuple[Node, Index]]:
        active_node, active_error = self._active_nodes(sae_name, node_idx, reverse)
        
        feat_list = [
            (sae_name, FeatureIndex(idx)) 
            for idx in active_node
        ]
        error_list = [
            (sae_name, ErrorIndex(idx))
            for idx in active_error
        ]
        return feat_list + error_list # type: ignore
        
    def add_node(
        self, 
        sae_name: Node, 
        node_idx: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
        
    def add_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def delete_node(
        self, 
        sae_name: Node,
        node_idx: Index, 
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def delete_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def _check_node_shape(self, sae_name: Node, sparseact: SparseAct):
        try:
            if self.token_wise:
                assert sparseact.act.shape == t.Size([self.seq_length, self.dict_saes[sae_name.name].cfg.d_sae]) # type: ignore
            else:
                assert sparseact.act.shape == t.Size([self.dict_saes[sae_name.name].cfg.d_sae])
            if self.use_error_term:
                if self.token_wise:
                    assert sparseact.resc.shape == t.Size([self.seq_length]) # type: ignore
                else:
                    assert sparseact.resc.numel() == 1 # type: ignore
        except:
            raise ValueError("Wrong input shape for nodes.")     
    
    def update_node(
        self, 
        sae_name: Node,
        node_idx: Index,  
        value_and_mask: Tuple[SparseAct, SparseAct],
    ) -> None:
        self._check_graph()
        value, mask = value_and_mask
        assert isinstance(value, SparseAct), "value is not an instance of SparseAct"
        assert isinstance(mask, SparseAct), "mask is not an instance of SparseAct"
        self._check_node_shape(sae_name, value)
        
        if not self.token_wise:
            # convert to the right format of (seq, d_sae) and (seq) 
            value_act = einops.repeat(value.act, 'd_sae -> seq d_sae', seq=self.seq_length)
            mask_act = einops.repeat(mask.act, 'd_sae -> seq d_sae', seq=self.seq_length)
            value_resc = einops.repeat(value.resc, ' -> seq', seq=self.seq_length) if self.use_error_term else None
            mask_resc = einops.repeat(mask.resc, ' -> seq', seq=self.seq_length) if self.use_error_term else None
            
            value = SparseAct(value_act, None, value_resc)
            mask = SparseAct(mask_act, None, mask_resc)
            
        self.nodes[(sae_name, node_idx)] = mask.to(t.float32)
        self.node_scores[(sae_name, node_idx)] = value
    
    def update_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
        value_and_mask: Tuple[Tensor, Tensor],
    ) -> None:
        self._check_graph()
        value, mask = value_and_mask
        assert isinstance(value, Tensor), "value is not an instance of Tensor"
        assert isinstance(mask, Tensor), "mask is not an instance of Tensor"
        self.edges[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = mask
        self.edge_scores[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = value
    
    def find_deleted_nodes(self, reverse = False) -> List[Tuple[Node, Index]]:
        self._check_graph()
        
        all_list = []
        for sae_name, node_idx in self.nodes.keys():
            inactive_node, inactive_error = self._active_nodes(sae_name, node_idx, not reverse)
            feat_list = [
                (sae_name, FeatureIndex(idx))
                for idx in inactive_node
            ]
            error_list = [
                (sae_name, ErrorIndex(idx))
                for idx in inactive_error
            ]
            all_list += feat_list + error_list
        
        return all_list
    
    def find_deleted_edges(self, reverse = False) -> List[Tuple[Node, Index, Node, Index]]:
        self._check_graph()
        
        all_list = []
        for end_node, end_idx in self.edges.keys():
            for start_node, start_idx in self.edges[(end_node, end_idx)].keys():
                edge = self.edges[(end_node, end_idx)][(start_node, start_idx)]
                if not reverse:
                    inactive_node = (edge.values() == 0) # (num_active, seq, d_sae+1) or (num_active, d_sae+1)
                else:
                    inactive_node = edge.values()
                    
                d_sae_end = self.dict_saes[end_node.name].cfg.d_sae
                d_sae_start = self.dict_saes[start_node.name].cfg.d_sae
                    
                active_idx = [index for _, index in self.active_nodes(start_node, start_idx)]
                revised_active_nodes = [
                    index.idx if isinstance(index, FeatureIndex) else index.idx + (d_sae_start,) # type: ignore
                    for index in active_idx
                ]
                
                for num_active_idx in range(edge.indices().shape[1]):
                    end_node_idx = edge.indices()[:, num_active_idx].tolist()
                    # check if the end_node_idx is the error term
                    check_error = True if self.use_error_term and end_node_idx[-1] == d_sae_end else False
                    # create the end_index class to be used in the tuple
                    end_index = FeatureIndex(end_node_idx) if not check_error else ErrorIndex(end_node_idx[:-1])
                    
                    for i, idx in enumerate(revised_active_nodes):
                        if inactive_node[(num_active_idx,) + idx] > 0:
                            all_list.append(
                                (
                                    start_node, 
                                    active_idx[i], 
                                    end_node, 
                                    end_index,
                                )
                            )
        
        return all_list
    
    def iterate_nodes(self) -> List[Tuple[Node, Index]]:
        self._check_graph()
        return list(self.nodes.keys())
    
    def iterate_edges(self) -> List[Tuple[Node, Index, Node, Index]]:
        self._check_graph()
        all_edges = []
        for end, edges in self.edges.items():
            for start, edge in edges.items():
                all_edges.append(start + end)
        return all_edges
    
    def add_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
    ) -> None:
        self._check_graph()
        self._set_feat_value(sae_name, node_idx, feat_idx, 1)
        
    def delete_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
    ) -> None:
        self._check_graph()
        self._set_feat_value(sae_name, node_idx, feat_idx, 0)
        
    def update_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
        value: int | float | Tensor,
    ) -> None:
        self._check_graph()
        assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor."
        self._set_feat_value(sae_name, node_idx, feat_idx, value)
        
    def add_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_value(sae_name, node_idx, error_idx, 1)
        
    def delete_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_value(sae_name, node_idx, error_idx, 0)
        
    def update_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        self._check_graph()
        assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor."
        self._set_error_value(sae_name, node_idx, error_idx, value)
        
    def _set_feat_value(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
        value: int | float | Tensor,
    ) -> None:
        self._check_graph()
        
        if isinstance(feat_idx, (tuple, list)):
            self._check_feat_idx(feat_idx)
            self.nodes[(sae_name, node_idx)].act[return_idx(feat_idx)] = value
        elif isinstance(feat_idx, FeatureIndex):
            self.nodes[(sae_name, node_idx)].act[feat_idx.as_index] = value
        else:
            raise ValueError(f"feat_idx of type {type(feat_idx)} is not supported.")
        
    def _set_error_value(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        if not self.use_error_term:  # nothing to delete
            return
        
        if isinstance(error_idx, (tuple, list)):
            self._check_error_idx(error_idx)
            self.nodes[(sae_name, node_idx)].resc[return_idx(error_idx, 1)] = value # type: ignore
        elif isinstance(error_idx, ErrorIndex):
            self.nodes[(sae_name, node_idx)].resc[error_idx.as_index] = value # type: ignore
        else:
            raise ValueError(f"error_idx of type {type(error_idx)} is not supported.")
        
    def remove_error_term(self) -> None:
        if self.use_error_term:
            for sae_name, node_idx in self.nodes.keys():
                if self.nodes[(sae_name, node_idx)].resc is not None:
                    self.nodes[(sae_name, node_idx)].resc = t.zeros(self.seq_length, requires_grad=False).to(self.device) # type: ignore
    def forward(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor] | None,
        patch_deleted_comp: bool = False,
        **kwargs,
    ) -> Tuple[Tensor, Dict[str, SparseAct]]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        '''
        self._check_graph()
        self.model.reset_hooks()
        self.model_setup()
        
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                if patch_deleted_comp and corrupt_cache is not None:
                    act_mask = self.nodes[(sae_name, (None,))].act == 0 # type: ignore
                    act[:, act_mask] = corrupt_cache[sae_hook_name(sae_name)][:, act_mask]
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif hook.name == error_term_name(sae_name) and self.use_error_term:
                if patch_deleted_comp and corrupt_cache is not None:
                    resc_mask = self.nodes[(sae_name, (None,))].resc == 0 # type: ignore
                    act[:, resc_mask] = corrupt_cache[error_term_name(sae_name)][:, resc_mask]
                fwd_cache[error_term_name(sae_name)] = act.detach()
                
            # if "resid_post" in hook.name: # type: ignore
            #     print(act.max())
                
            return act
        
        with t.no_grad():
            with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                with self.model.hooks(
                    fwd_hooks=[
                        (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                        for sae_name in self.dict_saes.keys()
                    ],
                ):
                    logits = self.model(clean_token)
        
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = SparseAct(
                act=fwd_cache[sae_hook_name(sae_name)],
                res=fwd_cache[error_term_name(sae_name)] if self.use_error_term else None,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
    
        return logits, cache
    
    def __call__(self, *args, **kwargs) -> Tuple[Tensor, Dict[str, SparseAct]]:
        return self.forward(*args, **kwargs)
    
    def forward_backward_gradient(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        mode: str = 'node',
        gradient_mode: str = 'standard',
        pass_through_grad: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects
    ]:
        if mode == 'node':
            if verbose:
                print("Calculating node gradients...")
            if gradient_mode == "standard":
                node_grads, clean_cache = self._gradient_wrt_nodes(
                    clean_token, metric, retain_graph, pass_through_grad, **kwargs
                )
            elif gradient_mode == "ig":
                node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                    clean_token, corrupt_cache, metric, retain_graph, verbose, **kwargs
                )
            else:
                raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
            node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache)
                
            return node_effect, {}
        
        elif mode == 'edge':
            if kwargs.get('node_grads', None) is None or kwargs.get('node_effect', None) is None:
                # run the _gradient_wrt_nodes to get node_grads, and use node_effect to prune out unimportant nodes
                if verbose:
                    print("Calculating node gradients...")
                if gradient_mode == "standard":
                    node_grads, clean_cache = self._gradient_wrt_nodes(
                        clean_token, metric, retain_graph, pass_through_grad, **kwargs
                    )
                elif gradient_mode == "ig":
                    node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                        clean_token, corrupt_cache, metric, retain_graph, verbose, **kwargs
                    )
                else:
                    raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
                node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache)
            else:
                node_grads: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_grads') # type: ignore
                node_effect: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_effect') # type: ignore
                
            # Pruning
            if kwargs.get('prune', False):
                if verbose:
                    print("Pruning nodes...")
                self._prune_nodes(node_effect, verbose, **kwargs)
                
            if kwargs.get('gradient_only', False):
                if verbose:
                    print("Returning edge gradients only...")
                for name, sparse_act in node_grads.items():
                    node_grads[name] = sparse_act.to_sparse_like_self(t.ones_like(sparse_act.to_tensor()))
                    del sparse_act
                    
            edge_grads, _ = self._gradient_wrt_edges(
                clean_token, corrupt_cache, node_grads, verbose, **kwargs
            )
            return node_effect, edge_grads
        
        else:
            raise NotImplementedError(f"mode {mode} is not supported")
        
    def _attrib_effect_node(
        self,
        grads: Dict,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        clean_cache: Dict[str, SparseAct],
    ) -> Dict:
        attrib_effect = {}
        aggregate_dim = [0] if self.token_wise else [0, 1]
        
        for (node, index), grad in grads.items():
            corrupt_sparse_act = cache_to_sparseact(
                corrupt_cache, 
                sae_hook_name(node.name), 
                error_term_name(node.name) if self.use_error_term else None,
            )
            attrib_effect[(node, index)] = ( # act: (seq, d_sae), resc: (seq) || act: (d_sae), resc: ()
                grad @
                (corrupt_sparse_act - clean_cache[node.name])
            ).sum(aggregate_dim)
            
        return attrib_effect
    
    def _gradient_wrt_nodes(
        self,
        clean_token: Tensor,
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        pass_through_grad: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        Backward pass on the graph wrt the metric
        Return the gradients wrt nodes, edges, and activation cache
        '''
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache = {}
        pass_through_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == output_hook_name(sae_name):
                if self.use_error_term:
                    # we have: output = recon + stop_grad(error_term)
                    # so, the TRUE error_grad (if not stop_grad) is output_grad 
                    # due to: output = recon + error_term
                    bwd_cache[error_term_name(sae_name)] = grad.detach()
                if pass_through_grad:
                    pass_through_cache[output_hook_name(sae_name)] = grad.detach()
                    
            elif hook.name == input_hook_name(sae_name):
                if pass_through_grad:
                    # we have to modify inplace instead of grad = ... and then return grad
                    # because, returning a tensor in bwd pass hook is buggy somehow
                    grad.copy_(pass_through_cache[output_hook_name(sae_name)])
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif error_term_name(sae_name) == hook.name and self.use_error_term:
                fwd_cache[error_term_name(sae_name)] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    with self.model.hooks(
                        fwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                            for sae_name in self.dict_saes.keys()
                        ],
                        bwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_bwd, sae_name=sae_name)) 
                            for sae_name in self.dict_saes.keys()
                        ],
                    ):
                        metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                error_term_name(node.name) if self.use_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                error_term_name(sae_name) if self.use_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
    
    def _gradient_wrt_nodes_ig(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        steps = kwargs.get('steps', 10)
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache: Dict[str, Tensor] = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                create_list(bwd_cache, sae_hook_name(sae_name))
                bwd_cache[sae_hook_name(sae_name)] += grad.detach()
                
            elif hook.name == output_hook_name(sae_name):
                if self.use_error_term:
                    # we have: output = recon + stop_grad(error_term)
                    # so, the TRUE error_grad (if not stop_grad) is output_grad 
                    # due to: output = recon + error_term
                    create_list(bwd_cache, error_term_name(sae_name))
                    bwd_cache[error_term_name(sae_name)] += grad.detach()
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(sae_name)], 
                        act, 
                        frac,
                    )
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif error_term_name(sae_name) == hook.name and self.use_error_term:
                # interpolate for integrated gradients
                if hook.name == error_term_name(target_name):
                    act = interpolate(
                        corrupt_cache[error_term_name(sae_name)], 
                        act, 
                        frac,
                    )
                fwd_cache[error_term_name(sae_name)] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    for target_name in self.dict_saes.keys():
                        for step in range(steps):
                            frac = step / steps
                            with self.model.hooks(
                                fwd_hooks=[
                                    (
                                        lambda name, sae_name=sae_name: sae_name in name, 
                                        partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac)
                                    ) for sae_name in self.dict_saes.keys()
                                ],
                                bwd_hooks=[
                                    (
                                        lambda name, target_name=target_name: target_name in name, 
                                        partial(hook_sae_bwd, sae_name=target_name)
                                    ) 
                                ],
                            ):
                                metric(self.model(clean_token)).backward(retain_graph=retain_graph)
               
        # average the gradients                 
        for key in bwd_cache.keys():
            bwd_cache[key] /= steps
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                error_term_name(node.name) if self.use_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                error_term_name(sae_name) if self.use_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
        
    def _gradient_wrt_edges(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        node_grads: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects
        Dict[str, SparseAct]
    ]:
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        gradient_mode = kwargs.get('edge_gradient_mode', 'gradient')
            
        _, clean_cache = self.forward(clean_token, corrupt_cache=None)
                    
        edge_grads: Dict[Tuple[Node, Index, Node, Index], Tensor] = {}
        for layer, connection in tqdm(self.connection.items(), disable=not verbose):
            if verbose:
                print(f"Layer {layer}:")
            for hook_position_end, list_hook_positions_start in tqdm(connection.items(), disable=not verbose):
                assert hook_position_end in node_grads, f"Node gradient of {hook_position_end} is not provided."
                
                for hook_position_start in list_hook_positions_start:
                    corrupt_sparse_act = cache_to_sparseact(
                        corrupt_cache, 
                        sae_hook_name(hook_position_start[0].name), 
                        error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                    )
                    right_vec = corrupt_sparse_act - clean_cache[hook_position_start[0].name]
                    
                    if gradient_mode == "gradient":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution(
                            clean_token,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            layer,
                            **kwargs,
                        )
                    elif gradient_mode == "ig":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution_ig(
                            clean_token,
                            corrupt_cache,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            layer,
                            **kwargs,
                        )
                    else:
                        raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
                                  
        return edge_grads, clean_cache
    
    def _edge_attribution(
        self,
        clean_token: Tensor,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach()
                
            elif hook.name == output_hook_name(hook_position_start[0].name):
                # we have: output = recon + stop_grad(error_term)
                # so, the TRUE error_grad (if not stop_grad) is output_grad 
                # due to: output = recon + error_term
                if self.use_error_term:
                    # only store the gradient at the start position
                    bwd_cache[error_term_name(hook_position_start[0].name)] = grad.detach()

            # IMPORTANT NOTE for reproducibility:
            # we zero grad of intermediate components
            # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads
            # so, we instead zero grad of the input hook of SAE 
            elif hook.name == input_hook_name(sae_name):
                if hook.name != input_hook_name(hook_position_end[0].name):
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict) -> Tensor:
            if hook.name == sae_hook_name(hook_position_end[0].name):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act
                
            elif error_term_name(hook_position_end[0].name) == hook.name and self.use_error_term:
                # store activation for backward later
                to_bwd_cache[error_term_name(hook_position_end[0].name)] = act
                    
            return act
        
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        
        to_bwd_cache = {}
        bwd_cache = {}
        edge_effect = OrderedDict()
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    with self.model.hooks(
                        fwd_hooks=[
                            (lambda name: True, partial(
                                hook_sae_fwd, 
                                to_bwd_cache=to_bwd_cache, 
                            ))
                        ],
                        bwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(
                                hook_sae_bwd, 
                                sae_name=sae_name, 
                                bwd_cache=bwd_cache,
                            )) for sae_name in self.dict_saes.keys()
                        ],
                    ):
                        self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                        
                        aggregate_dim = [0] if self.token_wise else [0, 1]
                        to_bwd = ( # (seq, d_sae+1) || (d_sae+1)
                            cache_to_sparseact(
                                to_bwd_cache, 
                                sae_hook_name(hook_position_end[0].name), 
                                error_term_name(hook_position_end[0].name) if self.use_error_term else None,
                            ) @ leftvec.detach()
                        ).sum(aggregate_dim).to_tensor()
                        del to_bwd_cache
                        
                        for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)):
                            if isinstance(end_index, ErrorIndex):
                                # the last index is error: shape (d_sae+1) so last index is d_sae
                                to_bwd[end_index.idx + (d_sae_end,)].backward(retain_graph=True)
                                index = t.tensor(list(end_index.idx + (d_sae_end,)), device=self.device)
                            elif isinstance(end_index, FeatureIndex):
                                to_bwd[end_index.idx].backward(retain_graph=True)
                                index = t.tensor(list(end_index.idx), device=self.device)
                            else:
                                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
                            '''
                            edge_effect shape (seq, d_sae+1, seq, d_sae+1) or (d_sae+1, d_sae+1) in sparse_coo tensor
                            
                            the sparse_coo will have the shape:
                            --> indices of shape (2, num_active) or (1, num_active)
                            --> values of shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
                            '''
                            edge_effect[active_idx] = (
                                index,
                                ( # (seq, d_sae+1) || (d_sae+1)
                                    cache_to_sparseact(
                                        bwd_cache, 
                                        sae_hook_name(hook_position_start[0].name), 
                                        error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                                    ) @ rightvec
                                ).sum(aggregate_dim).to_tensor()
                            )
                        
                        del bwd_cache
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
        # if no active nodes, return empty tensor
        else:
            indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device)
            values = t.empty((0, seq, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device)
            
        if self.token_wise:
            return t.sparse_coo_tensor(indices, values, size=(seq, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
        
    def _edge_attribution_ig(
        self,
        clean_token: Tensor,
        corrupt_cache: Dict[str, Tensor] | ActivationCache,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        steps = kwargs.get('steps', 10)
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach()
                
            elif hook.name == output_hook_name(hook_position_start[0].name):
                # we have: output = recon + stop_grad(error_term)
                # so, the TRUE error_grad (if not stop_grad) is output_grad 
                # due to: output = recon + error_term
                if self.use_error_term:
                    # only store the gradient at the start position
                    bwd_cache[error_term_name(hook_position_start[0].name)] = grad.detach()

            # IMPORTANT NOTE for reproducibility:
            # we zero grad of intermediate components
            # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads
            # so, we instead zero grad of the input hook of SAE 
            elif hook.name == input_hook_name(sae_name):
                if hook.name != input_hook_name(hook_position_end[0].name):
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict, frac: float) -> Tensor:
            if hook.name == sae_hook_name(hook_position_end[0].name):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act
                
            elif error_term_name(hook_position_end[0].name) == hook.name and self.use_error_term:
                # store activation for backward later
                to_bwd_cache[error_term_name(hook_position_end[0].name)] = act
                
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # interpolate activation for intergrated gradients
                act = interpolate(
                    corrupt_cache[sae_hook_name(hook_position_start[0].name)], 
                    act, 
                    frac,
                )
                
            elif error_term_name(hook_position_start[0].name) == hook.name and self.use_error_term:
                # interpolate activation for intergrated gradients
                act = interpolate(
                    corrupt_cache[error_term_name(hook_position_start[0].name)], 
                    act, 
                    frac,
                )
                    
            return act
        
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        
        edge_effect = OrderedDict()
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    for step in range(steps):
                        frac = step / steps
                        
                        to_bwd_cache = {}
                        bwd_cache = {}
                        with self.model.hooks(
                            fwd_hooks=[
                                (lambda name: True, partial(
                                    hook_sae_fwd, 
                                    to_bwd_cache=to_bwd_cache, 
                                    frac=frac,
                                ))
                            ],
                            bwd_hooks=[
                                (lambda name, sae_name=sae_name: sae_name in name, partial(
                                    hook_sae_bwd, 
                                    sae_name=sae_name, 
                                    bwd_cache=bwd_cache,
                                )) for sae_name in self.dict_saes.keys()
                            ],
                        ):
                            self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                            
                            aggregate_dim = [0] if self.token_wise else [0, 1]
                            to_bwd = ( # (seq, d_sae+1) || (d_sae+1)
                                cache_to_sparseact(
                                    to_bwd_cache, 
                                    sae_hook_name(hook_position_end[0].name), 
                                    error_term_name(hook_position_end[0].name) if self.use_error_term else None,
                                ) @ leftvec.detach()
                            ).sum(aggregate_dim).to_tensor()
                            del to_bwd_cache
                            
                            for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)):
                                if isinstance(end_index, ErrorIndex):
                                    # the last index is error: shape (d_sae+1) so last index is d_sae
                                    to_bwd[end_index.idx + (d_sae_end,)].backward(retain_graph=True)
                                    index = t.tensor(list(end_index.idx + (d_sae_end,)), device=self.device)
                                elif isinstance(end_index, FeatureIndex):
                                    to_bwd[end_index.idx].backward(retain_graph=True)
                                    index = t.tensor(list(end_index.idx), device=self.device)
                                else:
                                    raise ValueError(f"end_index of type {type(end_index)} is not supported.")
                                '''
                                edge_effect shape (seq, d_sae+1, seq, d_sae+1) or (d_sae+1, d_sae+1) in sparse_coo tensor
                                
                                the sparse_coo will have the shape:
                                --> indices of shape (2, num_active) or (1, num_active)
                                --> values of shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
                                '''
                                if active_idx not in edge_effect:
                                    edge_effect[active_idx] = (
                                        index,
                                        ( # (seq, d_sae+1) || (d_sae+1)
                                            cache_to_sparseact(
                                                bwd_cache, 
                                                sae_hook_name(hook_position_start[0].name), 
                                                error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                                            ) @ rightvec
                                        ).sum(aggregate_dim).to_tensor()
                                    )
                                else:
                                    # accumulate the gradients
                                    prev_index, prev_value = edge_effect[active_idx]
                                    edge_effect[active_idx] = (
                                        prev_index,
                                        prev_value + ( # (seq, d_sae+1) || (d_sae+1)
                                            cache_to_sparseact(
                                                bwd_cache, 
                                                sae_hook_name(hook_position_start[0].name), 
                                                error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                                            ) @ rightvec
                                        ).sum(aggregate_dim).to_tensor()
                                    )
                            
                            del bwd_cache
        
        # average the gradients                 
        for active_idx, (prev_index, prev_value) in edge_effect.items():
            edge_effect[active_idx] = (
                prev_index,
                prev_value / steps
            )
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
        # if no active nodes, return empty tensor
        else:
            indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device)
            values = t.empty((0, seq, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device)
            
        if self.token_wise:
            return t.sparse_coo_tensor(indices, values, size=(seq, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()

    def _prune_nodes(
        self,
        node_effect: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> None:
        if kwargs.get("node_threshold", None) is None:
            raise ValueError("Please provide the node_threshold")
        else:
            node_threshold = kwargs.get("node_threshold")
            assert isinstance(node_threshold, (float, int, Tensor)), "node_threshold must be a int, float, or Tensor"
            
            for node, index in tqdm(node_effect.keys(), disable=not verbose):
                if kwargs.get("reverse_prune_node", False):
                    effect_feat_mask = node_effect[(node, index)].act.abs() < node_threshold
                else:
                    effect_feat_mask = node_effect[(node, index)].act.abs() > node_threshold
                    
                if verbose:
                    num_pruned = (~effect_feat_mask).sum().item()
                    print(f"{(node, index)}: pruned ({num_pruned / effect_feat_mask.numel() * 100:.5f}%) features")
                
                if self.use_error_term:
                    if kwargs.get("reverse_prune_node", False):
                        effect_error_mask = node_effect[(node, index)].resc.abs() < node_threshold # type: ignore
                    else:
                        effect_error_mask = node_effect[(node, index)].resc.abs() > node_threshold # type: ignore
                        
                    if verbose:
                        num_pruned = (~effect_error_mask).sum().item()
                        print(f"{(node, index)}: pruned ({num_pruned / effect_error_mask.numel() * 100:.5f}%) errors")
                        
                # replace the nodes with mask to prune
                mask = SparseAct(
                    act=effect_feat_mask,
                    resc=effect_error_mask if self.use_error_term else None, # type: ignore
                )
                self.nodes[(node, index)] = self.nodes[(node, index)] * mask
            
    def model_setup(self):
        pass
    
    def run_model(
        self, 
        toks: Tensor, 
        use_error_term: bool | None = None,
    ) -> Tuple[Tensor, ActivationCache]:
        
        return self.model.run_with_cache_with_saes( # type: ignore
            toks,
            saes=self._saes_to_list(),
            use_error_term=self.use_error_term if use_error_term is None else use_error_term,
            names_filter=lambda name: "sae" in name,
        )
    
    def _saes_to_list(self) -> List[Any]:
        return [sae for _, sae in self.dict_saes.items()]

    @contextmanager
    def _detach_error_term(self, detach: bool, sae_name: str | None = None):
        orig_detach_error_term = {}
        orig_disable_error_grad = {}
        try:
            for name, sae in self.dict_saes.items():
                if sae_name is None or sae_name == name:
                    orig_detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = detach
                else:
                    orig_detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = True # default detach error term
                
                orig_disable_error_grad[name] = sae.disable_error_grad
                sae.disable_error_grad = False # allow grad flows through error hook
            yield
        finally:
            for name, sae in self.dict_saes.items():
                sae.detach_error_term = orig_detach_error_term[name]
                sae.disable_error_grad = orig_disable_error_grad[name]
    
class ESAE_FG(Feature_Graph):
    def __init__(
        self,
        model: HookedSAETransformer,
        saes: Dict[int, List[Tuple[str, Any]]],
        esaes: Dict[int, List[Tuple[str, Any]]],
        use_esae_error_term: bool = False,
    ) -> None:
        super().__init__(model, saes, use_error_term=True)
        self.esaes = esaes
        self.dict_esaes = self.extract_saes(esaes)
        self.use_esae_error_term = use_esae_error_term
        
        self._check_valid_esaes()
            
    def _build_nodes(self, seq_length: int, token_wise: bool) -> None:
        self.seq_length = seq_length
        self.token_wise = token_wise
        
        for sae_name, sae in self.dict_saes.items():
            node_shape = (self.seq_length, sae.cfg.d_sae)
            
            esae = self.dict_esaes[sae_name]
            error_shape = (self.seq_length, esae.cfg.d_sae)
            
            self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.ones(node_shape, requires_grad=False).to(self.device),
                t.ones(error_shape, requires_grad=False).to(self.device) if self.use_error_term else None,
                t.ones(self.seq_length, requires_grad=False).to(self.device) if self.use_esae_error_term else None,
            )
            self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.zeros(node_shape, requires_grad=False).to(self.device),
                t.zeros(error_shape, requires_grad=False).to(self.device) if self.use_error_term else None,
                t.zeros(self.seq_length, requires_grad=False).to(self.device) if self.use_esae_error_term else None,
            )
            
    def _check_error_feature_idx(self, error_idx: Tuple | List) -> None:
        if self.token_wise:
            assert len(error_idx) == 2, "Please provide the correct index"
        else:
            assert len(error_idx) == 1, "Please provide the correct index"
      
    def _active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> Tuple[List, ...]:
        self._check_graph()
        if self.token_wise:
            node = self.nodes[(sae_name, node_idx)] # act: (seq, d_sae), resc: (seq, d_esae) 
        else:
            node = self.nodes[(sae_name, node_idx)].mean(dim=0) # act: (d_sae), resc: (d_esae)
        
        if reverse:
            active_node = (node.act == 0).nonzero().tolist()
            active_feature_error = (node.res == 0).nonzero().tolist() if self.use_error_term else [] # type: ignore
            active_error = (node.resc == 0).nonzero().tolist() if self.use_esae_error_term else [] # type: ignore
        else:
            active_node = node.act.nonzero().tolist()
            active_feature_error = node.res.nonzero().tolist() if self.use_error_term else [] # type: ignore
            active_error = node.resc.nonzero().tolist() if self.use_esae_error_term else [] # type: ignore
            
        return active_node, active_feature_error, active_error
    
    def active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> List[Tuple[Node, Index]]:
        active_node, active_feature_error, active_error = self._active_nodes(sae_name, node_idx, reverse)
        
        feat_list = [
            (sae_name, FeatureIndex(idx)) 
            for idx in active_node
        ]
        feat_error_list = [
            (sae_name, FeatureErrorIndex(idx))
            for idx in active_feature_error
        ]
        error_list = [
            (sae_name, ErrorIndex(idx))
            for idx in active_error
        ]
        return feat_list + feat_error_list + error_list # type: ignore
    
    def _check_node_shape(self, sae_name: Node, sparseact: SparseAct):
        try:
            if self.token_wise:
                assert sparseact.act.shape == t.Size([self.seq_length, self.dict_saes[sae_name.name].cfg.d_sae]) # type: ignore
            else:
                assert sparseact.act.shape == t.Size([self.dict_saes[sae_name.name].cfg.d_sae])
            if self.use_error_term:
                if self.token_wise:
                    assert sparseact.res.shape == t.Size([self.seq_length, self.dict_esaes[sae_name.name].cfg.d_sae]) # type: ignore
                else:
                    assert sparseact.res.shape == t.Size([self.dict_esaes[sae_name.name].cfg.d_sae]) # type: ignore
            if self.use_esae_error_term:
                if self.token_wise:
                    assert sparseact.resc.shape == t.Size([self.seq_length]) # type: ignore
                else:
                    assert sparseact.resc.numel() == 1 # type: ignore
        except:
            raise ValueError("Wrong input shape for nodes.") 
    
    def update_node(
        self, 
        sae_name: Node,
        node_idx: Index,  
        value_and_mask: Tuple[SparseAct, SparseAct],
    ) -> None:
        self._check_graph()
        value, mask = value_and_mask
        assert isinstance(value, SparseAct), "value is not an instance of SparseAct"
        assert isinstance(mask, SparseAct), "mask is not an instance of SparseAct"
        self._check_node_shape(sae_name, value)
        
        if not self.token_wise:
            # convert to the right format of (seq, d_sae) and (seq) 
            value_act = einops.repeat(value.act, 'd_sae -> seq d_sae', seq=self.seq_length)
            mask_act = einops.repeat(mask.act, 'd_sae -> seq d_sae', seq=self.seq_length)
            value_res = einops.repeat(value.res, 'd_esae -> seq d_esae', seq=self.seq_length) if self.use_error_term else None
            mask_res = einops.repeat(mask.res, 'd_esae -> seq d_esae', seq=self.seq_length) if self.use_error_term else None
            value_resc = einops.repeat(value.resc, ' -> seq', seq=self.seq_length) if self.use_esae_error_term else None
            mask_resc = einops.repeat(mask.resc, ' -> seq', seq=self.seq_length) if self.use_esae_error_term else None
            
            value = SparseAct(value_act, value_res, value_resc)
            mask = SparseAct(mask_act, mask_res, mask_resc)
            
        self.nodes[(sae_name, node_idx)] = mask.to(t.float32)
        self.node_scores[(sae_name, node_idx)] = value
    
    def find_deleted_nodes(self, reverse = False) -> List[Tuple[Node, Index]]:
        self._check_graph()
        
        all_list = []
        for sae_name, node_idx in self.nodes.keys():
            inactive_node, inactive_feature_error, inactive_error = self._active_nodes(sae_name, node_idx, not reverse)
            feat_list = [
                (sae_name, FeatureIndex(idx))
                for idx in inactive_node
            ]
            feat_error_list = [
                (sae_name, FeatureErrorIndex(idx))
                for idx in inactive_feature_error
            ]
            error_list = [
                (sae_name, ErrorIndex(idx))
                for idx in inactive_error
            ]
            all_list += feat_list + feat_error_list + error_list
        
        return all_list
    
    def find_deleted_edges(self, reverse = False) -> List[Tuple[Node, Index, Node, Index]]:
        self._check_graph()
        
        all_list = []
        for end_node, end_idx in self.edges.keys():
            for start_node, start_idx in self.edges[(end_node, end_idx)].keys():
                edge = self.edges[(end_node, end_idx)][(start_node, start_idx)]
                if not reverse:
                    inactive_node = (edge.values() == 0) # (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+1)
                else:
                    inactive_node = edge.values()
                    
                d_sae_end = self.dict_saes[end_node.name].cfg.d_sae
                d_esae_end = self.dict_esaes[end_node.name].cfg.d_sae
                d_sae_start = self.dict_saes[start_node.name].cfg.d_sae
                d_esae_start = self.dict_esaes[start_node.name].cfg.d_sae
                    
                active_idx = [index for _, index in self.active_nodes(start_node, start_idx)]
                revised_active_nodes = []
                for index in active_idx:
                    if isinstance(index, FeatureIndex):
                        revised_active_nodes.append(index.idx)
                    elif isinstance(index, FeatureErrorIndex):
                        revised_index = list(index.idx)
                        revised_index[-1] += d_sae_start
                        revised_active_nodes.append(tuple(revised_index))
                    elif isinstance(index, ErrorIndex):
                        revised_active_nodes.append(index.idx + (d_sae_start + d_esae_start,))
                
                for num_active_idx in range(edge.indices().shape[1]):
                    end_node_idx = edge.indices()[:, num_active_idx].tolist()
                    # check if the end_node_idx is the error term
                    if self.use_esae_error_term and end_node_idx[-1] == d_sae_end + d_esae_end:
                        end_index = ErrorIndex(end_node_idx[:-1])
                    elif self.use_error_term and end_node_idx[-1] == d_sae_end:
                        end_index = FeatureErrorIndex(end_node_idx)
                    else:
                        end_index = FeatureIndex(end_node_idx)
                    
                    for i, idx in enumerate(revised_active_nodes):
                        if inactive_node[(num_active_idx,) + idx] > 0:
                            all_list.append(
                                (
                                    start_node, 
                                    active_idx[i], 
                                    end_node, 
                                    end_index,
                                )
                            )
        
        return all_list
    
    def add_single_feature_error(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_error_idx: Tuple[int, ...] | List[int] | FeatureErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_feature_value(sae_name, node_idx, feat_error_idx, 1)
        
    def delete_single_feature_error(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_error_idx: Tuple[int, ...] | List[int] | FeatureErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_feature_value(sae_name, node_idx, feat_error_idx, 0)
        
    def update_single_feature_error(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_error_idx: Tuple[int, ...] | List[int] | FeatureErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        self._check_graph()
        assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor."
        self._set_error_feature_value(sae_name, node_idx, feat_error_idx, value)
    
    def _set_error_feature_value(
        self,
        sae_name: Node,
        node_idx: Index,
        error_feat_idx: Tuple[int, ...] | List[int] | FeatureErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        if not self.use_error_term:  # nothing to delete
            return
        
        if isinstance(error_feat_idx, (tuple, list)):
            self._check_error_feature_idx(error_feat_idx)
            self.nodes[(sae_name, node_idx)].res[return_idx(error_feat_idx)] = value # type: ignore
        elif isinstance(error_feat_idx, FeatureErrorIndex):
            self.nodes[(sae_name, node_idx)].res[error_feat_idx.as_index] = value # type: ignore
        else:
            raise ValueError(f"error_feat_idx of type {type(error_feat_idx)} is not supported.")
    
    def forward(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor] | None,
        patch_deleted_comp: bool = False,
        **kwargs,
    ) -> Tuple[Tensor, Dict[str, SparseAct]]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        '''
        self._check_graph()
        self.model.reset_hooks()
        self.model_setup()
        
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                if patch_deleted_comp and corrupt_cache is not None:
                    act_mask = self.nodes[(sae_name, (None,))].act == 0 # type: ignore
                    act[:, act_mask] = corrupt_cache[sae_hook_name(sae_name)][:, act_mask]
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif hook.name == sae_hook_name(error_term_name(sae_name)) and self.use_error_term:
                if patch_deleted_comp and corrupt_cache is not None:
                    res_mask = self.nodes[(sae_name, (None,))].res == 0 # type: ignore
                    act[:, res_mask] = corrupt_cache[sae_hook_name(error_term_name(sae_name))][:, res_mask]
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif error_term_name(error_term_name(sae_name)) == hook.name and self.use_esae_error_term:
                if patch_deleted_comp and corrupt_cache is not None:
                    resc_mask: Tensor = self.nodes[(sae_name, (None,))].resc == 0 # type: ignore
                    act[:, resc_mask] = corrupt_cache[error_term_name(error_term_name(sae_name))][:, resc_mask].detach()
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                
            return act
        
        with t.no_grad():
            with self._hook_esaes_to_saes():
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    with self.model.hooks(
                        fwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                            for sae_name in self.dict_saes.keys()
                        ]
                    ):
                        logits = self.model(clean_token)
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = SparseAct(
                act=fwd_cache[sae_hook_name(sae_name)],
                res=fwd_cache[sae_hook_name(error_term_name(sae_name))] if self.use_error_term else None,
                resc=fwd_cache[error_term_name(error_term_name(sae_name))] if self.use_esae_error_term else None,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
    
        return logits, cache
    
    def _attrib_effect_node(
        self,
        grads: Dict,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        clean_cache: Dict[str, SparseAct],
    ) -> Dict:
        attrib_effect = {}
        aggregate_dim = [0] if self.token_wise else [0, 1]
        
        for (node, index), grad in grads.items():
            corrupt_sparse_act = cache_to_sparseact(
                corrupt_cache, 
                sae_hook_name(node.name), 
                sae_hook_name(error_term_name(node.name)) if self.use_error_term else None,
                error_term_name(error_term_name(node.name)) if self.use_esae_error_term else None,
            )
            attrib_effect[(node, index)] = ( # act: (seq, d_sae), res: (seq, d_esae), resc: (seq, d_model) || act: (d_sae), res: (d_esae), resc: (d_model)
                grad *  # this is NOT MATMUL, elementwise because the sparseact.res are now features
                (corrupt_sparse_act - clean_cache[node.name])
            ).sum(aggregate_dim)
            
            if self.use_esae_error_term:
                attrib_effect[(node, index)].contract() # resc: (seq) || resc: ()
            
        return attrib_effect
    
    def _gradient_wrt_nodes(
        self,
        clean_token: Tensor,
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        pass_through_grad: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        Backward pass on the graph wrt the metric
        Return the gradients wrt nodes, edges, and activation cache
        '''
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache = {}
        pass_through_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == output_hook_name(sae_name):
                if pass_through_grad:
                    pass_through_cache[output_hook_name(sae_name)] = grad.detach()
                    
            elif sae_hook_name(error_term_name(sae_name)) == hook.name and self.use_error_term:
                bwd_cache[sae_hook_name(error_term_name(sae_name))] = grad.detach()
                
            elif output_hook_name(error_term_name(sae_name)) == hook.name and self.use_esae_error_term:
                # the error_grad is the output_grad, due to: output = recon + error_term
                bwd_cache[error_term_name(error_term_name(sae_name))] = grad.detach()
                    
            elif hook.name == input_hook_name(sae_name):
                if pass_through_grad:
                    # we have to modify inplace instead of grad = ... and then return grad
                    # because, returning a tensor in bwd pass hook is buggy somehow
                    grad.copy_(pass_through_cache[output_hook_name(sae_name)])
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:# cache the feature post activation
            if hook.name == sae_hook_name(sae_name):
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            # modify the error of the sae
            elif sae_hook_name(error_term_name(sae_name)) == hook.name and self.use_error_term:
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif error_term_name(error_term_name(sae_name)) == hook.name and self.use_esae_error_term:
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._hook_esaes_to_saes():
                    with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                        with self.model.hooks(
                            fwd_hooks=[
                                (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                                for sae_name in self.dict_saes.keys()
                            ],
                            bwd_hooks=[
                                (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_bwd, sae_name=sae_name)) 
                                for sae_name in self.dict_saes.keys()
                            ],
                        ):
                            metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                sae_hook_name(error_term_name(node.name)) if self.use_error_term else None,
                error_term_name(error_term_name(node.name)) if self.use_esae_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
    
    def _gradient_wrt_nodes_ig(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        steps = kwargs.get("steps", 10)
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                create_list(bwd_cache, sae_hook_name(sae_name))
                bwd_cache[sae_hook_name(sae_name)] += grad.detach()
                    
            elif sae_hook_name(error_term_name(sae_name)) == hook.name and self.use_error_term:
                create_list(bwd_cache, sae_hook_name(error_term_name(sae_name)))
                bwd_cache[sae_hook_name(error_term_name(sae_name))] += grad.detach()
                
            elif output_hook_name(error_term_name(sae_name)) == hook.name and self.use_esae_error_term:
                # the error_grad is the output_grad, due to: output = recon + error_term
                create_list(bwd_cache, error_term_name(error_term_name(sae_name)))
                bwd_cache[error_term_name(error_term_name(sae_name))] += grad.detach()
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            # cache the feature post activation
            if hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(sae_name)], 
                        act, 
                        frac,
                    )
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            # modify the error of the sae
            elif sae_hook_name(error_term_name(sae_name)) == hook.name and self.use_error_term:
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(error_term_name(target_name)):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(error_term_name(sae_name))], 
                        act, 
                        frac,
                    )
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif error_term_name(error_term_name(sae_name)) == hook.name and self.use_esae_error_term:
                # interpolate for integrated gradients
                if hook.name == error_term_name(error_term_name(target_name)):
                    act = interpolate(
                        corrupt_cache[error_term_name(error_term_name(sae_name))], 
                        act, 
                        frac,
                    )
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._hook_esaes_to_saes():
                    with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                        for target_name in self.dict_saes.keys():
                            for step in range(steps):
                                frac = step / steps
                                with self.model.hooks(
                                    fwd_hooks=[
                                        (
                                            lambda name, sae_name=sae_name: sae_name in name, 
                                            partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac)
                                        ) for sae_name in self.dict_saes.keys()
                                    ],
                                    bwd_hooks=[
                                        (
                                            lambda name, target_name=target_name: target_name in name, 
                                            partial(hook_sae_bwd, sae_name=target_name)
                                        ) 
                                    ],
                                ):
                                    metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                                    
        # average the gradients                 
        for key in bwd_cache.keys():
            bwd_cache[key] /= steps
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                sae_hook_name(error_term_name(node.name)) if self.use_error_term else None,
                error_term_name(error_term_name(node.name)) if self.use_esae_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
     
    def _gradient_wrt_edges(
        self,
        clean_token: Tensor,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        node_grads: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects
        Dict[str, SparseAct]
    ]:
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        gradient_mode = kwargs.get('edge_gradient_mode', 'gradient')
            
        _, clean_cache = self.forward(clean_token, corrupt_cache=None)
                    
        edge_grads: Dict[Tuple[Node, Index, Node, Index], Tensor] = {}
        for layer, connection in tqdm(self.connection.items(), disable=not verbose):
            if verbose:
                print(f"Layer {layer}:")
            for hook_position_end, list_hook_positions_start in tqdm(connection.items(), disable=not verbose):
                assert hook_position_end in node_grads, f"Node gradient of {hook_position_end} is not provided."
                
                for hook_position_start in list_hook_positions_start:
                    corrupt_sparse_act = cache_to_sparseact(
                        corrupt_cache, 
                        sae_hook_name(hook_position_start[0].name), 
                        sae_hook_name(error_term_name(hook_position_start[0].name)) if self.use_error_term else None,
                        error_term_name(error_term_name(hook_position_start[0].name)) if self.use_esae_error_term else None,
                    )
                    right_vec = corrupt_sparse_act - clean_cache[hook_position_start[0].name]
                    
                    if gradient_mode == "gradient":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution(
                            clean_token,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            layer,
                            **kwargs,
                        )
                    elif gradient_mode == "ig":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution_ig(
                            clean_token,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            layer,
                            **kwargs,
                        )
                    else:
                        raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
                                  
        return edge_grads, clean_cache
    
    def _edge_attribution(
        self,
        clean_token: Tensor,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach()
                    
            elif sae_hook_name(error_term_name(hook_position_start[0].name)) == hook.name and self.use_error_term:
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(error_term_name(hook_position_start[0].name))] = grad.detach()
                    
            elif output_hook_name(error_term_name(hook_position_start[0].name)) == hook.name and self.use_esae_error_term:
                # only store the gradient at the start position
                # the error_grad is the output_grad, due to: output = recon + error_term
                bwd_cache[error_term_name(error_term_name(hook_position_start[0].name))] = grad.detach()
                
            # IMPORTANT NOTE for reproducibility:
            # we zero grad of intermediate components
            # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads
            # so, we instead zero grad of the input hook of SAE 
            elif hook.name == input_hook_name(sae_name):
                if hook.name != input_hook_name(hook_position_end[0].name):
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict) -> Tensor:
            # cache the feature of SAE post activation
            if hook.name == sae_hook_name(hook_position_end[0].name):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act
            
            # cache the feature of ESAE post activation
            elif hook.name == sae_hook_name(error_term_name(hook_position_end[0].name)):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(error_term_name(hook_position_end[0].name))] = act
                
            elif error_term_name(error_term_name(hook_position_end[0].name)) == hook.name and self.use_esae_error_term:
                to_bwd_cache[error_term_name(error_term_name(hook_position_end[0].name))] = act   
                    
            return act
        
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_esae_end = self.dict_esaes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        d_esae_start = self.dict_esaes[hook_position_start[0].name].cfg.d_sae
        
        to_bwd_cache = {}
        bwd_cache = {}
        edge_effect = OrderedDict()
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self._hook_esaes_to_saes():
                    with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                        with self.model.hooks(
                            fwd_hooks=[
                                (lambda name: True, partial(
                                    hook_sae_fwd, 
                                    to_bwd_cache=to_bwd_cache, 
                                ))
                            ],
                            bwd_hooks=[
                                (lambda name, sae_name=sae_name: sae_name in name, partial(
                                    hook_sae_bwd, 
                                    sae_name=sae_name, 
                                    bwd_cache=bwd_cache,
                                )) for sae_name in self.dict_saes.keys()
                            ],
                        ):
                            self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                            
                            aggregate_dim = [0] if self.token_wise else [0, 1]
                            to_bwd = ( 
                                cache_to_sparseact(
                                    to_bwd_cache, 
                                    sae_hook_name(hook_position_end[0].name), 
                                    sae_hook_name(error_term_name(hook_position_end[0].name)) if self.use_error_term else None,
                                    error_term_name(error_term_name(hook_position_end[0].name)) if self.use_esae_error_term else None,
                                ) * leftvec.detach()
                            ).sum(aggregate_dim)
                            if self.use_esae_error_term:
                                to_bwd = to_bwd.contract()
                            to_bwd = to_bwd.to_tensor() # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                            
                            del to_bwd_cache
                            
                            for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)):
                                if isinstance(end_index, ErrorIndex):
                                    # the last index is error: shape (d_sae+d_esae+1) so last index is d_sae+d_esae
                                    to_bwd[end_index.idx + (d_sae_end+d_esae_end,)].backward(retain_graph=True)
                                    index = t.tensor(list(end_index.idx + (d_sae_end+d_esae_end,)), device=self.device)
                                elif isinstance(end_index, FeatureErrorIndex):
                                    revised_index = list(end_index.idx)
                                    revised_index[-1] += d_sae_end
                                    to_bwd[tuple(revised_index)].backward(retain_graph=True)
                                    index = t.tensor(revised_index, device=self.device)
                                elif isinstance(end_index, FeatureIndex):
                                    to_bwd[end_index.idx].backward(retain_graph=True)
                                    index = t.tensor(list(end_index.idx), device=self.device)
                                else:
                                    raise ValueError(f"end_index of type {type(end_index)} is not supported.")
                                '''
                                edge_effect shape (seq, d_sae+d_esae+1, seq, d_sae+d_esae+1) or (d_sae+d_esae+1, d_sae+d_esae+1) in sparse_coo tensor
                                
                                the sparse_coo will have the shape:
                                --> indices of shape (2, num_active) or (1, num_active)
                                --> values of shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+1)
                                '''
                                effect = ( 
                                    cache_to_sparseact(
                                        bwd_cache, 
                                        sae_hook_name(hook_position_start[0].name), 
                                        sae_hook_name(error_term_name(hook_position_start[0].name)) if self.use_error_term else None,
                                        error_term_name(error_term_name(hook_position_start[0].name)) if self.use_esae_error_term else None,
                                    ) * rightvec
                                ).sum(aggregate_dim)
                                if self.use_esae_error_term:
                                    effect.contract()
                                edge_effect[active_idx] = (index, effect.to_tensor()) # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                            
                            del bwd_cache
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += d_esae_end
            num_start += d_esae_start
        if self.use_esae_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+1)
        # if no active nodes, return empty tensor
        else:
            indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device)
            values = t.empty((0, seq, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device)
            
        if self.token_wise:
            return t.sparse_coo_tensor(indices, values, size=(seq, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
        
    def _edge_attribution_ig(
        self,
        clean_token: Tensor,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        steps = kwargs.get('steps', 10)
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach()
                    
            elif sae_hook_name(error_term_name(hook_position_start[0].name)) == hook.name and self.use_error_term:
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(error_term_name(hook_position_start[0].name))] = grad.detach()
                    
            elif output_hook_name(error_term_name(hook_position_start[0].name)) == hook.name and self.use_esae_error_term:
                # only store the gradient at the start position
                # the error_grad is the output_grad, due to: output = recon + error_term
                bwd_cache[error_term_name(error_term_name(hook_position_start[0].name))] = grad.detach()
                
            # IMPORTANT NOTE for reproducibility:
            # we zero grad of intermediate components
            # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads
            # so, we instead zero grad of the input hook of SAE 
            elif hook.name == input_hook_name(sae_name):
                if hook.name != input_hook_name(hook_position_end[0].name):
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict) -> Tensor:
            # cache the feature of SAE post activation
            if hook.name == sae_hook_name(hook_position_end[0].name):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act
            
            # cache the feature of ESAE post activation
            elif hook.name == sae_hook_name(error_term_name(hook_position_end[0].name)):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(error_term_name(hook_position_end[0].name))] = act
                
            elif error_term_name(error_term_name(hook_position_end[0].name)) == hook.name and self.use_esae_error_term:
                to_bwd_cache[error_term_name(error_term_name(hook_position_end[0].name))] = act   
                    
            return act
        
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_esae_end = self.dict_esaes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        d_esae_start = self.dict_esaes[hook_position_start[0].name].cfg.d_sae
        
        edge_effect = OrderedDict()
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self._hook_esaes_to_saes():
                    with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                        for step in range(steps):
                            frac = step / steps
                            
                            to_bwd_cache = {}
                            bwd_cache = {}
                            with self.model.hooks(
                                fwd_hooks=[
                                    (lambda name: True, partial(
                                        hook_sae_fwd, 
                                        to_bwd_cache=to_bwd_cache, 
                                    ))
                                ],
                                bwd_hooks=[
                                    (lambda name, sae_name=sae_name: sae_name in name, partial(
                                        hook_sae_bwd, 
                                        sae_name=sae_name, 
                                        bwd_cache=bwd_cache,
                                    )) for sae_name in self.dict_saes.keys()
                                ],
                            ):
                                self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                                
                                aggregate_dim = [0] if self.token_wise else [0, 1]
                                to_bwd = ( 
                                    cache_to_sparseact(
                                        to_bwd_cache, 
                                        sae_hook_name(hook_position_end[0].name), 
                                        sae_hook_name(error_term_name(hook_position_end[0].name)) if self.use_error_term else None,
                                        error_term_name(error_term_name(hook_position_end[0].name)) if self.use_esae_error_term else None,
                                    ) * leftvec.detach()
                                ).sum(aggregate_dim)
                                if self.use_esae_error_term:
                                    to_bwd = to_bwd.contract()
                                to_bwd = to_bwd.to_tensor() # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                                
                                del to_bwd_cache
                                
                                for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)):
                                    if isinstance(end_index, ErrorIndex):
                                        # the last index is error: shape (d_sae+d_esae+1) so last index is d_sae+d_esae
                                        to_bwd[end_index.idx + (d_sae_end+d_esae_end,)].backward(retain_graph=True)
                                        index = t.tensor(list(end_index.idx + (d_sae_end+d_esae_end,)), device=self.device)
                                    elif isinstance(end_index, FeatureErrorIndex):
                                        revised_index = list(end_index.idx)
                                        revised_index[-1] += d_sae_end
                                        to_bwd[tuple(revised_index)].backward(retain_graph=True)
                                        index = t.tensor(revised_index, device=self.device)
                                    elif isinstance(end_index, FeatureIndex):
                                        to_bwd[end_index.idx].backward(retain_graph=True)
                                        index = t.tensor(list(end_index.idx), device=self.device)
                                    else:
                                        raise ValueError(f"end_index of type {type(end_index)} is not supported.")
                                    '''
                                    edge_effect shape (seq, d_sae+d_esae+1, seq, d_sae+d_esae+1) or (d_sae+d_esae+1, d_sae+d_esae+1) in sparse_coo tensor
                                    
                                    the sparse_coo will have the shape:
                                    --> indices of shape (2, num_active) or (1, num_active)
                                    --> values of shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+1)
                                    '''
                                    effect = ( 
                                        cache_to_sparseact(
                                            bwd_cache, 
                                            sae_hook_name(hook_position_start[0].name), 
                                            sae_hook_name(error_term_name(hook_position_start[0].name)) if self.use_error_term else None,
                                            error_term_name(error_term_name(hook_position_start[0].name)) if self.use_esae_error_term else None,
                                        ) * rightvec
                                    ).sum(aggregate_dim)
                                    if self.use_esae_error_term:
                                        effect.contract()
                                    if active_idx not in edge_effect:
                                        edge_effect[active_idx] = (index, effect.to_tensor()) # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                                    else:
                                        _, existing_effect = edge_effect[active_idx]
                                        edge_effect[active_idx] = (index, existing_effect + effect.to_tensor())
                                
                                del bwd_cache
                                
        # average the gradients                 
        for active_idx, (prev_index, prev_value) in edge_effect.items():
            edge_effect[active_idx] = (
                prev_index,
                prev_value / steps
            )
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += d_esae_end
            num_start += d_esae_start
        if self.use_esae_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+1)
        # if no active nodes, return empty tensor
        else:
            indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device)
            values = t.empty((0, seq, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device)
            
        if self.token_wise:
            return t.sparse_coo_tensor(indices, values, size=(seq, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
        
    def _prune_nodes(
        self,
        node_effect: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> None:
        if kwargs.get("node_threshold", None) is None:
            raise ValueError("Please provide the node_threshold")
        else:
            node_threshold = kwargs.get("node_threshold")
            assert isinstance(node_threshold, (float, int, Tensor)), "node_threshold must be a int, float, or Tensor"
            
            for node, index in tqdm(node_effect.keys(), disable=not verbose):
                if kwargs.get("reverse_prune_node", False):
                    effect_feat_mask = node_effect[(node, index)].act.abs() < node_threshold
                else:
                    effect_feat_mask = node_effect[(node, index)].act.abs() > node_threshold
                    
                if verbose:
                    num_pruned = (~effect_feat_mask).sum().item()
                    print(f"{(node, index)}: pruned ({num_pruned / effect_feat_mask.numel() * 100:.5f}%) features")
                    
                if self.use_error_term:
                    if kwargs.get("reverse_prune_node", False):
                        effect_feat_error_mask = node_effect[(node, index)].res.abs() < node_threshold # type: ignore
                    else:
                        effect_feat_error_mask = node_effect[(node, index)].res.abs() > node_threshold # type: ignore
                        
                    if verbose:
                        num_pruned = (~effect_feat_error_mask).sum().item()
                        print(f"{(node, index)}: pruned ({num_pruned / effect_feat_error_mask.numel() * 100:.5f}%) error features")
                
                if self.use_esae_error_term:
                    if kwargs.get("reverse_prune_node", False):
                        effect_error_mask = node_effect[(node, index)].resc.abs() < node_threshold # type: ignore
                    else:
                        effect_error_mask = node_effect[(node, index)].resc.abs() > node_threshold # type: ignore
                        
                    if verbose:
                        num_pruned = (~effect_error_mask).sum().item()
                        print(f"{(node, index)}: pruned ({num_pruned / effect_error_mask.numel() * 100:.5f}%) errors")
                        
                # replace the nodes with mask to prune
                mask = SparseAct(
                    act=effect_feat_mask,
                    res=effect_feat_error_mask if self.use_error_term else None, # type: ignore
                    resc=effect_error_mask if self.use_esae_error_term else None, # type: ignore
                )
                self.nodes[(node, index)] = self.nodes[(node, index)] * mask
    
    def run_model(
        self, 
        toks: Tensor, 
        use_error_term: bool | None = None
    ) -> Tuple[Tensor, ActivationCache]:
        
        with self._hook_esaes_to_saes(use_esae_error_term=use_error_term):
            out = self.model.run_with_cache_with_saes(
                toks,
                saes=self._saes_to_list(),
                use_error_term=True,
                names_filter=lambda name: "sae" in name,
            )
        
        return out # type: ignore
    
    def _saes_to_list(self) -> List[Any]:
        return [sae for _, sae in self.dict_saes.items()]
    
    def _check_valid_esaes(self):
        for hook_name in self.dict_saes.keys():
            assert hook_name in self.dict_esaes, f"Hook name {hook_name} is not found in ESAEs."
    
    @contextmanager
    def _hook_esaes_to_saes(self, reset_esaes_end: bool = True, reset_saes_error_term: bool = True, use_esae_error_term: bool | None = None):
        self._check_valid_esaes()
        
        use_esae_error_term = self.use_esae_error_term if use_esae_error_term is None else use_esae_error_term
        orig_use_error_term = {}
        try:
            for sae_name, sae in self.dict_saes.items():
                sae.add_sae(self.dict_esaes[sae_name], use_error_term=use_esae_error_term)
                
                orig_use_error_term[sae_name] = sae.use_error_term
                sae.use_error_term = self.use_error_term
            yield
        finally:
            for sae_name, sae in self.dict_saes.items():
                if reset_esaes_end:
                    sae.reset_saes()
                if reset_saes_error_term:
                    sae.use_error_term = orig_use_error_term[sae_name]
                    
    @contextmanager
    def _detach_error_term(self, detach: bool, sae_name: str | None = None):
        orig_detach_error_term = {}
        orig_disable_error_grad = {}
        orig_detach_esae_error_term = {}
        orig_disable_esae_error_grad = {}
        
        def __detach(dict_saes: Dict[str, Any], detach: bool, detach_error_term: Dict, disable_error_grad: Dict, sae_name: str | None = None):
            for name, sae in dict_saes.items():
                if sae_name is None or sae_name in name:
                    detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = detach
                else:
                    detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = True # default detach error term
                    
                disable_error_grad[name] = sae.disable_error_grad
                sae.disable_error_grad = False # allow grad flows through error hook

        def __recover(dict_saes: Dict[str, Any], detach_error_term: Dict, disable_error_grad: Dict):
            for name, sae in dict_saes.items():
                sae.detach_error_term = detach_error_term[name]
                sae.disable_error_grad = disable_error_grad[name]
                
        try:
            __detach(
                self.dict_saes, detach, orig_detach_error_term, orig_disable_error_grad, sae_name
            )
            __detach(
                self.dict_esaes, detach, orig_detach_esae_error_term, orig_disable_esae_error_grad, sae_name
            )
            yield
        finally:
            __recover(
                self.dict_saes, orig_detach_error_term, orig_disable_error_grad
            )
            __recover(
                self.dict_esaes, orig_detach_esae_error_term, orig_disable_esae_error_grad
            )
        
    

if __name__ == "__main__":
    '''
    To debug graph
    '''
    from sae_lens import SAE
    device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
    
    gpt2: HookedSAETransformer = HookedSAETransformer.from_pretrained("gpt2-small", device=device)
    
    list_hook_names = [utils.get_act_name("resid_pre", layer) for layer in range(gpt2.cfg.n_layers)]
    dict_saes = {}
    list_saes: List[SAE] = []
    for layer, hook_name in enumerate(list_hook_names):
        gpt2_sae, _, _ = SAE.from_pretrained(
            release="gpt2-small-res-jb",
            sae_id=hook_name,
            device=str(device),
        )
        dict_saes[layer] = [(hook_name, gpt2_sae)]
        list_saes.append(gpt2_sae)

    use_error_term = True
        
    clean_data = "hello, my name is T"
    corrupt_data = "hi, his name is T"
    clean_token = gpt2.to_tokens(clean_data)
    corrupt_token = gpt2.to_tokens(corrupt_data)
    
    fg = Feature_Graph(
        gpt2, 
        dict_saes, 
        use_error_term=use_error_term,
    )
    
    fg.build_default_connection(clean_token.shape[1], token_wise=True, inter=True, intra=False)
    
    import random
    d_sae = list_saes[0].cfg.d_sae
    seq = clean_token.shape[1]
    n_layers = gpt2.cfg.n_layers
    delete_feats = [
        (
            ConnectionNode(f"blocks.{random.randint(0, n_layers-1)}.hook_resid_pre"),
            ConnectionIndex(), 
            FeatureIndex((random.randint(0, seq-1), random.randint(0, d_sae-1))),
        ) for _ in range(5000)
    ]
    delete_errors = [
        (
            ConnectionNode(f"blocks.{random.randint(0, n_layers-1)}.hook_resid_pre"),
            ConnectionIndex(), 
            ErrorIndex((random.randint(0, seq-1),)),
        ) for _ in range(3)
    ]
    if use_error_term:
        for error in delete_errors:
            fg.delete_single_error(*error)
    for feat in delete_feats:
        fg.delete_single_feature(*feat)
            
    print(fg.active_nodes(ConnectionNode("blocks.0.hook_resid_pre"), ConnectionIndex())[-5])
    
    _, corrupt_cache = gpt2.run_with_cache_with_saes(
        corrupt_token,
        saes=list_saes,
        use_error_term=use_error_term,
        names_filter=lambda name: "sae" in name,
    )
    
    fg_logit, patched_cache = fg(clean_token, corrupt_cache, patch_deleted_comp=True)
    
       
    if not use_error_term: 
        for sae in list_saes:
            sae.use_error_term = False
            
        def ablate_sae_latent(
            sae_acts: Tensor,
            hook: HookPoint,
        ) -> Tensor:
            for sae_name, _, feat_idx in delete_feats:
                if hook.name == sae_hook_name(sae_name.name):
                    sae_acts[ConnectionIndex(return_idx(feat_idx.list_index, 3)).as_index] = corrupt_cache[hook.name][ConnectionIndex(return_idx(feat_idx.list_index, 3)).as_index] # type: ignore
            
            return sae_acts
        logit = gpt2.run_with_hooks_with_saes(
            clean_token,
            saes=list_saes,
            fwd_hooks=[(lambda name: "sae" in name, ablate_sae_latent)],
        )
    else:
        for sae in list_saes:
            sae.use_error_term = True
        def ablate_sae_latent(
            sae_acts: Tensor,
            hook: HookPoint,
        ) -> Tensor:
            for sae_name, _, feat_idx in delete_feats:
                if hook.name == sae_hook_name(sae_name.name):
                    sae_acts[ConnectionIndex(return_idx(feat_idx.list_index, 3)).as_index] = corrupt_cache[hook.name][ConnectionIndex(return_idx(feat_idx.list_index, 3)).as_index] # type: ignore
            for sae_name, _, error_idx in delete_errors:
                if hook.name == error_term_name(sae_name.name):
                    sae_acts[ConnectionIndex(return_idx(error_idx.list_index, 2)).as_index] = corrupt_cache[hook.name][ConnectionIndex(return_idx(error_idx.list_index, 2)).as_index] # type: ignore
            
            return sae_acts
        logit = gpt2.run_with_hooks_with_saes(
            clean_token,
            saes=list_saes,
            fwd_hooks=[(lambda name: "sae" in name, ablate_sae_latent)],
        )
    t.testing.assert_close(logit, fg_logit, rtol=1e-5, atol=1e-5)
    
    # '''
    # To debug gradients
    # '''
    
    # device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
    
    # gpt2: HookedSAETransformer = HookedSAETransformer.from_pretrained("gpt2-small", device=device)

    # list_hook_names = [utils.get_act_name("resid_pre", layer) for layer in range(gpt2.cfg.n_layers)]
    # dict_saes = {}
    # list_saes = []
    # for layer, hook_name in enumerate(list_hook_names):
    #     gpt2_sae, _, _ = SAE.from_pretrained(
    #         release="gpt2-small-res-jb",
    #         sae_id=hook_name,
    #         device=str(device),
    #     )
    #     dict_saes[layer] = [(hook_name, gpt2_sae)]
    #     list_saes.append(gpt2_sae)
        
    # import random
    # from itertools import product

    # names = [" John", " Mary"]
    # name_tokens = [gpt2.to_single_token(name) for name in names]

    # prompt_template = "When{A} and{B} went to the shops,{S} gave the bag to"
    # prompts = [
    #     prompt_template.format(A=names[i], B=names[1 - i], S=names[j]) for i, j in product(range(2), range(2))
    # ]
    # correct_answers = names[::-1] * 2
    # incorrect_answers = names * 2
    # correct_toks = gpt2.to_tokens(correct_answers, prepend_bos=False)[:, 0].tolist()
    # incorrect_toks = gpt2.to_tokens(incorrect_answers, prepend_bos=False)[:, 0].tolist()
    # corrupt_prompts = [
    #     prompt_template.format(A=names[i], B=names[1 - i], S=names[1 - j]) for i, j in product(range(2), range(2))
    # ]
    # clean_token = gpt2.to_tokens(prompts)
    # corrupt_token = gpt2.to_tokens(corrupt_prompts)


    # def logits_to_ave_logit_diff(
    #     logits: Tensor,
    #     correct_toks: list[int] = correct_toks,
    #     incorrect_toks: list[int] = incorrect_toks,
    #     reduction: str | None = "mean",
    # ) -> Tensor:
    #     """
    #     Returns the avg logit diff on a set of prompts, with fixed s2 pos and stuff.
    #     """
    #     correct_logits = logits[range(len(logits)), -1, correct_toks]
    #     incorrect_logits = logits[range(len(logits)), -1, incorrect_toks]
    #     logit_diff = correct_logits - incorrect_logits
    #     if reduction is not None:
    #         logit_diff = logit_diff.mean() if reduction == "mean" else logit_diff.sum()
    #     return logit_diff

    # names_filters = [gpt2_sae.cfg.hook_name + ".hook_sae_acts_post" for gpt2_sae in list_saes]

    # use_error_term = True

    # fg = Feature_Graph(
    #     gpt2, 
    #     dict_saes, 
    #     use_error_term=use_error_term,
    # )

    # fg.build_default_connection(clean_token.shape[1], token_wise=True, inter=True, intra=True)

    # corrupt_logit, corrupt_cache = gpt2.run_with_cache_with_saes(
    #     corrupt_token,
    #     saes=list_saes,
    #     use_error_term=use_error_term,
    #     names_filter=lambda name: "sae" in name,
    # )
    # clean_logit = gpt2.run_with_saes(clean_token, saes=list_saes, use_error_term=use_error_term)

    # ioi_metric = partial(logits_to_ave_logit_diff, reduction='sum')

    # node_grads, _, clean_cache = fg.forward_backward_gradient(
    #     clean_token, corrupt_cache, ioi_metric, show_warnings=True, retain_graph=True, mode='node', patch_deleted_comp=True
    # )

    # attrib_effect = {}
    # for (node, index), grad in node_grads.items():
    #     corrupt_sparse_act = SparseAct(
    #         act=corrupt_cache[sae_hook_name(node.name)],
    #         res=corrupt_cache[error_term_name(node.name)] if use_error_term else None,
    #     )
    #     a = clean_cache[node.name]
    #     attrib_effect[(node, index)] = (
    #         grad @
    #         (corrupt_sparse_act - a)
    #     ).sum(0)
        
    # name_to_try = "blocks.2.hook_resid_pre"
    
    # s2_pos = 10
    # assert gpt2.to_str_tokens(prompts[0])[s2_pos] == " John"

    # active_latents = (clean_cache[name_to_try].act[:, s2_pos] > 0.0).any(0).nonzero().cpu().tolist()

    # ablation_effect = []
    # clean_metric = ioi_metric(clean_logit) # type: ignore
    # all_list_to_try = []
    # with t.no_grad():
    #     for active_node_idx in tqdm(active_latents):
    #         list_to_try = [active_node_idx]
    #         all_list_to_try.append(list_to_try)
            
    #         fwd_cache = {}
    #         def ablate_sae_latent(
    #             sae_acts: Tensor,
    #             hook: HookPoint,
    #         ) -> Tensor:
    #             if hook.name == sae_hook_name(name_to_try):
    #                 for node_idx in list_to_try:
    #                     sae_acts[:, s2_pos, node_idx] = corrupt_cache[sae_hook_name(name_to_try)][:, s2_pos, node_idx]
    #             if hook.name == output_hook_name(f"blocks.{hook.layer()}.hook_resid_pre"):
    #                 if use_error_term:
    #                     sae_acts += clean_cache[f"blocks.{hook.layer()}.hook_resid_pre"].res
    #             return sae_acts

    #         # NOTE: use_error_term = True might be wrong
    #         gpt2.reset_hooks()
    #         patched_logits = gpt2.run_with_hooks_with_saes(
    #             clean_token,
    #             saes=list_saes,
    #             # fwd_hooks=[(names_filters[i], ablate_sae_latent) for i in range(len(names_filters))],
    #             fwd_hooks=[(lambda name: "sae" in name, ablate_sae_latent)],
    #         )
    #         ablation_effect.append((ioi_metric(patched_logits) - clean_metric).cpu().item())
            
            
    # attrib_list = [
    #     sum([
    #         attrib_effect[ConnectionNode(name_to_try), ConnectionIndex()].act[s2_pos, active_node_idx].cpu().item() 
    #         for active_node_idx in list_to_try
    #     ])
    #     for list_to_try in all_list_to_try
    # ]

    # from plotly import express as px
    # import pandas as pd
    # df = pd.DataFrame({
    #     'ablation': ablation_effect, 
    #     'attribution': attrib_list,
    # })

    # # Create scatter plot
    # fig = px.scatter(df, x='attribution', y='ablation', labels={'attribution': 'attribution', 'ablation': 'ablation'},
    #                 title='Scatter Plot of Two Dictionaries')

    # # Add text labels
    # fig.update_traces(textposition='top center')

    # # Add y = x line
    # fig.add_shape(type='line', x0=min(df['ablation']), y0=min(df['ablation']), 
    #             x1=max(df['ablation']), y1=max(df['ablation']),
    #             line=dict(color='red', dash='dash'))

    # # Show plot
    # fig.show()