from collections.abc import Callable, Mapping, Sequence
import dataclasses
import re
from typing import 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.models import tokenizer as _tokenizer
from openpi.shared import array_typing as at
from openpi.shared import normalize as _normalize

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


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


@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 and "prompt" not in data:
            data["prompt"] = np.asarray(self.prompt)
        return data


@dataclasses.dataclass(frozen=True)
class InjectDatasetName(DataTransformFn):
    dataset_name: str | None

    def __call__(self, data: DataDict) -> DataDict:
        name = self.dataset_name.split("/")[-1] if "/" in self.dataset_name else self.dataset_name
        data["prompt"] = f"[dataset:{name}] {data['prompt']}"
        return data


@dataclasses.dataclass(frozen=True)
class Normalize(DataTransformFn):
    data_mask: list
    norm_stats: at.PyTree[NormStats] | None
    # 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:
            new_stats = dict(self.norm_stats)  # 浅拷贝，避免副作用
            if 'action' in new_stats:
                new_stats['actions'] = new_stats.pop('action')
            if 'observation.state' in new_stats:
                new_stats['state'] = new_stats.pop('observation.state')
            object.__setattr__(self, 'norm_stats', new_stats)

            if "actions" not in self.norm_stats or "state" not in self.norm_stats:
                raise ValueError("norm_stats must contain 'actions' and 'state' keys for normalization.")

        else:
            raise ValueError("norm_stats cannot be None. Please provide valid normalization statistics.")

    def __call__(self, data: DataDict) -> DataDict:
        if self.norm_stats is None:
            raise ValueError("norm_stats cannot be None. Please provide valid normalization statistics.")

        # # special case for OXE Dataset
        # if self.norm_stats["actions"].mean.shape[-1] == 7 and self.norm_stats["state"].mean.shape[-1] == 8:
        #     return data
        
        return apply_tree(
            data,
            self.norm_stats,
            self._normalize,
            strict=self.strict,
        )

    def _normalize(self, x, stats: NormStats):
        half_len = len(self.data_mask)
        data_mask_x = self.data_mask + [0] * half_len
        mask_x = np.array(data_mask_x, dtype=bool)
        count_true = np.sum(mask_x)

        pad_mean = np.zeros(x.shape[-1])
        pad_std = np.ones(x.shape[-1])

        if x.ndim == 1:
            if stats.mean.shape[-1] == 8 and stats.std.shape[-1] == 8 and count_true == 7:
                pad_mean[:8] = stats.mean
                pad_std[:8] = stats.std
                x[:8] = np.clip(x[:8], stats.q01, stats.q99)
            else:
                x[mask_x] = np.clip(x[mask_x], stats.q01, stats.q99)
                pad_mean[mask_x] = stats.mean
                pad_std[mask_x] = stats.std
        elif x.ndim == 2:
            x[:, mask_x[:half_len]] = np.clip(
                x[:, mask_x[:half_len]],
                stats.q01,
                stats.q99
            )
            pad_mean[..., mask_x[:half_len]] = stats.mean
            pad_std[..., mask_x[:half_len]] = stats.std
        else:
            raise ValueError(f"Unsupported x.ndim: {x.ndim}")
        
        return (x - pad_mean) / (pad_std + 1e-6)

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


