
import pickle
import fnmatch
import cv2
cv2.setNumThreads(1)
from aloha_scripts.utils import *
import time
from torch.utils.data import TensorDataset, DataLoader
import torchvision.transforms as transforms
import os
import json
import numpy as np
from aloha_scripts.lerobot_constants import LEROBOT_TASK_CONFIGS
import torch

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata

from typing import Protocol, SupportsIndex, TypeVar
T_co = TypeVar("T_co", covariant=True)
from tqdm import tqdm




class Dataset(Protocol[T_co]):
    """Interface for a dataset with random access."""

    def __getitem__(self, index: SupportsIndex) -> T_co:
        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")

    def __len__(self) -> int:
        raise NotImplementedError("Subclasses of Dataset should implement __len__.")

class TransformedDataset(Dataset[T_co]):
    def __init__(self, dataset: Dataset, norm_stats, camera_names,policy_class, robot=None, rank0_print=print, vla_data_post_process=None, data_args=None):
        self._dataset = dataset
        self.norm_stats = norm_stats
        self.camera_names = camera_names
        self.data_args = data_args
        self.robot = robot
        self.vla_data_post_process = vla_data_post_process
        self.rank0_print = rank0_print
        self.policy_class = policy_class
        # augment images for training (default for dp and scaledp)
        self.augment_images = True

        original_size = (480, 640)
        new_size = eval(self.data_args.image_size_stable) # 320, 240
        new_size = (new_size[1], new_size[0])
        ratio = 0.95
        self.transformations = [
            # todo resize
            # transforms.Resize(size=original_size, antialias=True),
            transforms.RandomCrop(size=[int(original_size[0] * ratio), int(original_size[1] * ratio)]),
            transforms.Resize(original_size, antialias=True),
            transforms.RandomRotation(degrees=[-5.0, 5.0], expand=False),
            transforms.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5),  # , hue=0.08)
            transforms.Resize(size=new_size, antialias=True),
        ]

        if 'diffusion' in self.policy_class.lower() or 'scale_dp' in self.policy_class.lower():
            self.augment_images = True
        else:
            self.augment_images = False

        # self.rank0_print(f"########################Current Image Size is [{self.data_args.image_size_stable}]###################################")
        # self.rank0_print(f"{RED}policy class: {self.policy_class}; augument: {self.augment_images}{RESET}")
        # a=self.__getitem__(100) # initialize self.is_sim and self.transformations
        # if len(self.camera_names) > 2:
            # self.rank0_print("%"*40)
            # self.rank0_print(f"The robot is {RED} {self.robot} {RESET} | The camera views: {RED} {self.camera_names} {RESET} | The history length: {RED} {self.data_args.history_images_length} {RESET}")
        self.is_sim = False

    def __getitem__(self, index: SupportsIndex) -> T_co:
        data = self._dataset[index]

        is_pad = data['action_is_pad']
        # sub_reason = data.meta.

        language_raw = self._dataset.meta.episodes[data['episode_index']]["language_dict"]['language_raw']
        if self.data_args.use_reasoning:
            none_counter = 0
            for k in ['substep_reasonings', 'reason']:
                vals = self._dataset.meta.episodes[data['episode_index']]["language_dict"][k]
                if vals is not None:
                    if k == 'substep_reasonings':
                        sub_reasoning = vals[data['frame_index']]
                    else:
                        sub_reasoning = vals
                # else:
                #     sub_reasoning = 'Next action:'
                else:
                    none_counter += 1
            if none_counter == 2:
                self.rank0_print(f"{RED} In {self._dataset.meta.repo_id}-{index}:{k} is None {RESET}")

        else:
            sub_reasoning = 'Default outputs no reasoning'

        all_cam_images = []
        for cam_name in self.camera_names:
            # Check if image is available
            image = data[cam_name].numpy()

            # Transpose image to (height, width, channels) if needed
            if image.shape[0] == 3:  # If image is in (channels, height, width)
                image = np.transpose(image, (1, 2, 0))  # Now it's (height, width, channels

                # image_dict[cam_name] = image  # resize

            all_cam_images.append(image)

        all_cam_images = np.stack(all_cam_images, axis=0)

        # construct observations, and scale 0-1 to 0-255
        image_data = torch.from_numpy(all_cam_images) * 255
        image_data = image_data.to(dtype=torch.uint8)
        # construct observations
        qpos_data = data['observation.state'].float()
        action_data = data['action'].float()

        # channel last
        image_data = torch.einsum('k h w c -> k c h w', image_data)

        if self.augment_images:
            for transform in self.transformations:
                image_data = transform(image_data)

        norm_stats = self.norm_stats
        # normalize to [-1, 1]
        action_data = ((action_data - norm_stats["action_min"]) / (norm_stats["action_max"] - norm_stats["action_min"])) * 2 - 1

        qpos_data = (qpos_data - norm_stats["qpos_mean"]) / norm_stats["qpos_std"]
        # std = 0.05
        # noise = std * torch.randn_like(qpos_data)
        # qpos_noise = qpos_data + noise
        # new_std = torch.sqrt(torch.tensor(1 ** 2 + std ** 2))
        # normalized_qpos = qpos_noise / new_std
        # qpos_data = normalized_qpos.float()
        sample = {
            'image': image_data,
            'state': qpos_data,
            'action': action_data,
            'is_pad': is_pad,
            'raw_lang': language_raw,
            'reasoning': sub_reasoning
        }

        return self.vla_data_post_process.forward_process(sample, use_reasoning=self.data_args.use_reasoning)

    def __len__(self) -> int:
        return len(self._dataset)
