src/openpi_value/transforms.py

from collections.abc import Callable, Mapping, Sequence
import dataclasses
import re
from typing import List, Protocol, TypeAlias, TypeVar, runtime_checkable

import flax.traverse_util as traverse_util
import jax
import numpy as np
from openpi_client import image_tools
import torch

from openpi_value.models import tokenizer as _tokenizer
from openpi_value.shared import array_typing as at
from openpi_value.shared import normalize as _normalize

import numpy as np
from PIL import Image
from typing import Union
import copy
import json

DataDict: TypeAlias = at.PyTree
NormStats: TypeAlias = _normalize.NormStats


T = TypeVar("T")
S = TypeVar("S")



def _resize_pil(image: Image.Image, size: Union[int, tuple[int, int]], method: int) -> Image.Image:
    """Resizes a single image using PIL without padding.

    Args:
        image: The PIL Image to resize.
        size: The target size.
            - If int: Resizes the shortest edge to this size, maintaining aspect ratio.
            - If (height, width): Resizes directly to this size, potentially distorting.
        method: The interpolation method.

    Returns:
        The resized PIL Image.
    """
    if isinstance(size, int):
        # Case 1: Resize shortest edge to `size`
        cur_width, cur_height = image.size

        # Optimization: If shortest edge is already `size`, no resize needed.
        if min(cur_width, cur_height) == size:
            return image

        # Calculate new dimensions
        short_edge = min(cur_width, cur_height)
        ratio = size / short_edge
        new_width = int(round(cur_width * ratio))
        new_height = int(round(cur_height * ratio))

        # PIL resize expects (width, height)
        return image.resize((new_width, new_height), resample=method)

    elif isinstance(size, (tuple, list)) and len(size) == 2:
        # Case 2: Resize directly to (height, width)
        target_height, target_width = size
        cur_width, cur_height = image.size

        # Optimization: If size is already correct, no resize needed.
        if (cur_width, cur_height) == (target_width, target_height):
            return image

        # PIL resize expects (width, height)
        return image.resize((target_width, target_height), resample=method)

    else:
        raise ValueError(f"Unsupported size format: {size}. Must be int or (height, width).")


def resize_no_pad(images: np.ndarray, size: Union[int, tuple[int, int]], method=Image.BILINEAR) -> np.ndarray:
    """Resizes a batch of images without padding, following tf.image.resize logic.

    Args:
        images: A batch of images in [..., height, width, channel] format.
        size: The target size.
            - If int: Resizes the shortest edge to this size, maintaining aspect ratio.
            - If (height, width): Resizes directly to this size, potentially distorting.
        method: The interpolation method to use. Default is bilinear.

    Returns:
        The resized images.
        - If size=(h, w), shape is [..., h, w, channel].
        - If size=int, shape is [..., new_h, new_w, channel]. Note: This
          will raise an error if images in the batch have different aspect
          ratios, as they cannot be stacked into a single tensor.
    """
    original_shape = images.shape
    
    # --- Optimization: Check if resize is a no-op ---
    if isinstance(size, int):
        # Check shortest edge
        if min(original_shape[-3], original_shape[-2]) == size:
            return images
    elif isinstance(size, (tuple, list)) and len(size) == 2:
        # Check exact (h, w)
        if original_shape[-3:-1] == size:
            return images
    # -------------------------------------------------

    # Reshape to a flat batch of images
    images = images.reshape(-1, *original_shape[-3:])

    # Handle empty batch case
    if images.shape[0] == 0:
        if isinstance(size, (tuple, list)):
            h, w = size
            return np.empty((*original_shape[:-3], h, w, original_shape[-1]), dtype=images.dtype)
        else:
            # Cannot determine output shape for 'int' size on empty batch,
            # so return the original empty batch.
            return images.reshape(original_shape)

    # Apply the helper function to each image in the batch
    resized_list = [_resize_pil(Image.fromarray(im), size, method=method) for im in images]

    # Stack them back into a single tensor
    # This will automatically convert PIL Images to np.ndarray
    resized = np.stack(resized_list)

    # Reshape back to original batch dimensions
    new_shape = resized.shape[-3:]  # (new_h, new_w, c)
    return resized.reshape(*original_shape[:-3], *new_shape)



