from unittest.mock import patch
from contextlib import contextmanager, ExitStack

# 1. A helper to generate the wrapper closure
# We pass 'original_func' here so it is "baked in" to the wrapper.
def create_avg_wrapper(original_func):
    def wrapper(inputs, *args, **kwargs):
        # Your custom logic
        if len(inputs.shape) == 3:
            inputs = inputs.mean(dim=1, keepdim=True)
        elif len(inputs.shape) == 2:
            inputs = inputs.unsqueeze(1)
        else:
            raise
        
        # Call the CAPTURED original function
        return original_func(inputs, *args, **kwargs)
    return wrapper

@contextmanager
def vision_sae_wrapper(saes):
    # Ensure we can handle a single SAE or a list of them
    if not isinstance(saes, list):
        saes = [saes]

    # ExitStack is designed exactly for "I need N context managers"
    with ExitStack() as stack:
        for sae in saes:
            # 1. Capture the specific original method for THIS sae
            original_forward = sae.forward
            
            # 2. Create a wrapper specific to this sae (closing over original_forward)
            patched_forward = create_avg_wrapper(original_forward)
            
            # 3. Enter the patch context and add it to the stack
            stack.enter_context(patch.object(sae, 'forward', side_effect=patched_forward))
        
        # Yield control back to your code
        yield

# --- Usage Example ---
# sae_list = [sae1, sae2, sae3]
# with vision_sae_wrapper(sae_list):
#     # All SAEs are now patched
#     output1 = sae1(input_tensor)
#     output2 = sae2(input_tensor)