def get_norm_stats(dataset_list):
    """
    caculate all data action and qpos(robot state ) mean and std
    """
    key_name_list=["observation.state","action"]

    all_qpos_data = []
    mean_list = []
    std_list = []
    length_list = []
    state_min_list = []
    state_max_list = []
    action_mean_list = []
    action_std_list = []
    action_max_list = []
    action_min_list = []

    # Collect data from each dataset
    for dataset in tqdm(dataset_list):

        mean_tensor = dataset.meta.stats["observation.state"]["mean"]
        std_tensor = dataset.meta.stats["observation.state"]["std"]
        state_max = dataset.meta.stats["observation.state"]["max"]
        state_min = dataset.meta.stats["observation.state"]["min"]

        action_mean = dataset.meta.stats["action"]["mean"]
        action_std = dataset.meta.stats["action"]["std"]
        action_min = dataset.meta.stats["action"]["min"]
        action_max = dataset.meta.stats["action"]["max"]
        # Ensure the tensors are on CPU and convert to numpy arrays
        mean_array = mean_tensor.cpu().numpy() if mean_tensor.is_cuda else mean_tensor.numpy()
        std_array = std_tensor.cpu().numpy() if std_tensor.is_cuda else std_tensor.numpy()
        state_max = state_max.cpu().numpy() if state_max.is_cuda else state_max.numpy()
        state_min = state_min.cpu().numpy() if state_min.is_cuda else state_min.numpy()

        action_mean = action_mean.cpu().numpy() if action_mean.is_cuda else action_mean.numpy()
        action_std = action_std.cpu().numpy() if action_std.is_cuda else action_std.numpy()
        action_min = action_min.cpu().numpy() if action_min.is_cuda else action_min.numpy()
        action_max = action_max.cpu().numpy() if action_max.is_cuda else action_max.numpy()

        # Append the arrays and the length of the dataset (number of samples)
        mean_list.append(mean_array)
        std_list.append(std_array)
        state_max_list.append(state_max)
        state_min_list.append(state_min)
        action_mean_list.append(action_mean)
        action_std_list.append(action_std)
        action_max_list.append(action_max)
        action_min_list.append(action_min)

        length_list.append(len(dataset))  # This is a single number, representing the number of samples

    # Convert lists to numpy arrays for easier manipulation
    mean_array = np.array(mean_list)  # Shape should be (num_datasets, 14)
    std_array = np.array(std_list)  # Shape should be (num_datasets, 14)
    length_array = np.array(length_list)  # Shape should be (num_datasets,)

    action_mean = np.array(action_mean_list)
    action_std = np.array(action_std_list)

    state_max = np.max(state_max_list, axis=0)
    state_min = np.min(state_min_list, axis=0)
    action_max = np.max(action_max_list, axis=0)
    action_min = np.min(action_min_list, axis=0)

    state_mean = np.sum(mean_array.T * length_array, axis=1) / np.sum(length_array)

    # To calculate the weighted variance (pooled variance):

    state_weighted_variance = np.sum(((length_array[:, None] - 1) * std_array ** 2 + (length_array[:, None] - 1) *mean_array**2),axis=0)/np.sum(length_array) - state_mean**2

    # Calculate the overall standard deviation (square root of variance)
    state_std = np.sqrt(state_weighted_variance)

    action_weighted_mean = np.sum(action_mean.T * length_array, axis=1) / np.sum(length_array)
    action_weighted_variance = np.sum(((length_array[:, None] - 1) * action_std ** 2 + (length_array[:, None] - 1) *action_mean**2),axis=0)/np.sum(length_array) - action_weighted_mean**2
    action_weighted_std = np.sqrt(action_weighted_variance)
    # Output the results
    print(f"Overall Weighted Mean: {state_mean}")
    print(f"Overall Weighted Std: {state_std}")

    eps = 0.0001
    stats = {"action_mean": action_weighted_mean, "action_std": action_weighted_std,
             "action_min": action_min - eps, "action_max": action_max + eps,
             "qpos_mean": state_mean, "qpos_std": state_std,
             }

    all_episode_len = len(all_qpos_data)
    return stats, all_episode_len

