from typing import Any, Callable

import equinox as eqx
from jaxtyping import PyTree


class AuxData:
    """A simple container for auxiliary data"""

    def __init__(self):
        self.data = None

    def update(self, x: Any):
        """**Arguments:**

        - `x`: Any output generated by an intermediate layer that is meant to be stored
        """
        self.data = x


def _make_intermediate_layer_wrapper():
    aux = AuxData()

    class IntermediateWrapper(eqx.Module):
        layer: eqx.Module

        def __call__(self, x, *, key=None):
            out = self.layer(x, key=key)
            aux.update(out)
            return out

    return aux, IntermediateWrapper


def intermediate_layer_getter(
    model: PyTree, get_target_layers: Callable
) -> "eqx.Module":
    """Wraps intermediate layers of a model for accessing intermediate activations. Based on a discussion
    [here](https://github.com/patrick-kidger/equinox/issues/186).

    !!! info
        Only supports storing the result of the most recent call. So, if the forward utilises the same layer multiple
        times, the returned intermediate value will be of the last call

    **Arguments:**

    - `model`: A PyTree representing the neural network model
    - `get_target_layers`: A callable function which returns a sequence
        of layers from the `model`

    **Returns:**
    A `PyTree`, encapsulating `model` for storing intermediate outputs from target layers.

    !!! info
        The returned model will now return a `tuple` with

            1. The final output of `model`
            2. An ordered list of intermediate activations
    """
    target_layers = get_target_layers(model)
    auxs, wrappers = zip(
        *[_make_intermediate_layer_wrapper() for _ in range(len(target_layers))]
    )
    model = eqx.tree_at(
        where=get_target_layers,
        pytree=model,
        replace=[
            wrapper(target_layer)
            for (wrapper, target_layer) in zip(wrappers, target_layers)
        ],
    )

    class IntermediateLayerGetter(eqx.Module):
        model: eqx.Module

        def __call__(self, x, *, key=None):
            out = self.model(x, key=key)
            return out, [aux.data for aux in auxs]

    return IntermediateLayerGetter(model)
