import numpy as np
import tensorflow_datasets as tfds

from simpler_env.utils.visualization import write_video

DATASETS = ["fractal20220817_data", "bridge"]


def dataset2path(dataset_name):
    if dataset_name == "robo_net":
        version = "1.0.0"
    elif dataset_name == "language_table":
        version = "0.0.1"
    else:
        version = "0.1.0"
    return f"gs://gresearch/robotics/{dataset_name}/{version}"


if __name__ == "__main__":
    dataset_name = DATASETS[0]
    dset = tfds.builder_from_directory(builder_dir=dataset2path(dataset_name))

    dset = dset.as_dataset(split="train[:50]", read_config=tfds.ReadConfig(add_tfds_id=True))
    dset = list(dset)
    for i, episode in enumerate(dset):
        gt_images = []
        episode_steps = list(episode["steps"])
        for j in range(len(episode_steps) - 1):
            gt_images.append(episode_steps[j]["observation"]["image"])
        write_video(f"{dataset_name}_vis/{i}_gt.mp4", gt_images, fps=5)

        # from matplotlib import pyplot as plt

        # images = gt_images
        # ACTION_DIM_LABELS = ['x', 'y', 'z', 'yaw', 'pitch', 'roll', 'grasp']

        # img_strip = np.concatenate(np.array(images[::3]), axis=1)

        # # set up plt figure
        # figure_layout = [
        #     ['image'] * len(ACTION_DIM_LABELS),
        #     ACTION_DIM_LABELS
        # ]
        # plt.rcParams.update({'font.size': 12})
        # fig, axs = plt.subplot_mosaic(figure_layout)
        # fig.set_size_inches([45, 10])

        # # plot actions
        # pred_actions = np.array([np.concatenate([episode_step['action']['world_vector'], episode_step['action']['rotation_delta'], episode_step['action']['open_gripper'][None]], axis=-1) for episode_step in episode_steps])
        # for action_dim, action_label in enumerate(ACTION_DIM_LABELS):
        #     # actions have batch, horizon, dim, in this example we just take the first action for simplicity
        #     axs[action_label].plot(pred_actions[:, action_dim], label='predicted action')
        #     axs[action_label].set_title(action_label)
        #     axs[action_label].set_xlabel('Time in one episode')

        # axs['image'].imshow(img_strip)
        # axs['image'].set_xlabel('Time in one episode (subsampled)')
        # plt.legend()
        # plt.show()