# Path: ravens/utils/utils.py
# def get_heightmap(points, colors, bounds, pixel_size):
# def get_pointcloud(depth, intrinsics):
# def transform_pointcloud(points, transform):
# def reconstruct_heightmaps(color, depth, configs, bounds, pixel_size):
# def pix_to_xyz(pixel, height, bounds, pixel_size, skip_height=False):
# def xyz_to_pix(position, bounds, pixel_size):
# def unproject_vectorized(uv_coordinates, depth_values,
#                          intrinsic,
#                          distortion):
# def unproject_depth_vectorized(im_depth, depth_dist,
#                                camera_mtx,
#                                camera_dist):
# def sample_distribution(prob, n_samples=1):
# def invert(pose):
# def multiply(pose0, pose1):
# def apply(pose, position):
# def eulerXYZ_to_quatXYZW(rotation):  # pylint: disable=invalid-name
# def quatXYZW_to_eulerXYZ(quaternion_xyzw):  # pylint: disable=invalid-name
# def apply_transform(transform_to_from, points_from):
# def preprocess(img):
# def get_fused_heightmap(obs, configs, bounds, pix_size):
# def get_image_transform(theta, trans, pivot=(0, 0)):
# def check_transform(image, pixel, transform):
# def get_se3_from_image_transform(theta, trans, pivot, heightmap, bounds,
#                                  pixel_size):
# def get_random_image_transform_params(image_size):
# def perturb(input_image, pixels, set_theta_zero=False):
# def plot(fname,  # pylint: disable=dangerous-default-value
#          title,
#          ylabel,
#          xlabel,
#          data,
#          xlim=[-np.inf, 0],
#          xticks=None,
#          ylim=[np.inf, -np.inf],
#          show_std=True):
# def create_visualizer(clear=True):
# def make_frame(vis, name, h, radius, o=1.0):
# def meshcat_visualize(vis, obs, act, info):
# COLORS = {
#     'blue': [078.0 / 255.0, 121.0 / 255.0, 167.0 / 255.0],
#     'red': [255.0 / 255.0, 087.0 / 255.0, 089.0 / 255.0],
#     'green': [089.0 / 255.0, 169.0 / 255.0, 079.0 / 255.0],
#     'orange': [242.0 / 255.0, 142.0 / 255.0, 043.0 / 255.0],
#     'yellow': [237.0 / 255.0, 201.0 / 255.0, 072.0 / 255.0],
#     'purple': [176.0 / 255.0, 122.0 / 255.0, 161.0 / 255.0],
#     'pink': [255.0 / 255.0, 157.0 / 255.0, 167.0 / 255.0],
#     'cyan': [118.0 / 255.0, 183.0 / 255.0, 178.0 / 255.0],
#     'brown': [156.0 / 255.0, 117.0 / 255.0, 095.0 / 255.0],
#     'gray': [186.0 / 255.0, 176.0 / 255.0, 172.0 / 255.0]
# }


import os
import pickle
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from ravens.utils import utils
  name = f'{FLAGS.task}-{FLAGS.agent}-{FLAGS.n_demos}-'
  print(name)

  # Load and print results to console.
  path = FLAGS.root_dir
  curve = []
  for fname in tf.io.gfile.listdir(path):
    fname = os.path.join(path, fname)
    if name in fname and '.pkl' in fname:
      n_steps = int(fname[(fname.rfind('-') + 1):-4])
      data = pickle.load(open(fname, 'rb'))
      rewards = []
      for reward, _ in data:
        rewards.append(reward)
      score = np.mean(rewards)
      std = np.std(rewards)
      curve.append((n_steps, score, std))
  curve.sort()
  for log in curve:
    print(f'  {log[0]} steps:\t{log[1]:.2f}%\t± {log[2]:.2f}%')

  # Plot results over training steps.
  title = f'{FLAGS.agent} on {FLAGS.task} w/ {FLAGS.n_demos} demos'
  ylabel = 'Testing Task Success (%)'
  xlabel = '# of Training Steps'
  if FLAGS.disp:
    logs = {}
    curve = np.array(curve)
    logs[name] = (curve[:, 0], curve[:, 1], curve[:, 2])
    fname = os.path.join(path, f'{name}-plot.png')

