import os
import json
import numpy as np
import lzma
import sys
from glob import glob
from tqdm import tqdm

fx = 389.07278
fy = 389.07278
cx = 321.61887
cy = 238.43630

# # path1 = '/data/cv_datasets/processed/LBMv12/BimanualClearKitchenCounter/davis/real/BimanualClearKitchenCounter__davis__real__2025-01-22T11-28-30-08-00__000000/lowdim/scene_left/0000000000.npz'
# # path2 = '/data/processed/HumanoidEveryday/Loco-manipulation__walk_towards_a_desk_and_place_a_cube_on_a_tray__episode_6/lowdim/0/0000000000.npz'

# path1 = '/data/cv_datasets/processed/LBMv12/BimanualClearKitchenCounter/davis/real/BimanualClearKitchenCounter__davis__real__2025-01-22T11-28-30-08-00__000000/depth/scene_left/0000000000.npz'
# path2 = '/data/processed/HumanoidEveryday/Loco-manipulation__walk_towards_a_desk_and_place_a_cube_on_a_tray__episode_6/depth/0/0000000000.npz'

# aa1 = np.load(path1)
# aa2 = np.load(path2)

# print(aa1['data'].shape, aa1['data'].min(), aa1['data'].max())
# print(aa2['data'].shape, aa2['data'].min(), aa2['data'].max())

# # print(aa1['intrinsics'])
# # print(aa2['intrinsics'])

# # print(aa1['prompt'])
# # print(aa2['prompt'])

# import sys
# sys.exit()


src = '/data/cv_datasets/downloads/HumanoidEveryday'
dst = '/data/processed/HumanoidEveryday'


def _load_depth_lzma(depth_lzma_path):
    with open(depth_lzma_path, "rb") as f:
        compressed_data = f.read()
        decompressed = lzma.decompress(compressed_data)
        depth_array = np.frombuffer(decompressed, dtype=np.uint16).reshape(
            (480, 640)
        )
        return depth_array / 1000


folders1 = glob(f'{src}/*')
folders1 = [f for f in folders1 if not f.endswith('zip')]
folders1 = [f for f in folders1 if f.endswith(sys.argv[1])]

for f1 in tqdm(folders1):
    folders2 = glob(f'{f1}/*')
    folders2 = [f for f in folders2 if not f.endswith('zip')]
    for f2 in tqdm(folders2):
        folders3 = sorted(glob(f'{f2}/episode_*'))
        if len(folders3) == 0:
            f2 = sorted(glob(f'{f2}/*'))[0]
            folders3 = sorted(glob(f'{f2}/episode_*'))

        try:
            with open(f'{f2}/metadata/metadata.json', 'r') as file:
                metadata = json.load(file)
        except:
            continue

        for f3 in tqdm(folders3):
            path = dst + '/' + f3.replace(src + '/', '').replace('/', '__').replace(' ', '_').replace('-','_')

            try:
                with open(f'{f3}/data.json', 'r') as file:
                    data = json.load(file)
            except:
                continue

            for i, val in enumerate(data):

                rgb_file = f'{f3}/{val["image"]}'
                lowdim_file = rgb_file.replace('/rgb/', '/lowdim/')

                rgb_dst_val = f'{path}/rgb/0'
                lowdim_dst_val = rgb_dst_val.replace('/rgb/', '/lowdim/')

                os.makedirs(rgb_dst_val, exist_ok=True)
                os.makedirs(lowdim_dst_val, exist_ok=True)

                try:
                    depth_file = f'{f3}/{val["depth"]}'
                    depth = _load_depth_lzma(depth_file)
                    depth_dst_val = f'{path}/depth/0'
                    os.makedirs(depth_dst_val, exist_ok=True)
                    depth_dst_val = f'{depth_dst_val}/%010d.npz' % i
                    np.savez_compressed(depth_dst_val, data=depth)
                except:
                    pass

                rgb_dst_val = f'{rgb_dst_val}/%010d.jpg' % i
                lowdim_dst_val = f'{lowdim_dst_val}/%010d.npz' % i

                command = f'cp "{rgb_file}" {rgb_dst_val}'
                # command = f'cp {rgb_file} ./aaa.jpg'
                os.system(command)

                lowdim = {
                    'prompt': [metadata['description']],
                    'intrinsics': np.array([
                        [ fx, 0.0,  cx],
                        [0.0,  fy,  cy],
                        [0.0, 0.0, 1.0],
                    ])
                }

                np.savez_compressed(lowdim_dst_val, **lowdim)

                # import sys
                # sys.exit()