@in_cpu
def as_dict(output: "ModelOutputType" | dict | Mapping | tf.experimental.ExtensionType) -> dict[str, ...]:
    """
    Converts an output with ExtensionType or dataclass type to dict.
    Annotations are preserved in the dict.
    Use RecursiveNamespace to load it back with __annotations__.
    """
    if isinstance(output, (dict, Mapping)):
        return {k: as_dict(v) for k, v in output.items()}
    elif isinstance(output, tf.experimental.ExtensionType):
        return {k: as_dict(v) for k, v in tf.experimental.extension_type.as_dict(output).items()} | {
            "__annotations__": getattr(type(output), "annotations", getattr(type(output), "__annotations__", {}))
        }
    elif is_dataclass(output):
        return {k: as_dict(v) for k, v in dataclass_asdict(output).items()} | {
            "__annotations__": getattr(type(output), "__annotations__", {})
        }
    elif isinstance(output, SimpleNamespace):
        return {k: as_dict(v) for k, v in vars(output).items()} | {
            "__annotations__": getattr(type(output), "__annotations__", {})
        }
    else:
        return output

def calculate_curves_diff(
    prev_output: "EconomicIntermediariesWithTrackers",
    current_output: "EconomicIntermediariesWithTrackers",
    encodings: Encodings,
) -> dict[str, NDArray]:
    """Takes 2 model outputs with trackers and then calculates _calculate_curves_diff."""
    if not (hasattr(prev_output, "trackers") and hasattr(current_output, "trackers")):
        raise ValueError("`trackers` not found! Please run WFO with `calculate_trackers=True`.")

    _, prev_roicurve = get_curve_matrix(
        to_numpy(prev_output.trackers.curve_tracker.spends),
        to_numpy(prev_output.trackers.curve_tracker.impact),
        to_numpy(prev_output.trackers.curve_tracker.slope),
        encodings["normalization_factor"],
        dynamic_range=True,
    )
    _, prev_global_me = get_mixed_effect_matrix(prev_output.trackers.global_me_tracker)
    _, prev_pricing_me = get_mixed_effect_matrix(prev_output.trackers.pricing_me_tracker)

    curr_roi_idx, curr_roicurve = get_curve_matrix(
        to_numpy(current_output.trackers.curve_tracker.spends),
        to_numpy(current_output.trackers.curve_tracker.impact),
        to_numpy(current_output.trackers.curve_tracker.slope),
        encodings["normalization_factor"],
        dynamic_range=True,
    )
    curr_global_idx, curr_global_me = get_mixed_effect_matrix(current_output.trackers.global_me_tracker)
    curr_pricing_idx, curr_pricing_me = get_mixed_effect_matrix(current_output.trackers.pricing_me_tracker)

    metric_combinations = {
        "roi_diffs": (prev_roicurve, curr_roicurve, curr_roi_idx),
        "global_me_impacts_diffs": (prev_global_me, curr_global_me, curr_global_idx),
        "pricing_me_impacts_diffs": (prev_pricing_me, curr_pricing_me, curr_pricing_idx),
    }

    return {
        metric: calculate_curves_diff_utils(prev[:, :, 1, :], curr[:, :, 1, :], idx[:, None, None])
        for metric, (prev, curr, idx) in metric_combinations.items()
    }

def calculate_curves_diff_utils(y_a: TensorLike, y_b: TensorLike, x_vals: TensorLike, axis: int = 0) -> NDArray:
    """Approximate $\\int_{\\min{spend}}^{\\max{spend}} \\left(roi_i'(x) - roi_{i+1}'(x)\\right)^2$
       to get difference of rois or any two curves in general. Units of the result are the same as
       the units of the y values passed in.

    Args:
        y_a (np.ndarray): Curve 1 y values at each x. Shape: num_points, batch, vehicles.
        y_b (np.ndarray): Curve 2 y values at each x. Shape: num_points, batch, vehicles.
        x_vals: The common x axis values of the two curves. Shape: spend_points,.
        axis (int, optional): The axis of integration, defaults to the 1st axis (spend_points).
    Returns:
        np.ndarray: Computed integral of squared differences between two curves.
    """
    # spend_points, batch, vehicles
    squared_diff = np.square(y_a - y_b)
    # batch, vehicles
    areas = np.trapz(squared_diff, x_vals, axis=axis) / (np.max(x_vals, axis=axis) - np.min(x_vals, axis=axis))
    return np.sqrt(areas)

@in_cpu
@warn_once(ConcatWarning)
def concat_intermediaries(
    intermediaries: list[T],
    axis: int = 0,
    axis_type: Axis = Axis.Batch,
) -> T:
    """Concatenates tf.experimental.ExtensionType objects in the batch_axis.
    Make sure that both `axis` and `axis_type` matches else it leads to unexpected results.

    Args:
        intermediaries (list[tf.experimental.ExtensionType]): tf.experimental.ExtensionType type instance.
        axis (int, optional): The axis where we need to concatenate. Defaults to 0 which is assumed to be batch.
        axis_type(Axis, optional): The annotated axis where we need to concatenate. Defaults to Axis.Batch.

    Returns:
        tf.experimental.ExtensionType: The concatenated output which is the same type that is passed in.
    """
    if not hasattr(intermediaries, "__len__") or len(intermediaries) < 1:
        # when we don't pass a list. Only avoid an error.
        return intermediaries
    elif len(intermediaries) == 1:
        # Handling edge case where its full batch.
        return intermediaries[0]

    OutputType = type(intermediaries[0])
    annotations: dict[str, ...] = getattr(OutputType, "annotations", getattr(OutputType, "__annotations__", {}))
    concated_intermediaries = {}
    # ExtensionType supports Mapping, so this detects that for us. else we assume its an ExtensionType
    dict_type = issubclass(OutputType, (dict, Mapping))
    attributes = intermediaries[0].keys() if dict_type else vars(intermediaries[0])
    for attr in attributes:
        annotation = annotations.get(attr, None)
        batch_values = [batch.get(attr) if dict_type else getattr(batch, attr) for batch in intermediaries]
        value = batch_values[0]
        if value is None:
            concatenated_values = None
        elif isinstance(value, (tf.experimental.ExtensionType, dict, Mapping, SimpleNamespace)):
            concatenated_values = concat_intermediaries(batch_values, axis=axis, axis_type=axis_type)
        else:
            concatenated_values = _concatenate_intermediaries_leaf(
                annotation, batch_values, value, axis, axis_type, attr
            )
        concated_intermediaries[attr] = concatenated_values

    return OutputType(**concated_intermediaries)

@in_cpu
def gather(
    output: T,
    indices: NDArray[np.int_],
    axis: int = 0,
    axis_type: Axis = Axis.Batch,
) -> T:
    """
    Gathers slices from different params in `output` according to `indices`.
    NOTE: this can lead to unexpected behaviour when tensors are not annotated as the `axis` might be available but
    it's not meant to be gathered.

    Args:
        output (list[tf.experimental.ExtensionType]): tf.experimental.ExtensionType type instance.
        indices (NDArray[np.int32]): The index Tensor. Must be one of the following types: int32, int64.
            The values must be in range [0, params.shape[axis]).
        axis (int, optional): The axis where we need to concatenate. Defaults to 0 which is assumed to be batch.
        axis_type(Axis, optional): The annotated axis where we need to concatenate. Defaults to Axis.Batch.

    Returns:
        tf.experimental.ExtensionType: The gathered output which is the same type that is passed in.
    """
    OutputType = type(output)
    # an AnnotatedExtensionType will have `annotations` to handle recursion
    annotations: dict[str, Any] = getattr(OutputType, "annotations", getattr(OutputType, "__annotations__", {}))
    gathered_output = {}
    # ExtensionType supports Mapping, so this detects that for us. else we assume its an ExtensionType
    dict_type = issubclass(OutputType, (dict, Mapping))

    if dict_type:
        output_dict = output
    elif isinstance(output, tf.experimental.ExtensionType):
        output_dict = tf.experimental.extension_type.as_dict(output)
    else:
        output_dict = dict(vars(output))

    for attr, values in output_dict.items():
        annotation = annotations.get(attr, None)
        gathered_values = _get_gathered_values(annotation, values, indices, axis, axis_type)

        gathered_output[attr] = gathered_values

    return OutputType(**gathered_output)

