"""
Modify AssemblyHands dataset by splitting long episodes into shorter overlapping segments while in Processed format.
This is done as webdatasetisation was failing for very long episodes (more than 7k frames), but splitting them into smaller episodes of ~3.6k frames with 100 frame overlap works well.
"""

# DGX input path: /home/vguizilini/workspace/data/predict2/data/cv_datasets/processed
# corresponding input path inside docker: /workspace/cv_datasets/processed/AssemblyHands/

# Output path DGX: /home/vguizilini/workspace/data/predict2/data/cv_datasets/processed/AssemblyHands_V2/
# corresponding output path inside docker: /workspace/cv_datasets/processed/AssemblyHands_V2/

# Script runs inside docker container
# Usage: python3 assembly_modify_parallelized.py

local_path = "/workspace/cv_datasets/processed/AssemblyHands/"
write_path = "/workspace/cv_datasets/processed/AssemblyHands_V2/"

"""
Dataset structure:
AssemblyHands/
episode_{}/
    rgb/
        frame_{}.png
        ...
    lowdim/
        frame_{}.png
        ...
"""

import os
import shutil
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import threading

camera_names = ['camera_0', 'camera_1', 'camera_2', 'camera_3', 'camera_4', 'camera_5', 'camera_6', 'camera_7']
split_length = 3600  # number of frames per split
overlap = 100  # number of overlapping frames between splits

# Thread-safe counter for episode numbering
class EpisodeCounter:
    def __init__(self, value=0):
        self._value = value
        self._lock = threading.Lock()
    
    def get_next(self):
        with self._lock:
            current = self._value
            self._value += 1
            return current

def copy_camera_frame(args):
    """Copy a single camera frame (both RGB and lowdim)"""
    camera, i, camera_rgb_path, camera_lowdim_path, new_rgb_cam_path, new_lowdim_cam_path, rgb_frames, lowdim_frames = args
    
    # Ensure directories exist
    os.makedirs(new_rgb_cam_path, exist_ok=True)
    os.makedirs(new_lowdim_cam_path, exist_ok=True)
    
    # Copy frames
    shutil.copy(os.path.join(camera_rgb_path, rgb_frames[i]), os.path.join(new_rgb_cam_path, rgb_frames[i]))
    shutil.copy(os.path.join(camera_lowdim_path, lowdim_frames[i]), os.path.join(new_lowdim_cam_path, lowdim_frames[i]))

def process_segment(args):
    """Process a single segment (split) of frames"""
    episode_path, rgb_path, lowdim_path, frames_per_camera, start_frame, end_frame, episode_counter, write_path = args
    
    # Get episode number
    new_ep_num = episode_counter.get_next()
    
    # Create new episode directory and rgb/lowdim subdirs
    new_episode_name = f'episode_{new_ep_num:04d}'
    new_episode_path = os.path.join(write_path, new_episode_name)
    os.makedirs(new_episode_path, exist_ok=True)
    new_rgb_path = os.path.join(new_episode_path, 'rgb')
    new_lowdim_path = os.path.join(new_episode_path, 'lowdim')
    os.makedirs(new_rgb_path, exist_ok=True)
    os.makedirs(new_lowdim_path, exist_ok=True)
    
    # Prepare arguments for parallel frame copying
    copy_args = []
    for i in range(start_frame, end_frame):
        for camera in camera_names:
            camera_rgb_path = os.path.join(rgb_path, camera)
            camera_lowdim_path = os.path.join(lowdim_path, camera)
            new_rgb_cam_path = os.path.join(new_rgb_path, camera)
            new_lowdim_cam_path = os.path.join(new_lowdim_path, camera)
            rgb_frames, lowdim_frames = frames_per_camera[camera]
            
            copy_args.append((camera, i, camera_rgb_path, camera_lowdim_path, 
                            new_rgb_cam_path, new_lowdim_cam_path, rgb_frames, lowdim_frames))
    
    # Parallel copy frames using ThreadPoolExecutor (I/O bound operations)
    with ThreadPoolExecutor(max_workers=min(32, len(copy_args))) as executor:
        executor.map(copy_camera_frame, copy_args)
    
    print(f'Modified data saved in {new_episode_path}')
    return new_episode_name

def process_episode(args):
    """Process a single episode"""
    episode, local_path, write_path, episode_counter = args
    
    print(f'Processing {episode}...')
    episode_path = os.path.join(local_path, episode)
    rgb_path = os.path.join(episode_path, 'rgb')
    lowdim_path = os.path.join(episode_path, 'lowdim')
    
    len_frames = 0
    frames_per_camera = {}
    
    # Load frame information for all cameras
    for camera in camera_names:
        camera_rgb_path = os.path.join(rgb_path, camera)
        camera_lowdim_path = os.path.join(lowdim_path, camera)
        rgb_frames = sorted([f for f in os.listdir(camera_rgb_path) if f.endswith('.jpg')])
        if len_frames == 0:
            len_frames = len(rgb_frames)
        lowdim_frames = sorted([f for f in os.listdir(camera_lowdim_path) if f.endswith('.npz')])
        frames_per_camera[camera] = (rgb_frames, lowdim_frames)
    
    # Prepare segments for parallel processing
    segment_args = []
    start_frame = 0
    curr_ep_count = 0
    
    while start_frame < len_frames:
        end_frame = min(start_frame + split_length, len_frames)
        segment_args.append((episode_path, rgb_path, lowdim_path, frames_per_camera, 
                           start_frame, end_frame, episode_counter, write_path))
        start_frame += (split_length - overlap)
        curr_ep_count += 1
    
    # Process segments in parallel using ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=min(4, len(segment_args))) as executor:
        results = list(executor.map(process_segment, segment_args))
    
    print(f'Finished processing {episode} -> generated {curr_ep_count} new episodes. Results: {len(results)}')
    return curr_ep_count

def modify_assemblyhands_dataset(local_path, write_path):
    episodes = [d for d in os.listdir(local_path) if d.startswith('episode_')]
    episodes = sorted(episodes)

    exsiting_eps = [d for d in os.listdir(write_path) if d.startswith('episode_')]
    if exsiting_eps:
        print(f'Warning: Found {len(exsiting_eps)} existing episodes in {write_path}. Starting episode numbering from {len(exsiting_eps)}.')
    
    # print (episodes)
    # episodes = episodes[0:1]  # for testing, process only the first episode
    print(f'Found {len(episodes)} episodes in the dataset.')

    # Thread-safe episode counter
    episode_counter = EpisodeCounter(len(exsiting_eps) if exsiting_eps else 0)
    
    # Prepare arguments for parallel episode processing
    episode_args = [(episode, local_path, write_path, episode_counter) for episode in episodes]
    
    # Process episodes in parallel using ProcessPoolExecutor for CPU-bound work
    # Use ThreadPoolExecutor if you prefer threads (better for I/O bound operations)
    total_new_episodes = 0
    with ThreadPoolExecutor(max_workers=min(len(episodes), 2)) as executor:
        results = list(executor.map(process_episode, episode_args))
        total_new_episodes = sum(results)
    
    return total_new_episodes

if __name__ == "__main__":
    new_ep_count = modify_assemblyhands_dataset(local_path, write_path)
    print(f'Total new episodes created: {new_ep_count}')

    # upload to S3 command (run separately):
    # aws s3 cp --recursive /workspace/cv_datasets/processed/AssemblyHands_V2/ s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/processed/AssemblyHands_V2/