# Path: rl_teacher/envs.py
# def make_with_torque_removed(env_id):
#     if '-v' in env_id:
#         env_id = env_id[:env_id.index('-v')].lower()
#     if env_id.startswith('short'):
#         env_id = env_id[len('short'):]
#         short = True
#     else:
#         short = False
#     return task_by_name(env_id, short)  # Use our task_by_name function to get the env
# 
# Path: rl_teacher/segment_sampling.py
# def segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds, workers):
#     """ Generate a list of path segments by doing random rollouts. Can use multiple processes. """
#     if workers < 2:  # Default to basic segment collection
#         return basic_segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds)
# 
#     pool = Pool(processes=workers)
#     segments_per_worker = int(math.ceil(n_desired_segments / workers))
#     # One job per worker. Only the first worker is verbose.
#     jobs = [
#         (env_id, make_env, segments_per_worker, clip_length_in_seconds, i, i == 0, workers)
#         for i in range(workers)]
#     results = pool.starmap(basic_segments_from_rand_rollout, jobs)
#     pool.close()
#     return [segment for sublist in results for segment in sublist]
# 
# Path: rl_teacher/teach.py
# CLIP_LENGTH = 1.5
# 
# Path: rl_teacher/video.py
# def write_segment_to_video(segment, fname, env):
#     os.makedirs(osp.dirname(fname), exist_ok=True)
#     frames = [env.render_full_obs(x) for x in segment["human_obs"]]
#     for i in range(int(env.fps * 0.2)):
#         frames.append(frames[-1])
#     export_video(frames, fname, fps=env.fps)


import os.path as osp
import uuid
from rl_teacher.envs import make_with_torque_removed
from rl_teacher.segment_sampling import segments_from_rand_rollout
from rl_teacher.teach import CLIP_LENGTH
from rl_teacher.video import write_segment_to_video


TEST_RENDER_DIR = '/tmp/rl_teacher_media_test'

def test_render_videos():
    env_id = "Hopper-v1"
    env = make_with_torque_removed(env_id)

# Path: rl_teacher/envs.py
# def make_with_torque_removed(env_id):
#     if '-v' in env_id:
#         env_id = env_id[:env_id.index('-v')].lower()
#     if env_id.startswith('short'):
#         env_id = env_id[len('short'):]
#         short = True
#     else:
#         short = False
#     return task_by_name(env_id, short)  # Use our task_by_name function to get the env
# 
# Path: rl_teacher/segment_sampling.py
# def segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds, workers):
#     """ Generate a list of path segments by doing random rollouts. Can use multiple processes. """
#     if workers < 2:  # Default to basic segment collection
#         return basic_segments_from_rand_rollout(env_id, make_env, n_desired_segments, clip_length_in_seconds)
# 
#     pool = Pool(processes=workers)
#     segments_per_worker = int(math.ceil(n_desired_segments / workers))
#     # One job per worker. Only the first worker is verbose.
#     jobs = [
#         (env_id, make_env, segments_per_worker, clip_length_in_seconds, i, i == 0, workers)
#         for i in range(workers)]
#     results = pool.starmap(basic_segments_from_rand_rollout, jobs)
#     pool.close()
#     return [segment for sublist in results for segment in sublist]
# 
# Path: rl_teacher/teach.py
# CLIP_LENGTH = 1.5
# 
# Path: rl_teacher/video.py
# def write_segment_to_video(segment, fname, env):
#     os.makedirs(osp.dirname(fname), exist_ok=True)
#     frames = [env.render_full_obs(x) for x in segment["human_obs"]]
#     for i in range(int(env.fps * 0.2)):
#         frames.append(frames[-1])
#     export_video(frames, fname, fps=env.fps)


# Path: rl_teacher/tests/video_render_test.py
import os.path as osp
import uuid
from rl_teacher.envs import make_with_torque_removed
from rl_teacher.segment_sampling import segments_from_rand_rollout
from rl_teacher.teach import CLIP_LENGTH
from rl_teacher.video import write_segment_to_video



TEST_RENDER_DIR = '/tmp/rl_teacher_media_test'

def test_render_videos():
    env_id = "Hopper-v1"
    env = make_with_torque_removed(env_id)

    segments = segments_from_rand_rollout(env_id, make_with_torque_removed,