from pathlib import Path
from omegaconf import DictConfig

from calvin_agent.datasets.disk_dataset import DiskDataset

DATASETS_DIR = Path("../calvin_debug_dataset/training")
OBSERVATION_SPACE = DictConfig({"rgb_obs": ['rgb_static', 'rgb_gripper'],  # conf/datamodule/observation_space/lang_rgb_static_gripper_rel_act.yaml
                    "depth_obs": [],
                    "state_obs": ['robot_obs'],
                    "actions": ['rel_actions'],
                    "language": ['language'],
                    })

PROPRIO_STATE = DictConfig({"n_state_obs": 8,  # conf/datamodule/proprioception_dims/robot_no_joints.yaml
                            "keep_indices": [[0, 7], [14,15]],
                            "robot_orientation_idx": [3, 6],
                            "normalize": True,
                            "normalize_robot_orientation": True,
                            })


def main():
    dataset = DiskDataset(datasets_dir=DATASETS_DIR,
                          obs_space=OBSERVATION_SPACE,
                          proprio_state=PROPRIO_STATE,
                          key="vis",
                          lang_folder="lang_annotations",
                          num_workers=0)

    print(len(dataset))

    for data in dataset:
        print(data)


if __name__ == "__main__":
    main()