def create_dataset(repo_id, chunk_size, home_lerobot=None, local_debug=False) ->  Dataset:
    with open(os.path.join(home_lerobot, repo_id, "meta", 'info.json'), 'r') as f:
            data = json.load(f)
    fps = data['fps']
    delta_timestamps = {
        # "observation.state": [t / fps for t in range(args['chunk_size'])],
        "action": [t / fps for t in range(chunk_size)],
    }

    if local_debug:
        print(f"{RED} Warning only using first two episodes {RESET}")
        dataset = LeRobotDataset(repo_id, episodes=[0,1], delta_timestamps=delta_timestamps, local_files_only=True)
    else:
        dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, local_files_only=True)
    return dataset
def load_data(camera_names, chunk_size, config, rank0_print=print, policy_class=None, vla_data_post_process=None, **kwargs):
    repo_id_list = LEROBOT_TASK_CONFIGS[config['data_args'].task_name]['dataset_dir']
    dataset_list = []
    for repo_id in repo_id_list:
        dataset = create_dataset(repo_id, chunk_size, home_lerobot=config['data_args'].home_lerobot, local_debug=config['training_args'].local_debug)
        dataset_list.append(dataset)
    norm_stats, all_episode_len = get_norm_stats(dataset_list)
    train_dataset_list =[]
    robot = 'aloha' if config['action_head_args'].action_dim == 14 or ('aloha' in config['training_args'].output_dir) else 'franka'

    rank0_print(
        f"########################Current Image Size is [{config['data_args'].image_size_stable}]###################################")
    rank0_print(f"{RED}policy class: {policy_class};{RESET}")
    for dataset in dataset_list:
      train_dataset_list.append(TransformedDataset(
          dataset, norm_stats, camera_names, policy_class=policy_class, robot=robot,
          rank0_print=rank0_print, vla_data_post_process=vla_data_post_process, data_args=config['data_args']))

        # self.rank0_print("%"*40)
    rank0_print(
    f"The robot is {RED} {robot} {RESET} | The camera views: {RED} {camera_names} {RESET} | "
    f"The history length: {RED} {config['data_args'].history_images_length} | Data augmentation: {train_dataset_list[0].augment_images} {RESET}")


    train_dataset = torch.utils.data.ConcatDataset(train_dataset_list)
    # train_dataloder = DataLoader(train_dataset, batch_size=batch_size_train, shuffle=True, num_workers=8, pin_memory=True,prefetch_factor=2)
    # val_dataloader = None
    rank0_print(f"{RED}All images: {len(train_dataset)} {RESET}")

    return train_dataset, None, norm_stats

def get_norm_stats_by_tasks(dataset_path_list,args):
    data_tasks_dict = dict(
        fold_shirt=[],
        clean_table=[],
        others=[],
    )
    for dataset_path in dataset_path_list:
        if 'fold' in dataset_path or 'shirt' in dataset_path:
            key = 'fold_shirt'
        elif 'clean_table' in dataset_path and 'pick' not in dataset_path:
            key = 'clean_table'
        else:
            key = 'others'
            base_action = preprocess_base_action(base_action)
        data_tasks_dict[key].append(dataset_path)
    norm_stats_tasks = {k: None for k in data_tasks_dict.keys()}
    for k, v in data_tasks_dict.items():
        if len(v) > 0:
            norm_stats_tasks[k], _ = get_norm_stats(v)
    return norm_stats_tasks

def smooth_base_action(base_action):
    return np.stack([
        np.convolve(base_action[:, i], np.ones(5) / 5, mode='same') for i in range(base_action.shape[1])
    ], axis=-1).astype(np.float32)


def preprocess_base_action(base_action):
    # base_action = calibrate_linear_vel(base_action)
    base_action = smooth_base_action(base_action)

    return base_action


def postprocess_base_action(base_action):
    linear_vel, angular_vel = base_action
    linear_vel *= 1.0
    angular_vel *= 1.0
    # angular_vel = 0
    # if np.abs(linear_vel) < 0.05:
    #     linear_vel = 0
    return np.array([linear_vel, angular_vel])

def compute_dict_mean(epoch_dicts):
    result = {k: None for k in epoch_dicts[0]}
    num_items = len(epoch_dicts)
    for k in result:
        value_sum = 0
        for epoch_dict in epoch_dicts:
            value_sum += epoch_dict[k]
        result[k] = value_sum / num_items
    return result


def detach_dict(d):
    new_d = dict()
    for k, v in d.items():
        new_d[k] = v.detach()
    return new_d


def set_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)