@runtime_checkable
class DataTransformFn(Protocol):
    def __call__(self, data: DataDict) -> DataDict:
        """Apply transformation to the data.

        Args:
            data: The data to apply the transform to. This is a possibly nested dictionary that contains
                unbatched data elements. Each leaf is expected to be a numpy array. Using JAX arrays is allowed
                but not recommended since it may result in extra GPU memory usage inside data loader worker
                processes.

        Returns:
            The transformed data. Could be the input `data` that was modified in place, or a new data structure.
        """


@dataclasses.dataclass(frozen=True)
class Group:
    """A group of transforms."""

    # Transforms that are applied to the model input data.
    inputs: Sequence[DataTransformFn] = ()

    # Transforms that are applied to the model output data.
    outputs: Sequence[DataTransformFn] = ()

    def push(self, *, inputs: Sequence[DataTransformFn] = (), outputs: Sequence[DataTransformFn] = ()) -> "Group":
        """Append transforms to the group and return a new group.

        Args:
            inputs: Appended to the *end* of the current input transforms.
            outputs: Appended to the *beginning* of the current output transforms.

        Returns:
            A new group with the appended transforms.
        """
        return Group(inputs=(*self.inputs, *inputs), outputs=(*outputs, *self.outputs))


@dataclasses.dataclass(frozen=True)
class CompositeTransform(DataTransformFn):
    """A composite transform that applies a sequence of transforms in order."""

    transforms: Sequence[DataTransformFn]

    def __call__(self, data: DataDict) -> DataDict:
        for transform in self.transforms:
            data = transform(data)
        return data


def compose(transforms: Sequence[DataTransformFn]) -> DataTransformFn:
    """Compose a sequence of transforms into a single transform."""
    return CompositeTransform(transforms)


@dataclasses.dataclass(frozen=True)
class RepackTransform(DataTransformFn):
    """Repacks an input dictionary into a new dictionary.

    Repacking is defined using a dictionary where the keys are the new keys and the values
    are the flattened paths to the old keys. We use '/' as the separator during flattening.

    Example:
    {
        "images": {
            "cam_high": "observation.images.top",
            "cam_low": "observation.images.bottom",
        },
        "state": "observation.state",
        "actions": "action",
    }
    """

    structure: at.PyTree[str]

    def __call__(self, data: DataDict) -> DataDict:
        flat_item = flatten_dict(data)




        return jax.tree.map(lambda k: flat_item[k], self.structure)


@dataclasses.dataclass(frozen=True)
class InjectDefaultPrompt(DataTransformFn):
    prompt: str | None

    def __call__(self, data: DataDict) -> DataDict:

        if self.prompt is not None:
            data["prompt"] = np.asarray(self.prompt) # ! use default prompt, if exists
        
        return data


@dataclasses.dataclass(frozen=True)
class Normalize(DataTransformFn):
    norm_stats: at.PyTree[NormStats] | None
    # If true, will use quantile normalization. Otherwise, normal z-score normalization will be used.
    use_quantiles: bool = False
    # If true, will raise an error if any of the keys in the norm stats are not present in the data.
    strict: bool = False

    def __post_init__(self):
        if self.norm_stats is not None and self.use_quantiles:
            _assert_quantile_stats(self.norm_stats)

    def __call__(self, data: DataDict) -> DataDict:
        if self.norm_stats is None:
            return data

        out = apply_tree(
            data,
            self.norm_stats,
            self._normalize_quantile if self.use_quantiles else self._normalize,
            strict=self.strict,
        )

        return out

    @staticmethod
    def _cast_like(value, reference):
        """Keep normalization in the dtype supplied by the dataset/model."""
        if hasattr(value, "to") and hasattr(reference, "dtype"):
            return value.to(dtype=reference.dtype)
        if hasattr(value, "astype") and hasattr(reference, "dtype"):
            return value.astype(reference.dtype)
        return value

    def _normalize(self, x, stats: NormStats):
        mean, std = stats.mean[..., : x.shape[-1]], stats.std[..., : x.shape[-1]]
        return self._cast_like((x - mean) / (std + 1e-6), x)

    def _normalize_quantile(self, x, stats: NormStats):
        assert stats.q01 is not None
        assert stats.q99 is not None
        q01, q99 = stats.q01[..., : x.shape[-1]], stats.q99[..., : x.shape[-1]]
        return self._cast_like((x - q01) / (q99 - q01 + 1e-6) * 2.0 - 1.0, x)


@dataclasses.dataclass(frozen=True)
class Unnormalize(DataTransformFn):
    norm_stats: at.PyTree[NormStats] | None
    # If true, will use quantile normalization. Otherwise, normal z-score normalization will be used.
    use_quantiles: bool = False

    def __post_init__(self):
        if self.norm_stats is not None and self.use_quantiles:
            _assert_quantile_stats(self.norm_stats)

    def __call__(self, data: DataDict) -> DataDict:
        if self.norm_stats is None:
            return data

        # Make sure that all the keys in the norm stats are present in the data.
        return apply_tree(
            data,
            self.norm_stats,
            self._unnormalize_quantile if self.use_quantiles else self._unnormalize,
            strict=True,
        )

    @staticmethod
    def _cast_like(value, reference):
        if hasattr(value, "to") and hasattr(reference, "dtype"):
            return value.to(dtype=reference.dtype)
        if hasattr(value, "astype") and hasattr(reference, "dtype"):
            return value.astype(reference.dtype)
        return value

    def _unnormalize(self, x, stats: NormStats):
        mean = pad_to_dim(stats.mean, x.shape[-1], axis=-1, value=0.0)
        std = pad_to_dim(stats.std, x.shape[-1], axis=-1, value=1.0)
        return self._cast_like(x * (std + 1e-6) + mean, x)

    def _unnormalize_quantile(self, x, stats: NormStats):
        assert stats.q01 is not None
        assert stats.q99 is not None
        q01, q99 = stats.q01, stats.q99
        if (dim := q01.shape[-1]) < x.shape[-1]:
            return self._cast_like(
                np.concatenate([(x[..., :dim] + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01, x[..., dim:]], axis=-1), x
            )
        return self._cast_like((x + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01, x)


@dataclasses.dataclass(frozen=True)
class ResizeImages(DataTransformFn):
    height: int
    width: int

    def __call__(self, data: DataDict) -> DataDict:

        data["image_original"] = {k: v.copy() for k, v in data["image"].items()}
        data["image_original"] = {k: resize_no_pad(v, size=self.height) for k, v in data["image_original"].items()}  # * resize without padding, resize the shorter side to height

        data["image"] = {k: image_tools.resize_with_pad(v, self.height, self.width) for k, v in data["image"].items()}
        return data


@dataclasses.dataclass(frozen=True)
class SubsampleActions(DataTransformFn):
    stride: int

    def __call__(self, data: DataDict) -> DataDict:
        data["actions"] = data["actions"][:: self.stride]
        return data


@dataclasses.dataclass(frozen=True)
class DeltaActions(DataTransformFn):
    """Repacks absolute actions into delta action space."""

    # Boolean mask for the action dimensions to be repacked into delta action space. Length
    # can be smaller than the actual number of dimensions. If None, this transform is a no-op.
    # See `make_bool_mask` for more details.
    mask: Sequence[bool] | None

    def __call__(self, data: DataDict) -> DataDict:
        if "actions" not in data or self.mask is None:
            return data

        state, actions = data["state"], data["actions"]
        mask = np.asarray(self.mask)
        dims = mask.shape[-1]
        actions[..., :dims] -= np.expand_dims(np.where(mask, state[..., :dims], 0), axis=-2)
        data["actions"] = actions

        return data


@dataclasses.dataclass(frozen=True)
class AbsoluteActions(DataTransformFn):
    """Repacks delta actions into absolute action space."""

    # Boolean mask for the action dimensions to be repacked into absolute action space. Length
    # can be smaller than the actual number of dimensions. If None, this transform is a no-op.
    # See `make_bool_mask` for more details.
    mask: Sequence[bool] | None

    def __call__(self, data: DataDict) -> DataDict:
        if "actions" not in data or self.mask is None:
            return data

        state, actions = data["state"], data["actions"]
        mask = np.asarray(self.mask)
        dims = mask.shape[-1]
        actions[..., :dims] += np.expand_dims(np.where(mask, state[..., :dims], 0), axis=-2)
        data["actions"] = actions

        return data


@dataclasses.dataclass(frozen=True)
class TokenizePrompt(DataTransformFn):
    tokenizer: _tokenizer.PaligemmaTokenizer
    discrete_state_input: bool = False
    # advantage_bins: int = 10
    advantage_bins: Union[int, str, List[float], List[int]] = 10

    def __call__(self, data: DataDict) -> DataDict:
        if (prompt := data.pop("prompt", None)) is None:
            raise ValueError("Prompt is required")

        if self.discrete_state_input:
            if (state := data.get("state", None)) is None:
                raise ValueError("State is required.")
        else:
            state = None

        if not isinstance(prompt, str):
            prompt = prompt.item()

        
        action_advantage = data.get("action_advantage", None)
        
        action_advantage_original = copy.deepcopy(action_advantage)

        if action_advantage is not None:

            if len(action_advantage.shape) == 0:  # * True, get in

                if type(action_advantage) is torch.Tensor:
                    action_advantage = action_advantage.cpu().numpy()

                # * discretize advantages are in ints of [1, self.advantage_bins], 1, 2, 3, xxx, self.advantage_bins
                

                if isinstance(self.advantage_bins, list):
                    bins = np.array(self.advantage_bins)  # * manual advantage.
                elif isinstance(self.advantage_bins, int):
                    bins = np.linspace(-1, 1, self.advantage_bins + 1)[:-1]  # * add 0.1 offset to set 0 to low bins.
                else:
                    raise NotImplementedError(f"advantage_bins type {type(self.advantage_bins)} not supported.")
                
                action_advantage = np.digitize(action_advantage, bins=bins)  # * add 0.1 offset to set 0 to low bins.
                
                
        tokens, token_masks = self.tokenizer.tokenize(prompt, state, action_advantage=action_advantage)
        
        return {**data, "tokenized_prompt": tokens, "tokenized_prompt_mask": token_masks,

                # * Custom
                "action_advantage": action_advantage,
                "action_advantage_original": action_advantage_original,
                }




@dataclasses.dataclass(frozen=True)
class PromptFromLeRobotTask(DataTransformFn):
    """Extracts a prompt from the current LeRobot dataset task."""

    # Contains the LeRobot dataset tasks (dataset.meta.tasks).
    tasks: dict[int, str]

    def __call__(self, data: DataDict) -> DataDict:

        
        
        if self.tasks is not None:
            if "task_index" not in data:
                raise ValueError('Cannot extract prompt without "task_index"')
            
            task_index = int(data["task_index"])
            if (prompt := self.tasks.get(task_index)) is None:
                raise ValueError(f"{task_index=} not found in task mapping: {self.tasks}")
        else:
            assert "task" in data, 'Cannot extract prompt without "task"'
            prompt = data["task"]

            # * For gelexea dataset, only keep the part after '@'

        
        if '@' in prompt:
            prompt = prompt.split('@')[-1].strip()
        return {**data, "prompt": prompt}


@dataclasses.dataclass(frozen=True)
class PadStatesAndActions(DataTransformFn):
    """Zero-pads states and actions to the model action dimension."""

    model_action_dim: int

    def __call__(self, data: DataDict) -> DataDict:
        data["state"] = pad_to_dim(data["state"], self.model_action_dim, axis=-1)
        if "actions" in data:
            data["actions"] = pad_to_dim(data["actions"], self.model_action_dim, axis=-1)
        return data


def flatten_dict(tree: at.PyTree) -> dict:
    """Flatten a nested dictionary. Uses '/' as the separator."""
    return traverse_util.flatten_dict(tree, sep="/")


def unflatten_dict(tree: dict) -> at.PyTree:
    """Unflatten a flattened dictionary. Assumes that '/' was used as a separator."""
    return traverse_util.unflatten_dict(tree, sep="/")


def transform_dict(patterns: Mapping[str, str | None], tree: at.PyTree) -> at.PyTree:
    """Transform the structure of a nested dictionary using a set of patterns.

    The transformation is defined using the `patterns` dictionary. The keys are the
    input keys that should be matched and the values are the new names inside the output
    dictionary. If the value is None, the input key is removed.

    Both keys and values should represent flattened paths using '/' as the separator.
    Keys can be regular expressions and values can include backreferences to the
    matched groups (see `re.sub` for more details). Note that the regular expression
    must match the entire key.

    The order inside the `patterns` dictionary is important. Only the first pattern that
    matches the input key will be used.

    See unit tests for more examples.

    Args:
        patterns: A mapping from old keys to new keys.
        tree: The nested dictionary to transform.

    Returns:
        The transformed nested dictionary.
    """
    data = flatten_dict(tree)

    # Compile the patterns.
    compiled = {re.compile(k): v for k, v in patterns.items()}

    output = {}
    for k in data:
        for pattern, repl in compiled.items():
            if pattern.fullmatch(k):
                new_k = pattern.sub(repl, k, count=1) if repl is not None else None
                break
        else:
            # Use the original key if no match is found.
            new_k = k

        if new_k is not None:
            if new_k in output:
                raise ValueError(f"Key '{new_k}' already exists in output")
            output[new_k] = data[k]

    # Validate the output structure to make sure that it can be unflattened.
    names = sorted(output)
    for i in range(len(names) - 1):
        name, next_name = names[i : i + 2]
        if next_name.startswith(name + "/"):
            raise ValueError(f"Leaf '{name}' aliases a node of '{next_name}'")

    return unflatten_dict(output)


def apply_tree(
    tree: at.PyTree[T], selector: at.PyTree[S], fn: Callable[[T, S], T], *, strict: bool = False
) -> at.PyTree[T]:
    tree = flatten_dict(tree)
    selector = flatten_dict(selector)

    def transform(k: str, v: T) -> T:
        if k in selector:
            return fn(v, selector[k])
        return v

    if strict:
        for k in selector:
            if k not in tree:
                raise ValueError(f"Selector key {k} not found in tree")

    return unflatten_dict({k: transform(k, v) for k, v in tree.items()})


def pad_to_dim(x: np.ndarray, target_dim: int, axis: int = -1, value: float = 0.0) -> np.ndarray:
    """Pad an array to the target dimension with zeros along the specified axis."""
    current_dim = x.shape[axis]
    if current_dim < target_dim:
        pad_width = [(0, 0)] * len(x.shape)
        pad_width[axis] = (0, target_dim - current_dim)
        return np.pad(x, pad_width, constant_values=value)
    return x


def make_bool_mask(*dims: int) -> tuple[bool, ...]:
    """Make a boolean mask for the given dimensions.

    Example:
        make_bool_mask(2, -2, 2) == (True, True, False, False, True, True)
        make_bool_mask(2, 0, 2) == (True, True, True, True)

    Args:
        dims: The dimensions to make the mask for.

    Returns:
        A tuple of booleans.
    """
    result = []
    for dim in dims:
        if dim > 0:
            result.extend([True] * (dim))
        else:
            result.extend([False] * (-dim))
    return tuple(result)


def _assert_quantile_stats(norm_stats: at.PyTree[NormStats]) -> None:
    for k, v in flatten_dict(norm_stats).items():
        if v.q01 is None or v.q99 is None:
            raise ValueError(
                f"quantile stats must be provided if use_quantile_norm is True. Key {k} is missing q01 or q99."
            )