def get_index(
    level: tuple[str, ...] | str | None,
    inputs: EconomicModelInput,
    encodings: Encodings,
) -> pd.Index | pd.MultiIndex:
    """For given `level`, pd.Index | pd.MultiIndex is created.

    Args:
        level (tuple[str, ...] | str | None): The level(s) on which we need to create the index from.
        inputs (EconomicModelInput): Inputs intermediearies that contains the level(s) index values.
        encodings (dict[str, dict[str, int]]): Encodings dict which will be used to decode the index values.

    Returns:
        pd.Index | pd.MultiIndex: The decoded index created based on level(s).
    """
    if level is None:
        level = ("brand", "wholesaler")

    if isinstance(level, str) and (level == "country"):
        # NOTE: add 'country' in encodings so that this can be generalized as well.
        group_index = pd.Index(["US" for _ in range(len(inputs.brand_index))], name="country")
    elif isinstance(level, str):
        group_index = pd.Index(getattr(inputs, f"{level}_index"), name=level)
    elif isinstance(level, (tuple, list)):
        group_index = pd.MultiIndex.from_arrays(
            [np.array(getattr(inputs, f"{lvl}_index")) for lvl in level], names=level
        )
    else:
        raise ValueError(f"Invalid level {level}")

    if level != "country":
        # NOTE: add 'country' in encodings so that this can be generalized as well.
        group_index = _index_mapper(group_index, encodings)
    return group_index

def get_indexer(index: pd.MultiIndex | pd.Index, target: Sequence[tuple | Hashable]) -> NDArray[np.intp]:
    """
    Compute indexer and mask for new index given the current index.
    The indexer should be then used as an input to ndarray.take to align the current data to the new index.
    This is based on pd.Index.get_indexer. Pandas implementation doesn't support if the `index` is not unique.

    Args:
        index (pd.MultiIndex | pd.Index): Current index.
        target (Sequence[tuple | Hashable]): New target index, we like to compute indexes on current index.

    Returns:
        NDArray[np.intp]: Integers from 0 to n - 1 indicating that the index at these positions matches the
            corresponding target values. Missing values in the target are marked by -1.
    """
    index_mapping = defaultdict(lambda: -1, zip(target, range(len(target))))
    return index.map(index_mapping).values.astype(np.intp)

def process_period_curves_diff(differences: list[dict[str, NDArray]]) -> dict[str, NDArray]:
    """Processes the differences calculated for each period and curves."""
    metrics = defaultdict(list)
    for period_curves in differences:
        for curve_name, value in period_curves.items():
            metrics[curve_name].append(value)
    return {curve_name: _process_period_curve(period_curves) for curve_name, period_curves in metrics.items()}

@in_cpu
def sort_batch_index(
    output: EconomicIntermediaries | dict[str, EconomicIntermediaries],
    encodings: Encodings,
    level: tuple[str, ...] = ("brand", "wholesaler"),
    axis: int = 0,
    axis_type: Axis = Axis.Batch,
) -> EconomicIntermediaries | dict[str, EconomicIntermediaries]:
    """
    Sorts the given intermediaries on the batch axis, based on level indexes.
    The sort order is based on the level(s) provided. This is useful when we need to have the outputs in the same order.

    Args:
        output (EconomicIntermediaries | dict[str, EconomicIntermediaries]): Outputs of the model.
        encodings (dict[str, dict[str, int]]): Encodings dict which will be used to decode the index values.
        level (tuple[str, ...], optional): The level(s) on which the batches is ordered.
            Defaults to ("brand", "wholesaler").
        axis (int, optional): The batch axis. Defaults to 0 which is assumed to be batch.
        axis_type(Axis, optional): The annotated batch axis. Defaults to Axis.Batch.

    Returns:
        EconomicIntermediaries | dict[str, EconomicIntermediaries]: The sorted intermediares based on level(s).
    """
    if isinstance(output, dict):
        return {
            key: sort_batch_index(
                output[key],
                encodings=encodings,
                level=level,
                axis=axis,
                axis_type=axis_type,
            )
            for key in output.keys()
        }

    if TYPE_CHECKING:
        assert isinstance(output, EconomicIntermediaries)
        assert output.inputs is not None
    index = get_index(level=level, inputs=output.inputs, encodings=encodings)
    sorted_idx = tf.convert_to_tensor(index.argsort(), dtype=tf.int32)

    return gather(output, sorted_idx, axis=axis, axis_type=axis_type)