# coding=utf-8
# Copyright 2023 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""RoboNet dataset.

RoboNet: Large-Scale Multi-Robot Learning

Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bernadette Bucher,
Karl Schmeckpeper, Siddharth Singh, Sergey Levine, Chelsea Finn

https://www.robonet.wiki/
"""

import os

import numpy as np
from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
import tensorflow_datasets.public_api as tfds

DATA_URL_SAMPLE = (
    'https://drive.google.com/uc?export=download&'
    'id=1YX2TgT8IKSn9V4wGCwdzbRnS53yicV2P'
)
DATA_URL = (
    'https://drive.google.com/uc?export=download&'
    'id=1BkqHzfRkfzgzCfc73NbNnPMK_rg3i1n9'
)

STATES_DIM = 5
ACTIONS_DIM = 5


class RobonetConfig(tfds.core.BuilderConfig):
  """ "Configuration for RoboNet video rescaling."""

  def __init__(
      self, *, sample_dataset=False, width=None, height=None, **kwargs
  ):
    """The parameters specifying how the dataset will be processed.

    The dataset comes with three separate splits. You can specify which split
    you want in `split_number`. If `width` and `height` are set, the videos
    will be rescaled to have those heights and widths (using ffmpeg).

    Args:
      sample_dataset: Whether or not to use the sample dataset.
      width: An integer with the width or None.
      height: An integer with the height or None.
      **kwargs: Passed on to the constructor of `BuilderConfig`.
    """
    super(RobonetConfig, self).__init__(
        version=tfds.core.Version('4.0.1'), **kwargs
    )
    if (width is None) ^ (height is None):
      raise ValueError('Either both dimensions should be set, or none of them')
    self.sample_dataset = sample_dataset
    self.width = width
    self.height = height


class Builder(tfds.core.BeamBasedBuilder):
  """RoboNet: Large-Scale Multi-Robot Learning."""

  BUILDER_CONFIGS = [
      RobonetConfig(
          name='robonet_sample_64',
          description='64x64 RoboNet Sample.',
          sample_dataset=True,
          width=64,
          height=64,
      ),
      RobonetConfig(
          name='robonet_sample_128',
          description='128x128 RoboNet Sample.',
          sample_dataset=True,
          width=128,
          height=128,
      ),
      RobonetConfig(
          name='robonet_64',
          description='64x64 RoboNet.',
          sample_dataset=False,
          width=64,
          height=64,
      ),
      RobonetConfig(
          name='robonet_128',
          description='128x128 RoboNet.',
          sample_dataset=False,
          width=128,
          height=128,
      ),
  ]

  def _info(self):
    if self.builder_config.width is not None:
      if self.builder_config.height is None:
        raise ValueError('Provide either both height and width or none.')
      ffmpeg_extra_args = (
          '-vf',
          'scale={}x{}'.format(
              self.builder_config.height, self.builder_config.width
          ),
      )
    else:
      ffmpeg_extra_args = []

    video_shape = (
        None,
        self.builder_config.height,
        self.builder_config.width,
        3,
    )

    features = tfds.features.FeaturesDict({
        # Video frames: uint8 [None, Time, Width, Height, Channels]
        'video': tfds.features.Video(
            video_shape,
            ffmpeg_extra_args=ffmpeg_extra_args,
            encoding_format='png',
        ),  # pytype: disable=wrong-arg-types  # gen-stub-imports
        # Robot actions: float32, [None, ACTIONS_DIM]
        'actions': tfds.features.Tensor(
            shape=(None, ACTIONS_DIM), dtype=np.float32
        ),
        # Robot states: float32, [None, STATE_DIM]
        'states': tfds.features.Tensor(
            shape=(None, STATES_DIM), dtype=np.float32
        ),
        # Filename: Text
        'filename': tfds.features.Text(),
    })

    return self.dataset_info_from_configs(
        features=features,
        homepage='https://www.robonet.wiki/',
    )

  def _split_generators(self, dl_manager):
    files = dl_manager.download_and_extract(
        DATA_URL_SAMPLE if self.builder_config.sample_dataset else DATA_URL
    )
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,
            gen_kwargs={
                'filedir': os.path.join(files, 'hdf5'),
            },
        ),
    ]

  def _build_pcollection(self, pipeline, filedir):
    """Generate examples as dicts."""
    beam = tfds.core.lazy_imports.apache_beam

    def _process_example(filename):
      """Converts one video from hdf5 format."""
      h5py = tfds.core.lazy_imports.h5py
      with h5py.File(filename) as hf:
        video_bytes = hf['env']['cam0_video']['frames'][:].tobytes()
        states = hf['env']['state'][:].astype(np.float32)
        states = np.pad(
            states, ((0, 0), (0, STATES_DIM - states.shape[1])), 'constant'
        )
        actions = hf['policy']['actions'][:].astype(np.float32)
        actions = np.pad(
            actions, ((0, 0), (0, ACTIONS_DIM - actions.shape[1])), 'constant'
        )

      basename = os.path.basename(filename)
      features = {
          'video': video_bytes,
          'actions': actions,
          'states': states,
          'filename': basename,
      }
      return basename, features

    filenames = tf.io.gfile.glob(os.path.join(filedir, '*.hdf5'))
    return pipeline | beam.Create(filenames) | beam.Map(_process_example)
