# Path: weakly_supervised_control/experiment_utils.py
# def load_experiment(experiment_dir: str, use_cpu: bool = False):
#     # Load config file.
#     with gfile.GFile(os.path.join(experiment_dir, 'variant.json'), 'r') as f:
#         variant = json.load(f)
#     _replace_paths(variant)
# 
#     # Load policy, VAE, and environment.
#     with gfile.GFile(os.path.join(experiment_dir, 'params.pkl'), 'rb') as f:
#         if use_cpu:
#             params = torch.load(f, map_location=lambda storage, loc: storage)
#         else:
#             params = torch.load(f)
# 
#     with gfile.GFile(os.path.join(experiment_dir, 'vae.pkl'), 'rb') as f:
#         vae = pickle.load(f)
# 
#     return variant, params, vae
# 
# def process_image_trajectory(
#         images,
#         object_color: Tuple[float, float, float] = (0, 0.392, 0.784),
#         threshold: float = 1e-2,
#         draw_color: Tuple[float, float, float] = (1, 1, 1),
#         alpha: float = 0.6,
#         draw_arrow: bool = False,
#         arrow_length: float = 3,
#         arrow_width: float = 3):
#     """Draws a line over the tracked object_color along the image trajectory."""
#     object_color = np.array(object_color)
#     lower = np.maximum(0.0, object_color - threshold)
#     upper = np.minimum(1.0, object_color + threshold)
# 
#     centroids = []
#     for image in images:
#         image = image.reshape((3, 48, 48)).transpose((2, 1, 0))
#         mask = cv2.inRange(image, lower, upper)
#         centroid = np.mean(list(zip(*np.nonzero(mask))), axis=0)
#         centroids.append(centroid)
# 
#     draw_indices = set()
# 
#     def _to_index(p: Tuple[float, float]) -> Tuple[int, int]:
#         assert p.shape == (2,), p
#         return tuple(map(int, p))
# 
#     # Walk along the trajectory and collect pixels to draw.
#     pos = centroids[0].copy()
#     draw_indices.add(_to_index(pos))
#     directions = []
#     dest_index = 1
#     while dest_index < len(centroids):
#         dest_pos = centroids[dest_index]
#         dest_vec = dest_pos - pos
#         dest_distance = np.linalg.norm(dest_vec)
#         if dest_distance < 1e-2:
#             dest_index += 1
#             continue
#         dest_direction = dest_vec / dest_distance
#         step = np.minimum(dest_distance, 0.5)
#         pos += dest_direction * step
#         draw_indices.add(_to_index(pos))
#         directions.append(dest_direction)
# 
#     if draw_arrow and directions:
#         average_dir = np.mean(directions[-10:], axis=0)
#         average_dir /= np.linalg.norm(average_dir)
#         assert average_dir.shape == (2,)
# 
#         def _add_line(src, dest, iters: int = 20):
#             assert src.shape == (2,) and dest.shape == (2,), (src, dest)
#             for i in range(iters):
#                 a = i / (iters - 1)
#                 v = (1. - a) * src + a * dest
#                 draw_indices.add(_to_index(v))
# 
#         # Add an arrow head.
#         tri_base = centroids[-1] - average_dir * arrow_length
#         ortho_vec = np.array(
#             [average_dir[1], -average_dir[0]]) * arrow_width * 0.5
#         tri_right = tri_base + ortho_vec
#         tri_left = tri_base - ortho_vec
#         # _add_line(tri_left, tri_right)
#         _add_line(tri_left, centroids[-1])
#         _add_line(tri_right, centroids[-1])
# 
#     # Draw on the last image.
#     result = image.copy()
#     draw_color = np.array(draw_color)
#     for index in draw_indices:
#         if any(i < 0 or i >= result.shape[0] for i in index):
#             continue
#         result[index] = (1 - alpha) * result[index] + alpha * draw_color
# 
#     return result


from typing import Tuple
from tensorflow.io import gfile
from rlkit.samplers.rollout_functions import rollout
from weakly_supervised_control.experiment_utils import load_experiment, process_image_trajectory
        from scipy.stats import pearsonr
        from weak_disentangle.train_utils import load_dset
import click
import cv2
import json
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import os
import torch
import rlkit.torch.pytorch_util as ptu
            goals.append(g)
        goals = np.array(goals)

    fig = plt.figure(figsize=(6, 6))
    gs = gridspec.GridSpec(
        grid_size, grid_size,
        left=0.1, right=1, bottom=0.1, top=1, wspace=0.01, hspace=0.01)

    for i in range(goals.shape[0]):
        for j in range(goals.shape[1]):
            goal = goals[i, j]
            policy.reset()
            o = env.reset()

            images = []
            path_length = 0
            while path_length < max_path_length:
                new_obs = np.hstack((o['latent_observation'], goal))
                action = policy.get_action(new_obs)
                next_o, reward, done, env_info = env.step(action[0])
                images.append(next_o['image_observation'])

                if render:
                    env.render()
                if done:
                    break

                o = next_o
                path_length += 1


