"""
datasets.py

Lightweight PyTorch Dataset Definition for wrapping RLDS TFDS Pipeline; just defines transform from RLDS default
format to OpenVLA, IterableDataset shim.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Tuple, Type
import collections
import os
import numpy as np
import cv2
import torch
from PIL import Image
import torchvision
from torchvision.transforms import transforms
from torch.utils.data import Dataset, IterableDataset
from transformers import PreTrainedTokenizerBase
from data.utils.som_tom import som_prompting, tom_prompting

# from prismatic.models.backbones.llm.prompting import PromptBuilder
# from prismatic.models.backbones.vision import ImageTransform
from ..action_tokenizer import ActionTokenizer
from .rlds import make_interleaved_dataset, make_single_dataset
from .rlds.oxe import OXE_NAMED_MIXTURES, get_oxe_dataset_kwargs_and_weights
from .rlds.utils.data_utils import NormalizationType

# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
IGNORE_INDEX = -100
from typing import Callable, Dict, Sequence, Tuple

def tree_map(fn: Callable, tree: dict) -> dict:
    """Maps a function over a nested dictionary."""
    return {k: tree_map(fn, v) if isinstance(v, dict) else fn(v) for k, v in tree.items()}


def tree_map_with_key(fn: Callable, tree: dict, keys: Sequence = ()) -> dict:
    """Maps a function over a nested dictionary."""
    return {
        k: tree_map_with_key(fn, v, (*keys, k)) if isinstance(v, dict) else fn((*keys, k), v) for k, v in tree.items()
    }


@dataclass
class RLDSBatchTransform:
    action_tokenizer: ActionTokenizer
    base_tokenizer: PreTrainedTokenizerBase
    image_transform: None # ImageTransform
    prompt_builder_fn: None # Type[PromptBuilder]
    visual_tracker: None
    dataset_settings: None
    data_root_dir: str = "/mnt/vlpdatasets"
    predict_stop_token: bool = True
    trace_folder: str = "open-x-traces-v2"
    image_folder: str = "open-x-images-v2"
    local_run: bool = False

    def __call__(self, rlds_batch: Dict[str, Any]) -> Dict[str, Any]:
        """Converts a RLDS batch to the format expected by the OpenVLA collator/models."""
        dataset_name, action = rlds_batch["dataset_name"], rlds_batch["action"][0]
        img = Image.fromarray(rlds_batch["observation"]["image_primary"][0])
        imgs_future = [Image.fromarray(img) for img in rlds_batch["observation_future"]["image_primary"]]
        lang = rlds_batch["task"]["language_instruction"].decode().lower()        
        traj_index = rlds_batch['_traj_index']
        frame_index = rlds_batch['_frame_index']

        action_token_ids = self.action_tokenizer.encode_actions_to_token_ids(action)

        # Construct Chat-based Prompt =>> Input is default query + language instruction, output are the action tokens
        convs = [
            {"role": "system", "content": "You are agent that can see, talk and act."},
            {"role": "user", "content": f"<image>\nWhat action should the robot take to {lang}?"},
            {"role": "assistant", "content": "<action>"},
        ]
        prompt = self.base_tokenizer.apply_chat_template(convs, tokenize=False, add_generation_prompt=False)

        # Tokenize (w/ `base_tokenizer`)
        input_ids = self.base_tokenizer(prompt, add_special_tokens=True).input_ids

        action_token_len = len(action_token_ids)
        action_placeholder_token_id = self.base_tokenizer.convert_tokens_to_ids("<action>")
        # # replace the action_placeholder_token_id with action_token_ids in input_ids
        input_ids = list(input_ids)
        input_ids_filled = []
        for i, token_id in enumerate(input_ids):
            if token_id == action_placeholder_token_id:
                input_ids_filled.extend(action_token_ids.tolist())
            else:
                input_ids_filled.append(token_id)

        # Tensorize =>> Run Image Transform to get `pixel_values` =>> Return
        #   =>> IMPORTANT :: IF WE'RE USING HF LLM.forward(..., labels=labels), SHIFTING HAPPENS _INSIDE_ MODEL!
        input_ids, labels = torch.tensor(input_ids_filled), torch.tensor(input_ids_filled)

        pixel_values = transforms.Compose([transforms.ToTensor()])(img)
        image_pt = self.image_transform(img, return_tensors='pt')
        images = collections.defaultdict(list)
        for key, val in image_pt.items():
            images[key].append(val)

        pixel_values_future = torch.stack([transforms.Compose([transforms.ToTensor()])(item) for item in imgs_future], dim=0)

        # extract visual traces
        # trace_folder = self.trace_folder
        # if not os.path.exists(trace_folder):
        #     trace_folder = "./open-x-traces-v2"
        # trace_file = f"{dataset_name}/{traj_index}/{frame_index}.pth"
        # trace_path = os.path.join(trace_folder, trace_file)
        # if not os.path.exists(trace_path):        
        #     pixel_values_seq = torch.cat([pixel_values.unsqueeze(0), pixel_values_future], dim=0).unsqueeze(0)
        #     out = self.visual_tracker.extract_visual_trace(pixel_values_seq*255)
        #     # self.visual_tracker.visualize(*out)
        #     # save the visual trace to disk
        #     trace_info = {
        #         'dataset_name': dataset_name,
        #         'traj_index': traj_index,
        #         'frame_index': frame_index,
        #         'lang': lang,
        #         'action': action,
        #         'trace_info': out[1:]
        #     }
        #     os.makedirs(os.path.dirname(trace_path), exist_ok=True)
        #     torch.save(trace_info, trace_path)
        
        # save image
        # image_folder = self.image_folder
        # if not os.path.exists(image_folder):
        #     image_folder = "./open-x-images-v2"
        # image_file = f"{dataset_name}/{traj_index}/{frame_index}.jpg"
        # image_path = os.path.join(image_folder, image_file)      
        # if not os.path.exists(image_path):
        #     os.makedirs(os.path.dirname(image_path), exist_ok=True)
        #     img.save(image_path)

        # [CRITICAL] We do not want to take the loss for anything but the predicted action tokens!
        # NOTE: we add 2 to the length of the action to account for the \n\n and <|eot_id|> tokens!
        labels[: -(action_token_len + 2)] = IGNORE_INDEX
        if not self.predict_stop_token:
            labels[-1] = IGNORE_INDEX
                
        return dict(pixel_values=images['pixel_values'], image_sizes=images['image_sizes'], pixel_values_future=pixel_values_future, input_ids=input_ids, labels=labels, dataset_name=dataset_name) 
        # return dict(pixel_values=pixel_values, pixel_values_future=pixel_values_future, action=action, conversation=conversation, dataset_name=dataset_name)

class RLDSDataset(IterableDataset):
    def __init__(
        self,
        data_root_dir: Path,
        data_mix: str,
        batch_transform: RLDSBatchTransform,
        resize_resolution: Tuple[int, int],
        shuffle_buffer_size: int = 256_000,
        train: bool = True,
        image_aug: bool = False,
        future_action_window_size: int = 0,
    ) -> None:
        """Lightweight wrapper around RLDS TFDS Pipeline for use with PyTorch/OpenVLA Data Loaders."""
        self.data_root_dir, self.data_mix, self.batch_transform = data_root_dir, data_mix, batch_transform

        # Configure RLDS Dataset(s)
        if self.data_mix in OXE_NAMED_MIXTURES:
            mixture_spec = OXE_NAMED_MIXTURES[self.data_mix]
        else:
            # Assume that passed "mixture" name is actually a single dataset -- create single-dataset "mix"
            mixture_spec = [(self.data_mix, 1.0)]
        
        # fmt: off
        per_dataset_kwargs, weights = get_oxe_dataset_kwargs_and_weights(
            self.data_root_dir,
            mixture_spec,
            load_camera_views=("primary",),
            load_depth=False,
            load_proprio=False,
            load_language=True,
            action_proprio_normalization_type=NormalizationType.BOUNDS_Q99,
        )
        rlds_config = dict(
            traj_transform_kwargs=dict(
                window_size=1,                                      # If we wanted to feed / predict more than one step
                future_action_window_size=future_action_window_size,                        # For action chunking
                skip_unlabeled=True,                                # Skip trajectories without language labels
                goal_relabeling_strategy="uniform",                 # Goals are currently unused
            ),
            frame_transform_kwargs=dict(
                resize_size=resize_resolution,
                num_parallel_calls=16,                          # For CPU-intensive ops (decoding, resizing, etc.)
            ),
            dataset_kwargs_list=per_dataset_kwargs,
            shuffle_buffer_size=shuffle_buffer_size,
            sample_weights=weights,
            balance_weights=True,
            traj_transform_threads=len(mixture_spec),
            traj_read_threads=len(mixture_spec),
            train=train,
        )

        # If applicable, enable image augmentations
        if image_aug:
            rlds_config["frame_transform_kwargs"].update({"image_augment_kwargs" : dict(
                random_resized_crop=dict(scale=[0.9, 0.9], ratio=[1.0, 1.0]),
                random_brightness=[0.2],
                random_contrast=[0.8, 1.2],
                random_saturation=[0.8, 1.2],
                random_hue=[0.05],
                augment_order=[
                    "random_resized_crop",
                    "random_brightness",
                    "random_contrast",
                    "random_saturation",
                    "random_hue",
                ],
            )}),
        # fmt: on

        # Initialize RLDS Dataset
        self.dataset, self.dataset_length, self.dataset_statistics = self.make_dataset(rlds_config)

    def make_dataset(self, rlds_config):
        return make_interleaved_dataset(**rlds_config)

    def __iter__(self) -> Dict[str, Any]:
        for rlds_batch in self.dataset.as_numpy_iterator():
            yield self.batch_transform(rlds_batch)

    def __len__(self) -> int:
        return self.dataset_length

    # === Explicitly Unused ===
    def __getitem__(self, idx: int) -> None:
        raise NotImplementedError("IterableDataset does not implement map-style __getitem__; see __iter__ instead!")


class EpisodicRLDSDataset(RLDSDataset):
    """Returns full episodes as list of steps instead of individual transitions (useful for visualizations)."""

    def make_dataset(self, rlds_config):
        per_dataset_kwargs = rlds_config["dataset_kwargs_list"]
        assert len(per_dataset_kwargs) == 1, "Only support single-dataset `mixes` for episodic datasets."

        return make_single_dataset(
            per_dataset_kwargs[0],
            train=rlds_config["train"],
            traj_transform_kwargs=rlds_config["traj_transform_kwargs"],
            frame_transform_kwargs=rlds_config["frame_transform_kwargs"],
        )

    def __iter__(self) -> Dict[str, Any]:
        for rlds_batch in self.dataset.as_numpy_iterator():
            out = [
                self.batch_transform(tree_map(lambda x: x[i], rlds_batch))  # noqa: B023
                for i in range(rlds_batch["action"].shape[0])
            ]
            yield out


class DummyDataset(Dataset):
    def __init__(
        self,
        action_tokenizer: ActionTokenizer,
        base_tokenizer: PreTrainedTokenizerBase,
        image_transform: None, # ImageTransform,
        prompt_builder_fn: None # Type[PromptBuilder],
    ) -> None:
        self.action_tokenizer = action_tokenizer
        self.base_tokenizer = base_tokenizer
        self.image_transform = image_transform
        self.prompt_builder_fn = prompt_builder_fn

        # Note =>> We expect the dataset to store statistics for action de-normalization. Specifically, we store the
        # per-dimension 1st and 99th action quantile. The values below correspond to "no normalization" for simplicity.
        self.dataset_statistics = {
            "dummy_dataset": {
                "action": {"q01": np.zeros((7,), dtype=np.float32), "q99": np.ones((7,), dtype=np.float32)}
            }
        }

    def __len__(self):
        # TODO =>> Replace with number of elements in your dataset!
        return 10000

    def __getitem__(self, idx):
        # TODO =>> Load image, action and instruction from disk -- we use dummy values
        image = Image.fromarray(np.asarray(np.random.rand(224, 224, 3) * 255.0, dtype=np.uint8))
        action = np.asarray(np.random.rand(7), dtype=np.float32)
        instruction = "do something spectacular"

        # Add instruction to VLA prompt
        prompt_builder = self.prompt_builder_fn("openvla")
        conversation = [
            {"from": "human", "value": f"What action should the robot take to {instruction}?"},
            {"from": "gpt", "value": self.action_tokenizer(action)},
        ]
        for turn in conversation:
            prompt_builder.add_turn(turn["from"], turn["value"])

        # Tokenize (w/ `base_tokenizer`)
        input_ids = self.base_tokenizer(prompt_builder.get_prompt(), add_special_tokens=True).input_ids
        labels = list(input_ids)

        # Tensorize =>> Run Image Transform to get `pixel_values` =>> Return
        #   =>> IMPORTANT :: IF WE'RE USING HF .forward(..., labels=labels), SHIFTING HAPPENS _INSIDE_ MODEL!
        input_ids, labels = torch.tensor(input_ids), torch.tensor(labels)
        pixel_values = self.image_transform(image)

        # [CRITICAL] We do not want to take the loss for anything but the predicted action tokens!
        labels[: -(len(action) + 1)] = IGNORE_INDEX

        return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels)