from tracemalloc import start
from typing import Literal
from .Feature_Graph_Trans import *
from transformer_lens.utils import get_act_name
        
def _hook_name(sae_name: str, name: str | None):
    return sae_name + "." + str(name)

class Feature_Graph_Cross(Feature_Graph_Trans):
    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,
    ):
        super().__init__(model, saes, use_error_term)
                
    def process_transcoder(self):
        self.input_hooks = []
        self.output_hooks = []
        self.crosscoders = []
        self.non_crosscoders = []
        check_mlp_out = False
        self.check_attn_out = False
        self.check_resid_pre = False
        for sae in self.dict_saes.values():
            sae = sae.to(self.device)
            if getattr(sae, 'input_hook', False) and getattr(sae, 'output_hooks', False):
                self.input_hooks.append(sae.input_hook)
                self.output_hooks.append(sae.output_hooks[0])
                self.crosscoders.append(sae)
                
                if "mlp_out" in sae.output_hooks[0]:
                    check_mlp_out = True
            else:
                self.non_crosscoders.append(sae)
                
            if "attn_out" in sae.cfg.hook_name:
                self.check_attn_out = True
            if "resid_pre" in sae.cfg.hook_name:
                self.check_resid_pre = True
            
        assert check_mlp_out, "Crosscoder needs to be provided at the mlp_out hook."
     
    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()

        fwd_cache = {}
        bwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                for target_name in self.dict_saes.keys():
                    for step in range(steps):
                        frac = step / steps
                        with self._setup_fwd_bwd_grad_sae_hook_ig(
                            target_name=target_name,
                            frac=frac,
                            fwd_cache=fwd_cache,
                            bwd_cache=bwd_cache,
                            corrupt_cache=corrupt_cache,
                        ):
                            with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                                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_nodes_vw_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]
    ]:
        '''
        Using virtual weight to compute node grad
        '''
        steps = kwargs.get("steps", 10)
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        
        sink_hook_name = get_act_name("resid_post", self.n_layers-1) # last layer resid post hook
        
        sink_node_cache: Dict[str, Tensor] = {}
        def hook_bwd(tens: Tensor, hook: HookPoint):
            sink_node_cache[hook.name] = tens.detach() # type: ignore
        
        all_fwd_cache = {}
        all_bwd_cache = {}
        for target_name in self.dict_saes.keys():
            for step in range(steps):
                frac = step / steps

                fwd_cache = {}
                with t.set_grad_enabled(True):
                    with self._detach_error_term(True):
                        with self.model.hooks(
                            bwd_hooks=[(sink_hook_name, hook_bwd)]
                        ):
                            with self._setup_fwd_bwd_grad_sae_hook_ig(
                                target_name=target_name,
                                frac=frac,
                                fwd_cache=fwd_cache, 
                                bwd_cache={}, 
                                corrupt_cache=corrupt_cache, 
                            ):
                                with self._setup_forward_model_hook(transfer_grad=False): # save memory, we don't need transfer_grad here
                                    metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                                    
                # run the model on interpolate intervention
                with self._setup_forward_model_hook(transfer_grad=False): # save memory, we don't need transfer_grad here
                    with self._setup_virtual_weight_sae_hook_ig(
                        target_name=target_name,
                        frac=frac,
                        corrupt_cache=corrupt_cache, 
                    ):
                        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)   
                                
                current_grad = sink_node_cache[sink_hook_name] # (b, seq, d_model)
                
                bwd_cache = self._TE_using_virtual_weight(
                    current_grad,
                    fwd_cache=fwd_cache,
                    unpatch_clean_cache=unpatch_clean_cache,
                )
                
                if step == 0:
                    add_cache(all_fwd_cache, fwd_cache)
                
                add_cache(all_bwd_cache, bwd_cache)
                
        # average the gradients                 
        for key in all_bwd_cache.keys():
            all_bwd_cache[key] /= steps
            
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = all_bwd_cache[node.name]
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                all_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 _TE_using_virtual_weight(
        self,
        current_grad: Tensor,
        fwd_cache: ActivationCache | Dict[str, Tensor],
        unpatch_clean_cache: ActivationCache | Dict[str, Tensor],
    ) -> Dict[str, Tensor]:
        bwd_cache = {}
        resid_grad_cache = {}
        for layer in reversed(range(self.n_layers)):
            mlp_name = get_act_name("mlp_out", layer)
            attn_name = get_act_name("attn_out", layer)
            resid_pre_name = get_act_name("resid_pre", layer)
            
            resid_grad_cache[mlp_name] = current_grad # (b, seq, d_model)
            # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
            bwd_cache[mlp_name] = SparseAct(
                act=current_grad @ self.dict_saes[mlp_name].W_dec.T,
                res = current_grad if self.use_error_term else None,
            )
            for upper_layer in range(layer+1, self.n_layers):
                upper_mlp_name = get_act_name("mlp_out", upper_layer)
                bwd_cache[mlp_name] += SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                    act = resid_grad_cache[upper_mlp_name] @ self.dict_saes[mlp_name].crosscoder_decoders[upper_layer-layer-1].weight,
                    res = 0 if self.use_error_term else None, # type: ignore
                )
            
            grad_through_mlp = gradient_with_mlp(
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                transcoder_enc=self.dict_saes[mlp_name].W_enc, # (d_model, d_sae)
                transcoder_dec=self.dict_saes[mlp_name].W_dec, # (d_sae, d_model)
                transcoder_act=fwd_cache[sae_hook_name(mlp_name)], # (b, seq, d_sae)
                layer_end=layer,
                pos_end=None, 
                seq_length=self.seq_length, # type: ignore
                batch_size=current_grad.shape[0],
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_mlp # resid_mid_grad
            
            if self.check_attn_out:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[attn_name] = SparseAct(
                    act=current_grad @ self.dict_saes[attn_name].W_dec.T,
                    res = current_grad if self.use_error_term else None,
                )
            
            grad_through_attn = gradient_with_attn(
                model=self.model,
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=layer,
                pos_end=None, 
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_attn # resid_pre_grad
            
            if self.check_resid_pre:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[resid_pre_name] = SparseAct(
                    act=current_grad @ self.dict_saes[resid_pre_name].W_dec.T,
                    res = current_grad if self.use_error_term else None,
                )
                
        return bwd_cache 
    
    def _edge_attribution_trans(
        self,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        **kwargs,
    ) -> Tensor:
        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
        start_layer = int(hook_position_start[0].name.split(".")[1])
        end_layer = int(hook_position_end[0].name.split(".")[1])
        
        aggregate_dim = [0] if self.token_wise else [0, 1]
        edge_effect = {}        
        all_error = []
        
        for end_node, end_index in self.active_nodes(*hook_position_end):
            if isinstance(end_index, ErrorIndex):
                all_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureIndex):
                feat_id = end_index.idx[-1]
                pos_end = end_index.idx[-2]
                index = t.tensor(list(end_index.idx), device=self.device)
                end_node_grad = leftvec.act[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1)
                
                end_feature_vec = self.dict_saes[end_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                
                if "mlp_out" in hook_position_start[0].name:
                    list_grad_dot_leftvec_tensor = self._DE_using_virtual_weight_cross( # (b, seq, d_model)
                        grad=end_feature_vec,
                        pos_end=pos_end,
                        batch_size=leftvec.act.shape[0],
                        unpatched_clean_cache=unpatched_clean_cache,
                        hook_position_end=hook_position_end,
                        hook_position_start=hook_position_start,
                    )
                    
                    for i in range(len(list_grad_dot_leftvec_tensor)):
                        list_grad_dot_leftvec_tensor[i] *= end_node_grad  # (b, seq, d_model)
                        
                    grad_dot_leftvec = SparseAct(
                        # (b, seq, d_model) @ (d_model, d_sae)
                        act=list_grad_dot_leftvec_tensor[-1] @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                        res=list_grad_dot_leftvec_tensor[-1] if self.use_error_term else None
                    )
                    for i, upper_layer in enumerate(reversed(range(start_layer+1, end_layer))): # (start+1 -> end_layer-1)
                        rel_id = upper_layer-start_layer-1
                        grad_dot_leftvec += SparseAct(
                            # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                            act = (
                                list_grad_dot_leftvec_tensor[i] @ 
                                self.dict_saes[hook_position_start[0].name].crosscoder_decoders[rel_id].weight
                            ),
                            res = 0 if self.use_error_term else None, # type: ignore
                        )
                else:
                    grad_dot_leftvec_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                        grad=end_feature_vec,
                        pos_end=pos_end,
                        batch_size=leftvec.act.shape[0],
                        unpatched_clean_cache=unpatched_clean_cache,
                        hook_position_end=hook_position_end,
                        hook_position_start=hook_position_start,
                    ) * end_node_grad
                    grad_dot_leftvec = SparseAct(
                        # (b, seq, d_model) @ (d_model, d_sae)
                        act=grad_dot_leftvec_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                        res=grad_dot_leftvec_tensor if self.use_error_term else None
                    )
                
                '''
                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[index] = ( # (seq, d_sae+1) || (d_sae+1)
                    grad_dot_leftvec @ rightvec
                ).sum(aggregate_dim).to_tensor()
            else:
                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
            
        '''
        The gradient of error node to upstream node is 1 - gradient of sum end_feature_node
        We multiply error grad so that we only have to backward once (Jacobian vector product)
        The "end_feature_dependent" sums all of the gradient of end_feature_node.
        '''
        if self.use_error_term:
            all_end_node_grad: Tensor = leftvec.res # type: ignore
            feature_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_saes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_sae_end d_model -> b seq d_sae_end",
            )
            
            end_feature_dependent = einops.einsum(
                feature_coef_to_cal_error_edge, self.dict_saes[hook_position_end[0].name].W_enc,
                "b seq d_sae_end, d_model d_sae_end -> b seq d_model"
            )
            
            for end_error_node, end_error_index in all_error:
                pos_end = end_error_index.idx[0]
                index = t.tensor(list(end_error_index.idx + (d_sae_end,)), device=self.device)
                end_node_grad = leftvec.res[:, pos_end].unsqueeze(1) # (b, 1, d_model) # type: ignore
                
                end_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_node_grad - end_feature_dependent[:, pos_end].unsqueeze(1), # (b, 1, d_model)
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_error_grad_tensor if self.use_error_term else None
                )

                edge_effect[index] = ( # (seq, d_sae_start+1) | (d_sae_start+1)
                    end_error_grad @ rightvec
                ).sum(aggregate_dim).to_tensor()  
                    
        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(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value 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 _DE_using_virtual_weight_cross(
        self,
        grad: Tensor, # (1, 1, d_model) or (b, 1, d_model)
        pos_end: int,
        batch_size: int,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
    ) -> List[Tensor]:
        start_layer = int(hook_position_start[0].name.split(".")[1])
        end_layer = int(hook_position_end[0].name.split(".")[1])
        
        path = [] # path attention for gradient, if we have attention circuit then the path is empty
        if not self.check_attn_out:
            if not "resid_pre" in hook_position_end[0].name:
                path.append(("attn", end_layer)) # attn or mlp at end hook -> goes through attn at end_layer
                
            for layer in reversed(range(start_layer+1, end_layer)):
                path.append(("attn", layer))
                
            if "resid_pre" in hook_position_start[0].name:
                path.append(("attn", start_layer)) # resid_pre at start hook -> goes through attn at start_layer
        else:
            for layer in reversed(range(start_layer+1, end_layer)):
                path.append(("no_attn_grad", layer))
                
        if "mlp_out" in hook_position_end[0].name:
            ln = "ln2"
        elif "attn_out" in hook_position_end[0].name:
            ln = "ln1"
        else:
            ln = None
        
        if self.check_attn_out and "attn_out" in hook_position_end[0].name:
            current_grad = gradient_with_attn(
                model=self.model,
                end_feature_vec=grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=end_layer,
                pos_end=pos_end, 
                use_error_term=self.use_error_term,
                cache=unpatched_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
        else:
            current_grad = gradient_ln_only( # (b, seq, d_model) 
                end_feature_vec=grad,
                start_feature_vec=None,
                layer_end=end_layer,
                pos_end=pos_end,
                seq_length=self.seq_length, # type: ignore
                batch_size=batch_size,
                use_error_term=self.use_error_term,
                cache=unpatched_clean_cache,
                device=self.device, # type: ignore
                ln=ln,
            ).act 
            
        all_resid_grad = [] # we don't need to cache the end_layer grad
        
        for name, layer in path:
            if name == "attn":
                grad_through_attn = gradient_with_attn(
                    model=self.model,
                    end_feature_vec=current_grad, # (b, seq, d_model)
                    start_feature_vec=None,
                    layer_end=layer,
                    pos_end=None, 
                    use_error_term=self.use_error_term,
                    cache=unpatched_clean_cache,
                    device=self.device # type: ignore
                ).act # (b, seq, d_model)
                
                current_grad = current_grad + grad_through_attn
            
            if layer in range(start_layer+1, end_layer):
                all_resid_grad.append(current_grad) # cache the gradient of each intermediate layer
        
        all_resid_grad.append(current_grad) # cache the gradient of the start layer
        
        return all_resid_grad
             
    @contextmanager
    def _setup_forward_model_hook(self, use_error_term: bool | None = None, transfer_grad: bool = True):
        cross_cache = {}
        # Hook function at transcoder input: caches the activations before transcoder
        def hook_crosscoder_input(activations: Tensor, hook: HookPoint, crosscoder_idx: int):
            cross_cache[crosscoder_idx] = activations.clone()
        
        all_cross_recons = [0.0 for _ in range(len(self.crosscoders))]
        def hook_crosscoder_output(activations: Tensor, hook: HookPoint, crosscoder_idx: int):
            recons, cross_recons = self.crosscoders[crosscoder_idx].forward_crosscoder(
                (cross_cache[crosscoder_idx], activations, all_cross_recons[crosscoder_idx])
            )
            for j in range(crosscoder_idx+1, self.n_layers):
                all_cross_recons[j] += cross_recons[j-crosscoder_idx-1]
            
            if transfer_grad:
                return recons + (activations - activations.detach())
            else:
                return recons
        
        fwd_hooks = []
        for i in range(len(self.crosscoders)):
            fwd_hooks.append((self.input_hooks[i], partial(hook_crosscoder_input, crosscoder_idx=i)))
            fwd_hooks.append((self.output_hooks[i], partial(hook_crosscoder_output, crosscoder_idx=i)))
            
        use_error_term = use_error_term if use_error_term is not None else self.use_error_term
        try: 
            for sae in self._saes_to_list():
                self.model.add_sae(sae, use_error_term)
            for hook, func in fwd_hooks:
                self.model.add_hook(hook, func, dir="fwd")
            
            yield
        finally:
            self.model.reset_saes()
            self.model.reset_hooks() 
            
    def _saes_to_list(self) -> List[Any]:
        return self.non_crosscoders
    

class ESAE_FG_Cross(
    ESAE_FG_Trans,
    Feature_Graph_Cross,
):
    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, esaes, use_esae_error_term)
        
    def process_transcoder(self):
        Feature_Graph_Cross.process_transcoder(self)
                
    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 = {}
        
        with t.no_grad():
            with self._setup_forward_model_hook(transfer_grad=False):
                with self._setup_fwd_sae_hook(
                    fwd_cache=fwd_cache, corrupt_cache=corrupt_cache, patch_deleted_comp=patch_deleted_comp
                ):
                    logits = self.model(clean_token)
        
        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,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
        
        return logits, 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()

        fwd_cache = {}
        bwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                for target_name in self.dict_saes.keys():
                    for step in range(steps):
                        frac = step / steps
                        with self._setup_fwd_bwd_grad_sae_hook_ig(
                            target_name=target_name,
                            frac=frac,
                            fwd_cache=fwd_cache,
                            bwd_cache=bwd_cache,
                            corrupt_cache=corrupt_cache,
                        ):
                            with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                                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 _TE_using_virtual_weight(
        self,
        current_grad: Tensor,
        fwd_cache: ActivationCache | Dict[str, Tensor],
        unpatch_clean_cache: ActivationCache | Dict[str, Tensor],
    ) -> Dict[str, Tensor]:
        
        return Feature_Graph_Cross._TE_using_virtual_weight(
            self,
            current_grad=current_grad,
            fwd_cache=fwd_cache,
            unpatch_clean_cache=unpatch_clean_cache,
        )
         
    def _edge_attribution_trans(
        self,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        **kwargs,
    ) -> Tensor:
        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
        d_esae_end = self.dict_esaes[hook_position_end[0].name].cfg.d_sae
        d_esae_start = self.dict_esaes[hook_position_start[0].name].cfg.d_sae
        start_layer = int(hook_position_start[0].name.split(".")[1])
        end_layer = int(hook_position_end[0].name.split(".")[1])
        
        aggregate_dim = [0] if self.token_wise else [0, 1]
        edge_effect = {}        
        all_error = []
        all_feature_error = []
        
        for end_node, end_index in self.active_nodes(*hook_position_end):
            if isinstance(end_index, ErrorIndex):
                all_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureErrorIndex):
                all_feature_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureIndex):
                feat_id = end_index.idx[-1]
                pos_end = end_index.idx[-2]
                index = t.tensor(list(end_index.idx), device=self.device)
                end_node_grad = leftvec.act[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1)
                
                end_feature_vec = self.dict_saes[end_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                
                if "mlp_out" in hook_position_start[0].name:
                    list_grad_dot_leftvec_tensor = self._DE_using_virtual_weight_cross( # (b, seq, d_model)
                        grad=end_feature_vec,
                        pos_end=pos_end,
                        batch_size=leftvec.act.shape[0],
                        unpatched_clean_cache=unpatched_clean_cache,
                        hook_position_end=hook_position_end,
                        hook_position_start=hook_position_start,
                    )
                    
                    for i in range(len(list_grad_dot_leftvec_tensor)):
                        list_grad_dot_leftvec_tensor[i] *= end_node_grad  # (b, seq, d_model)
                        
                    grad_dot_leftvec = SparseAct(
                        # (b, seq, d_model) @ (d_model, d_sae)
                        act=list_grad_dot_leftvec_tensor[-1] @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                        res=list_grad_dot_leftvec_tensor[-1] @ self.dict_esaes[hook_position_start[0].name].W_dec.T if self.use_error_term else None,
                        resc=list_grad_dot_leftvec_tensor[-1] if self.use_esae_error_term else None
                    )
                    for i, upper_layer in enumerate(reversed(range(start_layer+1, end_layer))): # (start+1 -> end_layer-1)
                        rel_id = upper_layer-start_layer-1
                        grad_dot_leftvec += SparseAct(
                            # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                            act = (
                                list_grad_dot_leftvec_tensor[i] @ 
                                self.dict_saes[hook_position_start[0].name].crosscoder_decoders[rel_id].weight
                            ),
                            res = 0 if self.use_error_term else None, # type: ignore
                            resc = 0 if self.use_esae_error_term else None, # type: ignore
                        )
                else:
                    grad_dot_leftvec_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                        grad=end_feature_vec,
                        pos_end=pos_end,
                        batch_size=leftvec.act.shape[0],
                        unpatched_clean_cache=unpatched_clean_cache,
                        hook_position_end=hook_position_end,
                        hook_position_start=hook_position_start,
                    ) * end_node_grad
                    grad_dot_leftvec = SparseAct(
                        # (b, seq, d_model) @ (d_model, d_sae)
                        act=grad_dot_leftvec_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                        res=grad_dot_leftvec_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T if self.use_error_term else None,
                        resc=grad_dot_leftvec_tensor if self.use_esae_error_term else None,
                    )
                '''
                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)
                '''
                effect = (
                    grad_dot_leftvec * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
            else:
                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
            
        '''
        The gradient of feature error node to upstream node is f_esae_enc - sum gradient of end_feature_node
        The "grad_through_end_feat_error" computes the jacobian of f_esae_enc going through f_sae_dec and f_sae_enc
        '''
        if self.use_error_term:            
            grad_through_end_feat_error = einops.einsum( 
                self.dict_saes[hook_position_end[0].name].W_dec.T, self.dict_saes[hook_position_end[0].name].W_enc.T,
                "d_model1 d_sae_end, d_sae_end d_model2 -> d_model1 d_model2"
            )
            grad_through_end_feat_error = einops.einsum( # gradient through sae_feature
                self.dict_esaes[hook_position_end[0].name].W_enc.T, grad_through_end_feat_error,
                "d_esae_end d_model1, d_model1 d_model2 -> d_esae_end d_model2"
            )
            
            for end_feature_error_node, end_feature_error_index in all_feature_error:
                pos_end = end_feature_error_index.idx[0]
                feat_id = end_feature_error_index.idx[-1]
                
                revised_index = list(end_feature_error_index.idx)
                revised_index[-1] += d_sae_end
                index = t.tensor(revised_index, device=self.device)
                
                end_node_grad = leftvec.res[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1) # type: ignore
                end_feature_vec = self.dict_esaes[end_feature_error_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                
                end_feature_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_feature_vec - grad_through_end_feat_error[feat_id].unsqueeze(0).unsqueeze(0), # (1, 1, d_model)
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_feature_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_feature_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_feature_error_grad_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T,
                    resc=end_feature_error_grad_tensor if self.use_esae_error_term else None,
                )
                effect = (
                    end_feature_error_grad * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
                
        '''
        The gradient of feature error node to upstream node is:
            f_esae_enc - sum gradient of end_feature_node - sum gradient of end_feaeture_error_node
            
        
        The gradient of feature error node to upstream node is f_esae_enc - sum gradient of end_feature_node (see above)
        The "grad_through_end_feat_error" computes the jacobian of f_esae_enc going through f_sae_dec and f_sae_enc
        The "feature_error_coef_to_cal_error_edge" calculates the leftvec (metric gradient) at the esae_error | esae_feature
        We can then have the contribution of feature_error_node by: 
            The contribution of feature_end_node is the gradient of f_esae_enc - sum grad_through_end_feat_error * feature_error_coef_to_cal_error_edge
        
        The "end_feature_dependent" sums all of the gradient of end_feature_node.
        '''
                
        if self.use_esae_error_term and self.use_error_term:
            all_end_node_grad: Tensor = leftvec.resc # type: ignore
            '''
            End feauture contribution
            '''
            feature_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_saes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_sae_end d_model -> b seq d_sae_end",
            )
            
            end_feature_dependent = einops.einsum(
                feature_coef_to_cal_error_edge, self.dict_saes[hook_position_end[0].name].W_enc,
                "b seq d_sae_end, d_model d_sae_end -> b seq d_model"
            )
            
            '''
            End feature error contribution
            '''
            grad_through_end_feat_error = einops.einsum( 
                self.dict_saes[hook_position_end[0].name].W_dec.T, self.dict_saes[hook_position_end[0].name].W_enc.T,
                "d_model1 d_sae_end, d_sae_end d_model2 -> d_model1 d_model2"
            )
            grad_through_end_feat_error = einops.einsum( # gradient through sae_feature
                self.dict_esaes[hook_position_end[0].name].W_enc.T, grad_through_end_feat_error,
                "d_esae_end d_model1, d_model1 d_model2 -> d_esae_end d_model2"
            )
            
            feature_error_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_esaes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_esae_end d_model -> b seq d_esae_end",
            )
            
            feature_error_dependent = einops.einsum(
                self.dict_esaes[hook_position_end[0].name].W_enc.T - grad_through_end_feat_error,
                feature_error_coef_to_cal_error_edge,
                "d_esae_end d_model, b seq d_esae_end -> b seq d_model",
            )
            
            for end_error_node, end_error_index in all_error:
                pos_end = end_error_index.idx[0]
                index = t.tensor(list(end_error_index.idx + (d_sae_end+d_esae_end,)), device=self.device)
                end_node_grad = leftvec.resc[:, pos_end].unsqueeze(1) # (b, 1, d_model) # type: ignore
                
                end_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    # (b, 1, d_model)
                    grad= end_node_grad - end_feature_dependent[:, pos_end].unsqueeze(1) - feature_error_dependent[:, pos_end].unsqueeze(1), 
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_error_grad_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T,
                    resc=end_error_grad_tensor,
                )

                effect = (
                    end_error_grad * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
                    
        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(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value 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()
    
    @contextmanager  
    def _setup_forward_model_hook(self, use_error_term: bool | None = None, transfer_grad: bool = True):
        cross_cache = {}
        # Hook function at transcoder input: caches the activations before transcoder
        def hook_crosscoder_input(activations: Tensor, hook: HookPoint, crosscoder_idx: int):
            cross_cache[crosscoder_idx] = activations.clone()
        
        all_cross_recons = [0.0 for _ in range(len(self.crosscoders))]
        def hook_crosscoder_output(activations: Tensor, hook: HookPoint, crosscoder_idx: int):
            recons, cross_recons = self.crosscoders[crosscoder_idx].forward_crosscoder(
                (cross_cache[crosscoder_idx], activations, all_cross_recons[crosscoder_idx])
            )
            for j in range(crosscoder_idx+1, self.n_layers):
                all_cross_recons[j] += cross_recons[j-crosscoder_idx-1]
            
            if transfer_grad:
                return recons + (activations - activations.detach())
            else:
                return recons
        
        fwd_hooks = []
        for i in range(len(self.crosscoders)):
            fwd_hooks.append((self.input_hooks[i], partial(hook_crosscoder_input, crosscoder_idx=i)))
            fwd_hooks.append((self.output_hooks[i], partial(hook_crosscoder_output, crosscoder_idx=i)))
            
        use_error_term = use_error_term if use_error_term is not None else self.use_error_term
        try: 
            for sae in self._saes_to_list():
                self.model.add_sae(sae, use_error_term)
            for hook, func in fwd_hooks:
                self.model.add_hook(hook, func, dir="fwd")
            
            yield
        finally:
            self.model.reset_saes()
            self.model.reset_hooks() 
                
    def _saes_to_list(self) -> List[Any]:
        return Feature_Graph_Cross._saes_to_list(self)