"""
Contains observation-level transforms used in the octo data pipeline. These transforms operate on the
"observation" dictionary, and are applied at a per-frame level.
"""

from typing import Mapping, Optional, Tuple, Union

import tensorflow as tf
import logging

import src.data.dlimp as dl

from .obs_transforms import (
    image_dropout
)

log = logging.getLogger(__name__)

def augment(
    obs: dict,
    seed: tf.Tensor,
    augment_kwargs: Union[dict, Mapping[str, dict]],
) -> dict:
    """Augments images, skipping padding images."""
    if not hasattr(augment_kwargs, "items"):
        raise ValueError(
            "augment_kwargs must be a dict with keys corresponding to image names, or a single dict "
            "with an 'augment_order' key."
        )
    
    if "interleaved_instruction" in obs and "interleaved_instruction" in augment_kwargs:
        kwargs = augment_kwargs["interleaved_instruction"]
        def augment_fn(i, img):
            return dl.transforms.augment_image(img, **kwargs, seed=seed + i)
        obs["interleaved_instruction"]["image_instruction"] = tf.map_fn(
            lambda x: augment_fn(x[0], x[1]),
            (tf.range(tf.shape(obs["interleaved_instruction"]["image_instruction"])[0]), obs["interleaved_instruction"]["image_instruction"]),
            fn_output_signature=tf.uint8
        )
    
    image_names = {key[6:] for key in obs if key.startswith("image_")}
    
    # "augment_order" is required in augment_kwargs, so if it's there, we can assume that the user has passed
    # in a single augmentation dict (otherwise, we assume that the user has passed in a mapping from image
    # name to augmentation dict)
    if "augment_order" in augment_kwargs:
        augment_kwargs = {name: augment_kwargs for name in image_names}

    for i, name in enumerate(image_names):
        if name not in augment_kwargs:
            continue
        kwargs = augment_kwargs[name]
        log.debug(f"Augmenting image_{name} with kwargs {kwargs}")
        obs[f"image_{name}"] = tf.cond(
            obs["pad_mask_dict"][f"image_{name}"],
            lambda: dl.transforms.augment_image(
                obs[f"image_{name}"],
                **kwargs,
                seed=seed + i,  # augment each image differently
            ),
            lambda: obs[f"image_{name}"],  # skip padding images
        )


    return obs

def decode_and_resize(
    obs: dict,
    resize_size: Union[Tuple[int, int], Mapping[str, Tuple[int, int]]],
    depth_resize_size: Union[Tuple[int, int], Mapping[str, Tuple[int, int]]],
) -> dict:
    """Decodes images and depth images, and then optionally resizes them."""
    
    if "interleaved_instruction" in obs:
        decoded_imgs = tf.map_fn(
            lambda img: tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8),
            obs["interleaved_instruction"]["image_instruction"],
            dtype=tf.uint8
        )
        obs["interleaved_instruction"]["image_instruction"] = tf.map_fn(
            lambda img: dl.transforms.resize_image(img, size=resize_size["interleaved_instruction"]),
            decoded_imgs,
            dtype=tf.uint8
        )

    # just gets the part after "image_" or "depth_"
    image_names = {key[6:] for key in obs if key.startswith("image_")}
    depth_names = {key[6:] for key in obs if key.startswith("depth_")}

    if isinstance(resize_size, tuple):
        resize_size = {name: resize_size for name in image_names}
    if isinstance(depth_resize_size, tuple):
        depth_resize_size = {name: depth_resize_size for name in depth_names}

    for name in image_names:
        if name not in resize_size:
            log.warning(
                f"No resize_size was provided for image_{name}. This will result in 1x1 "
                "padding images, which may cause errors if you mix padding and non-padding images."
            )
        image = obs[f"image_{name}"]
        if image.dtype == tf.string:
            if tf.strings.length(image) == 0:
                # this is a padding image
                image = tf.zeros((*resize_size.get(name, (1, 1)), 3), dtype=tf.uint8)
            else:
                image = tf.io.decode_image(
                    image, expand_animations=False, dtype=tf.uint8
                )
        elif image.dtype != tf.uint8:
            raise ValueError(
                f"Unsupported image dtype: found image_{name} with dtype {image.dtype}"
            )
        if name in resize_size:
            image = dl.transforms.resize_image(image, size=resize_size[name])
        obs[f"image_{name}"] = image

    for name in depth_names:
        if name not in depth_resize_size:
            log.warning(
                f"No depth_resize_size was provided for depth_{name}. This will result in 1x1 "
                "padding depth images, which may cause errors if you mix padding and non-padding images."
            )
        depth = obs[f"depth_{name}"]
        if depth.dtype == tf.string:
            if tf.strings.length(depth) == 0:
                # this is a padding image
                depth = tf.zeros(
                    (*depth_resize_size.get(name, (1, 1)), 1), dtype=tf.float32
                )
            else:
                depth = tf.io.decode_image(
                    depth, expand_animations=False, dtype=tf.float32
                )[..., 0]
        elif depth.dtype != tf.float32:
            raise ValueError(
                f"Unsupported depth dtype: found depth_{name} with dtype {depth.dtype}"
            )
        if name in depth_resize_size:
            depth = dl.transforms.resize_depth_image(
                depth, size=depth_resize_size[name]
            )
        obs[f"depth_{name}"] = depth

    return obs