# Path: ravens/utils/utils.py
# def get_heightmap(points, colors, bounds, pixel_size):
# def get_pointcloud(depth, intrinsics):
# def transform_pointcloud(points, transform):
# def reconstruct_heightmaps(color, depth, configs, bounds, pixel_size):
# def pix_to_xyz(pixel, height, bounds, pixel_size, skip_height=False):
# def xyz_to_pix(position, bounds, pixel_size):
# def unproject_vectorized(uv_coordinates, depth_values,
#                          intrinsic,
#                          distortion):
# def unproject_depth_vectorized(im_depth, depth_dist,
#                                camera_mtx,
#                                camera_dist):
# def sample_distribution(prob, n_samples=1):
# def invert(pose):
# def multiply(pose0, pose1):
# def apply(pose, position):
# def eulerXYZ_to_quatXYZW(rotation):  # pylint: disable=invalid-name
# def quatXYZW_to_eulerXYZ(quaternion_xyzw):  # pylint: disable=invalid-name
# def apply_transform(transform_to_from, points_from):
# def preprocess(img):
# def get_fused_heightmap(obs, configs, bounds, pix_size):
# def get_image_transform(theta, trans, pivot=(0, 0)):
# def check_transform(image, pixel, transform):
# def get_se3_from_image_transform(theta, trans, pivot, heightmap, bounds,
#                                  pixel_size):
# def get_random_image_transform_params(image_size):
# def perturb(input_image, pixels, set_theta_zero=False):
# def plot(fname,  # pylint: disable=dangerous-default-value
#          title,
#          ylabel,
#          xlabel,
#          data,
#          xlim=[-np.inf, 0],
#          xticks=None,
#          ylim=[np.inf, -np.inf],
#          show_std=True):
# def create_visualizer(clear=True):
# def make_frame(vis, name, h, radius, o=1.0):
# def meshcat_visualize(vis, obs, act, info):
# COLORS = {
#     'blue': [078.0 / 255.0, 121.0 / 255.0, 167.0 / 255.0],
#     'red': [255.0 / 255.0, 087.0 / 255.0, 089.0 / 255.0],
#     'green': [089.0 / 255.0, 169.0 / 255.0, 079.0 / 255.0],
#     'orange': [242.0 / 255.0, 142.0 / 255.0, 043.0 / 255.0],
#     'yellow': [237.0 / 255.0, 201.0 / 255.0, 072.0 / 255.0],
#     'purple': [176.0 / 255.0, 122.0 / 255.0, 161.0 / 255.0],
#     'pink': [255.0 / 255.0, 157.0 / 255.0, 167.0 / 255.0],
#     'cyan': [118.0 / 255.0, 183.0 / 255.0, 178.0 / 255.0],
#     'brown': [156.0 / 255.0, 117.0 / 255.0, 095.0 / 255.0],
#     'gray': [186.0 / 255.0, 176.0 / 255.0, 172.0 / 255.0]
# }


# Path: ravens/plot.py
import os
import pickle
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from ravens.utils import utils

  name = f'{FLAGS.task}-{FLAGS.agent}-{FLAGS.n_demos}-'
  print(name)

  # Load and print results to console.
  path = FLAGS.root_dir
  curve = []
  for fname in tf.io.gfile.listdir(path):
    fname = os.path.join(path, fname)
    if name in fname and '.pkl' in fname:
      n_steps = int(fname[(fname.rfind('-') + 1):-4])
      data = pickle.load(open(fname, 'rb'))
      rewards = []
      for reward, _ in data:
        rewards.append(reward)
      score = np.mean(rewards)
      std = np.std(rewards)
      curve.append((n_steps, score, std))
  curve.sort()
  for log in curve:
    print(f'  {log[0]} steps:\t{log[1]:.2f}%\t± {log[2]:.2f}%')

  # Plot results over training steps.
  title = f'{FLAGS.agent} on {FLAGS.task} w/ {FLAGS.n_demos} demos'
  ylabel = 'Testing Task Success (%)'
  xlabel = '# of Training Steps'
  if FLAGS.disp:
    logs = {}
    curve = np.array(curve)
    logs[name] = (curve[:, 0], curve[:, 1], curve[:, 2])
    fname = os.path.join(path, f'{name}-plot.png')

    utils.plot(fname, title, ylabel, xlabel, data=logs, ylim=[0, 1])