# Path: weakly_supervised_control/experiment_utils.py
# def load_experiment(experiment_dir: str, use_cpu: bool = False):
#     # Load config file.
#     with gfile.GFile(os.path.join(experiment_dir, 'variant.json'), 'r') as f:
#         variant = json.load(f)
#     _replace_paths(variant)
# 
#     # Load policy, VAE, and environment.
#     with gfile.GFile(os.path.join(experiment_dir, 'params.pkl'), 'rb') as f:
#         if use_cpu:
#             params = torch.load(f, map_location=lambda storage, loc: storage)
#         else:
#             params = torch.load(f)
# 
#     with gfile.GFile(os.path.join(experiment_dir, 'vae.pkl'), 'rb') as f:
#         vae = pickle.load(f)
# 
#     return variant, params, vae
# 
# def process_image_trajectory(
#         images,
#         object_color: Tuple[float, float, float] = (0, 0.392, 0.784),
#         threshold: float = 1e-2,
#         draw_color: Tuple[float, float, float] = (1, 1, 1),
#         alpha: float = 0.6,
#         draw_arrow: bool = False,
#         arrow_length: float = 3,
#         arrow_width: float = 3):
#     """Draws a line over the tracked object_color along the image trajectory."""
#     object_color = np.array(object_color)
#     lower = np.maximum(0.0, object_color - threshold)
#     upper = np.minimum(1.0, object_color + threshold)
# 
#     centroids = []
#     for image in images:
#         image = image.reshape((3, 48, 48)).transpose((2, 1, 0))
#         mask = cv2.inRange(image, lower, upper)
#         centroid = np.mean(list(zip(*np.nonzero(mask))), axis=0)
#         centroids.append(centroid)
# 
#     draw_indices = set()
# 
#     def _to_index(p: Tuple[float, float]) -> Tuple[int, int]:
#         assert p.shape == (2,), p
#         return tuple(map(int, p))
# 
#     # Walk along the trajectory and collect pixels to draw.
#     pos = centroids[0].copy()
#     draw_indices.add(_to_index(pos))
#     directions = []
#     dest_index = 1
#     while dest_index < len(centroids):
#         dest_pos = centroids[dest_index]
#         dest_vec = dest_pos - pos
#         dest_distance = np.linalg.norm(dest_vec)
#         if dest_distance < 1e-2:
#             dest_index += 1
#             continue
#         dest_direction = dest_vec / dest_distance
#         step = np.minimum(dest_distance, 0.5)
#         pos += dest_direction * step
#         draw_indices.add(_to_index(pos))
#         directions.append(dest_direction)
# 
#     if draw_arrow and directions:
#         average_dir = np.mean(directions[-10:], axis=0)
#         average_dir /= np.linalg.norm(average_dir)
#         assert average_dir.shape == (2,)
# 
#         def _add_line(src, dest, iters: int = 20):
#             assert src.shape == (2,) and dest.shape == (2,), (src, dest)
#             for i in range(iters):
#                 a = i / (iters - 1)
#                 v = (1. - a) * src + a * dest
#                 draw_indices.add(_to_index(v))
# 
#         # Add an arrow head.
#         tri_base = centroids[-1] - average_dir * arrow_length
#         ortho_vec = np.array(
#             [average_dir[1], -average_dir[0]]) * arrow_width * 0.5
#         tri_right = tri_base + ortho_vec
#         tri_left = tri_base - ortho_vec
#         # _add_line(tri_left, tri_right)
#         _add_line(tri_left, centroids[-1])
#         _add_line(tri_right, centroids[-1])
# 
#     # Draw on the last image.
#     result = image.copy()
#     draw_color = np.array(draw_color)
#     for index in draw_indices:
#         if any(i < 0 or i >= result.shape[0] for i in index):
#             continue
#         result[index] = (1 - alpha) * result[index] + alpha * draw_color
# 
#     return result


# Path: scripts/eval_wsc.py
from typing import Tuple
from tensorflow.io import gfile
from rlkit.samplers.rollout_functions import rollout
from weakly_supervised_control.experiment_utils import load_experiment, process_image_trajectory
        from scipy.stats import pearsonr
        from weak_disentangle.train_utils import load_dset
import click
import cv2
import json
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import os
import torch
import rlkit.torch.pytorch_util as ptu

            goals.append(g)
        goals = np.array(goals)

    fig = plt.figure(figsize=(6, 6))
    gs = gridspec.GridSpec(
        grid_size, grid_size,
        left=0.1, right=1, bottom=0.1, top=1, wspace=0.01, hspace=0.01)

    for i in range(goals.shape[0]):
        for j in range(goals.shape[1]):
            goal = goals[i, j]
            policy.reset()
            o = env.reset()

            images = []
            path_length = 0
            while path_length < max_path_length:
                new_obs = np.hstack((o['latent_observation'], goal))
                action = policy.get_action(new_obs)
                next_o, reward, done, env_info = env.step(action[0])
                images.append(next_o['image_observation'])

                if render:
                    env.render()
                if done:
                    break

                o = next_o
                path_length += 1


            image = process_image_trajectory(images)