@dataclasses.dataclass(frozen=True)
class Unnormalize(DataTransformFn):
    data_mask: list
    norm_stats: at.PyTree[NormStats] | None

    def __post_init__(self):
        if self.norm_stats is not None:
            new_stats = dict(self.norm_stats)  # 浅拷贝，避免副作用
            if 'action' in new_stats:
                new_stats['actions'] = new_stats.pop('action')
            if 'observation.state' in new_stats:
                new_stats['state'] = new_stats.pop('observation.state')
            object.__setattr__(self, 'norm_stats', new_stats)

            if "actions" not in self.norm_stats or "state" not in self.norm_stats:
                raise ValueError("norm_stats must contain 'actions' and 'state' keys for normalization.")
        else:
            raise ValueError("norm_stats cannot be None. Please provide valid normalization statistics.")

    def __call__(self, data: DataDict) -> DataDict:
        if self.norm_stats is None:
            raise ValueError("norm_stats cannot be None. Please provide valid normalization statistics.")

        # # special case for OXE Dataset
        # if self.norm_stats["actions"].mean.shape[-1] == 7 and self.norm_stats["state"].mean.shape[-1] == 8:
        #     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,
            strict=False,
        )

    def _unnormalize(self, x, stats: NormStats):
        half_len = len(self.data_mask)
        data_mask_x = self.data_mask + [0] * len(self.data_mask)
        mask_x = np.array(data_mask_x, dtype=bool)
        count_true = np.sum(mask_x)

        pad_mean = np.zeros(x.shape[-1])
        pad_std = np.ones(x.shape[-1])
        
        if x.ndim == 1:
            if stats.mean.shape[-1] == 8 and stats.std.shape[-1] == 8 and count_true == 7:
                pad_mean[:8] = stats.mean
                pad_std[:8] = stats.std
            else:
                pad_mean[..., mask_x] = stats.mean
                pad_std[..., mask_x] = stats.std
        elif x.ndim == 2:
            pad_mean[..., mask_x[:half_len]] = stats.mean
            pad_std[..., mask_x[:half_len]] = stats.std
        else:
            raise ValueError(f"Unsupported x.ndim: {x.ndim}")
        
        return x * (pad_std + 1e-6) + pad_mean

    def _unnormalize_quantile(self, x, stats: NormStats):
        assert stats.q01 is not None
        assert stats.q99 is not None
        return (x + 1.0) / 2.0 * (stats.q99 - stats.q01 + 1e-6) + stats.q01


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

    def __call__(self, data: DataDict) -> DataDict:
        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
    data_mask: list

    def __call__(self, data: DataDict) -> DataDict:
        if "actions" not in data or self.mask is None:
            assert "state" in data, "Cannot apply DeltaActions transform without 'actions' or 'mask'"
            return data

        state, actions = data["state"], data["actions"]

        mask = np.asarray(self.mask)
        half_len = len(self.data_mask)
        data_mask = np.array(self.data_mask + [0] * half_len, dtype=bool)

        actions[..., data_mask[:half_len]] -= np.expand_dims(np.where(mask, state[..., data_mask], 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
    data_mask: list
    def __call__(self, data: DataDict) -> DataDict:
        if "actions" not in data or self.mask is None:
            assert "state" in data, "Cannot apply AbsoluteActions transform without 'actions' or 'mask'"
            return data
        # print("************AbsoluteActions")
        state, actions = data["state"], data["actions"]
        mask = np.asarray(self.mask)
        half_len = len(self.data_mask)
        data_mask = np.array(self.data_mask + [0] * half_len, dtype=bool)

        actions[..., data_mask[:half_len]] += np.expand_dims(np.where(mask, state[..., data_mask], 0), axis=-2)
        data["actions"] = actions

        return data


@dataclasses.dataclass(frozen=True)
class DeltaActions_Libero(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_Libero(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

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

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

        tokens, token_masks = self.tokenizer.tokenize(prompt)
        return {**data, "tokenized_prompt": tokens, "tokenized_prompt_mask": token_masks}


@dataclasses.dataclass(frozen=True)
class TokenizeFASTInputs(DataTransformFn):
    tokenizer: _tokenizer.FASTTokenizer

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

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

        state, actions = data["state"], data.get("actions")
        tokens, token_mask, ar_mask, loss_mask = self.tokenizer.tokenize(prompt, state, actions)
        return {
            **data,
            "tokenized_prompt": tokens,
            "tokenized_prompt_mask": token_mask,
            "token_ar_mask": ar_mask,
            "token_loss_mask": loss_mask,
        }


@dataclasses.dataclass(frozen=True)
class ExtractFASTActions(DataTransformFn):
    tokenizer: _tokenizer.FASTTokenizer
    action_horizon: int
    action_dim: int

    def __call__(self, data: DataDict) -> DataDict:
        if "actions" not in data:
            return data
        # Model outputs are saved in "actions", but for FAST models they represent tokens.
        tokens = data.pop("actions")
        actions = self.tokenizer.extract_actions(tokens.astype(np.int32), self.action_horizon, self.action_dim)
        return {
            **data,
            "actions": actions,
        }


@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 "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}")

        return {**data, "prompt": prompt}


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, data_mask: list, axis: int = -1) -> np.ndarray:
    """Pad an array to the target dimension with zeros along the specified axis."""
    mask = np.array(data_mask, dtype=bool)
    count_ones = np.sum(mask)
    
    if x.shape[axis] != count_ones:
        raise ValueError(f"x 在轴 {axis} 的长度 ({x.shape[axis]}) 与 data_mask 中为1的数量 ({count_ones}) 不一致")
    new_shape = x.shape[:-1] + (target_dim,)
    result = np.zeros(new_shape)
    result[..., mask] = x

    return result


def pad_state_to_dim(x: np.ndarray, target_dim: int, data_mask: list, axis: int = -1) -> np.ndarray:
    """Pad an array to the target dimension with zeros along the specified axis."""
    mask = np.array(data_mask, dtype=bool)
    count_ones = np.sum(mask)
    
    if x.shape[axis] != count_ones:
        raise ValueError(f"x 在轴 {axis} 的长度 ({x.shape[axis]}) 与 data_mask 中为1的数量 ({count_ones}) 不一致")
    
    new_shape = x.shape[:-1] + (target_dim,)
    result = np.zeros(new_shape, dtype=np.float32)

    data_mask_x = data_mask + [0] * len(data_mask)
    mask_x = np.array(data_mask_x, dtype=bool)
    result[..., mask_x] = x

    data_mask_mask = [0] * len(data_mask) + data_mask
    mask_mask = np.array(data_mask_mask, dtype=bool)
    result[..., mask_mask] = 1

    return result


def pad_libero_oxe_state_to_dim(x: np.ndarray, target_dim: int, data_mask: list, axis: int = -1) -> np.ndarray:
    """Pad an array to the target dimension with zeros along the specified axis."""
    
    new_shape = x.shape[:-1] + (target_dim,)
    result = np.zeros(new_shape, dtype=np.float32)

    result[..., :8] = x

    data_mask_mask = [0] * len(data_mask) + data_mask
    mask_mask = np.array(data_mask_mask, dtype=bool)
    result[..., mask_mask] = 1

    return result


def pad_calvin_state_to_dim(x: np.ndarray, target_dim: int, data_mask: list, axis: int = -1) -> np.ndarray:
    """Pad an array to the target dimension with zeros along the specified axis."""
    # 将 data_mask 转换为布尔数组
    mask = np.array(data_mask, dtype=bool)
    if mask[0] == True:
        x = x[..., :7]
    else:
        x = x[..., 7:]

    count_ones = np.sum(mask)
    if x.shape[axis] != count_ones:
        raise ValueError(f"x 在轴 {axis} 的长度 ({x.shape[axis]}) 与 data_mask 中为1的数量 ({count_ones}) 不一致")

    new_shape = x.shape[:-1] + (target_dim,)
    result = np.zeros(new_shape, dtype=np.float32)

    data_mask_x = data_mask + [0] * len(data_mask)
    mask_x = np.array(data_mask_x, dtype=bool)
    result[..., mask_x] = x

    data_mask_mask = [0] * len(data_mask) + data_mask
    mask_mask = np.array(data_mask_mask, dtype=bool)
    result[..., mask_mask] = 1

    return result


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."
            )
