def dataset_to_dataframe(
    data: AnnotatedExtensionTypeWithShape,
    attr_name: str,
    encodings: dict[str, dict[str | int, int]],
    axis_types: Sequence[Axis] | None = None,
    index_types: Sequence[str] = ("state", "wholesaler", "brand", "product"),
    do_sort: bool = True,
) -> pd.DataFrame | pd.Series:
    """Create a dataframe for the specific attribute of the batched data.

    Args:
        data (EconomicModelInput): The batched data.
        attr_name (str): The attribute to make into a DataFrame.
        axis_types (Sequence[Axis]): The axis type labels for each axis of the index of the
                                      transposed DataFrame.
        encodings (dict[str, dict[str | int, int]]): The cached data and encodings about the full
                                                     dataset.

    Returns:
        A transformed and transposed DataFrame from the data with the columns equal to the batch
        dimension.
    """
    if axis_types is None:
        axis_types = data.datasets[attr_name].axes[1:]
    lookups_idx = [get_lookups(encodings, k) for k in index_types]
    index_tuples = get_tuples(data, index_types, lookups_idx, encodings)
    index = pd.MultiIndex.from_tuples(
        index_tuples,
        names=index_types,
    )
    if len(axis_types) > 0:
        lookups_col = [get_lookups(encodings, k) for k in axis_types]
        column_arrs = [
            [x[0] for x in get_tuples(data, (axis_type,), (lookups,), encodings)]
            for axis_type, lookups in zip(axis_types, lookups_col)
        ]
        columns = pd.MultiIndex.from_product(
            column_arrs,
            names=(
                ("time", "signal")
                if len(axis_types) == 2 and axis_types[0] == Axis.Time
                else (i.value for i in axis_types)
            ),
        )
        result = pd.DataFrame(
            get_data_attr(data, attr_name, encodings).reshape(len(index), -1), index=index, columns=columns
        ).T
        sort_axis = 1
    else:
        result = pd.Series(get_data_attr(data, attr_name, encodings).reshape(-1), index=index, name=attr_name)
        sort_axis = 0
    if do_sort:
        return result.sort_index(axis=sort_axis)
    else:
        return result

def denormalize_z_score(
    batch_df: pd.DataFrame,
    signal: str,
    encodings: dict[str, Any],
    normalized_on: str | None = None,
    mean_rescaled: bool = False,
) -> pd.DataFrame:
    """Denormalize z-score transformed batch_df.

    Args:
        batch_df (pd.DataFrame): The dataframe that needs denomralized.
        signal (str): Signal that is being transformed or name of df.
        encodings (dict[str, Any]): Encodings to get std and means dict for reverse transformation.
        normalized_on (str | None, optional): On what level, the df is normalized on.
            The given level wil be mapped by f"{signal}_means_row_lookup". Defaults to None.
        mean_rescaled (bool, optional): Specifies if the data is rescaled by the mean (transformed_mean=1)
            or centered (transformed_mean=0). Defaults to False.

    Returns:
        pd.DataFrame: DataFrame after applying reverse z score transformation.
    """
    if f"{signal}_stds_dict" in encodings:
        # applied on signal level
        stds = encodings[f"{signal}_stds_dict"]
        batch_df = batch_df * batch_df.columns.get_level_values("signal").map(defaultdict(lambda: 1, stds))

    if f"{signal}_means_row_lookup" in encodings:
        # it's mean normalized on a specific granularity
        if normalized_on not in batch_df.columns.names or normalized_on == "signal":
            raise ValueError(f"`normalized_on` must be one of these {set(batch_df.columns.names)-{'signal',}}")

        # mean applied on each col and then row level
        means_row_lookup = get_lookups(encodings[f"{signal}_means_row_lookup"])
        means_col_lookup = get_lookups(encodings[f"{signal}_means_col_lookup"])
        normalized_on_means_dict = {
            k: {name: vi for name, vi in zip(means_col_lookup, v)}
            for k, v in zip(means_row_lookup, encodings[f"{signal}_means"])
        }
        signal_level = batch_df.columns.get_level_values("signal")
        normalized_on_level = batch_df.columns.get_level_values(normalized_on)
        means_dict_index = normalized_on_level.map(normalized_on_means_dict)
        for i, (signal, means_dict) in enumerate(zip(signal_level, means_dict_index)):
            if signal in means_dict:
                if mean_rescaled:
                    batch_df.iloc[:, i] *= means_dict[signal] + EPSILON
                else:
                    batch_df.iloc[:, i] += means_dict[signal]
        batch_df = batch_df.round(3)
    elif signal == "regional_trend":
        brand_lookup = encodings["brand_lookup"]
        region_lookup = encodings["region_lookup"]
        brand_region_index = [(r, b, "regional_trend") for r in region_lookup for b in brand_lookup]
        means_series = pd.DataFrame(
            encodings["regional_trend_means_dict"]["regional_trend"],
            index=pd.MultiIndex.from_tuples(brand_region_index, names=["region", "brand", "signal"]),
        )
        wholesaler_region = pd.DataFrame.from_dict(encodings["wholesaler_region_lookup"], orient="index").reset_index(
            names=["wholesaler"]
        )
        wholesaler_region = wholesaler_region.rename(columns={0: "region"})
        original_order = list(batch_df.columns.names)
        means_series = (
            means_series.reset_index()
            .merge(wholesaler_region, on="region")
            .drop("region", axis=1)
            .set_index(["wholesaler", "brand", "signal"])
        )
        batch_df, means_series = batch_df.align(means_series.T, join="left", axis=1)
        if mean_rescaled:
            batch_df = batch_df * means_series.values
        else:
            batch_df = batch_df + means_series.values
        batch_df = batch_df.reorder_levels(original_order, axis=1)
    else:
        # it's simple mean normalized df
        means_series = pd.DataFrame(encodings[f"{signal}_means_dict"]).stack()
        means_series.index.names = ["brand", "signal"]
        original_order = list(batch_df.columns.names)
        batch_df, means_series = batch_df.align(means_series, join="left", axis=1)
        if mean_rescaled:
            batch_df = batch_df * means_series.values[None, :]
        else:
            batch_df = batch_df + means_series.values[None, :]
        batch_df = batch_df.reorder_levels(original_order, axis=1)
    return batch_df

def get_data_attr(data: EconomicModelInput, attr_name: str | Axis | list[str], encodings: dict[str, Any]):
    """
    if data has attribute attr_name, return the numpy (often temporal) object.
    If attr_name is the name of an index, not of a signal (like wholesaler, product, state, etc),
    then return data.{attr_name}_index. For instance, for wholesaler, return data.wholesaler_index.
    """
    if attr_name == "time" or attr_name == "time_index":
        date_lookups = get_lookups(encodings["date"])
        return pd.DatetimeIndex([date_lookups[i] for i in to_numpy(data.date_index)], name="date")
    elif isinstance(attr_name, str):
        attr_name = str(attr_name)
        if hasattr(data, attr_name):
            return to_numpy(getattr(data, attr_name))
        elif hasattr(data, f"{attr_name}_index"):
            result = to_numpy(getattr(data, f"{attr_name}_index"))
            if len(result.shape) > 1:
                result = result[0]
            return result
    else:
        return attr_name