import torch
from glob import glob
from vidar.utils.read import read_image, read_numpy
from vidar.datasets.augmentations.tensor import to_tensor_image, to_tensor
from vidar.geometry.pose_utils import invert_pose

src0 = '/data/cv_datasets/raw/DyCheck/iphone/zips'
dst0 = '/data/cv_datasets/processed/DyCheck'

folders = sorted(glob(f'{src0}/*'))

for f in folders:

    f = f.split('/')[-1]
    # path = '/data/cv_datasets/raw/DyCheck/iphone/zips/backpack-20250915T161627Z-1-001/backpack'
    src = f'{src0}/{f}'
    dst = f'{dst0}/{f}'

    print(src)
    print(dst)

    rgbs = sorted(glob(f'{src}/rgb/1x/*.png'))
    depths = sorted(glob(f'{src}/depth/1x/*.npy'))
    cameras = sorted(glob(f'{src}/camera/*.json'))

    names = {}
    rgbs_dict = {}
    for v in rgbs:
        st = int(v.split('/')[-1][0])
        if st not in rgbs_dict:
            rgbs_dict[st] = []
            names[st] = []
        rgbs_dict[st].append(to_tensor_image(read_image(v)))
        names[st].append(v.split('/')[-1].split('.')[0])
    rgbs = rgbs_dict

    depths_dict = {}
    for v in depths:
        st = int(v.split('/')[-1][0])
        if st not in depths_dict:
            depths_dict[st] = []
        depths_dict[st].append(to_tensor(read_numpy(v)))
    depths = depths_dict

    # rgbs = torch.stack([to_tensor_image(read_image(v)) for v in rgbs], 0)
    # depths = torch.stack([to_tensor(read_numpy(v)) for v in depths], 0).permute(0, 3, 1, 2)

    import json
    pose, intrinsics = {}, {}
    for cam in cameras:
        st = int(cam.split('/')[-1][0])
        with open(cam, 'r') as file:
            data = json.load(file)

        R = torch.tensor(data['orientation']).reshape(3, 3)
        t = torch.tensor(data['position']).reshape(3)
        T = torch.eye(4)
        T[:3, :3] = R
        T[:3, -1] = t
        if st not in pose:
            pose[st] = []
        pose[st].append(T)

        px, py = data['principal_point']
        foc = data['focal_length']
        K = torch.eye(3)
        K[0,0] = K[1,1] = foc
        K[0,2] = px
        K[1,2] = py
        if st not in intrinsics:
            intrinsics[st] = []
        intrinsics[st].append(K)

    for key in rgbs.keys():
        rgbs[key] = torch.stack(rgbs[key], 0)
    for key in depths.keys():
        depths[key] = torch.stack(depths[key], 0).permute(0, 3, 1, 2)
    for key in pose.keys():
        pose[key] = invert_pose(torch.stack(pose[key], 0))
    for key in intrinsics.keys():
        intrinsics[key] = torch.stack(intrinsics[key], 0)

    from vidar.geometry.camera import Camera
    points, camera = {}, {}
    for key in pose.keys():
        camera[key] = Camera(K=intrinsics[key], Tcw=pose[key], hw=(720,960))
        if key == 0:
            points[key] = camera[key].reconstruct_depth_map(depths[key], to_world=True)


    from vidar.utils.write import write_image, write_npz
    for key in rgbs.keys():
        for i in range(rgbs[key].shape[0]):
            write_image(f'{dst}/rgb/{key}/{names[key][i]}.jpg', rgbs[key][i])
    for key in depths.keys():
        for i in range(depths[key].shape[0]):
            write_npz(f'{dst}/depth/{key}/{names[key][i]}.npz', dict(data=depths[key][i, 0]))
    for key in pose.keys():
        for i in range(pose[key].shape[0]):
            write_npz(f'{dst}/lowdim/{key}/{names[key][i]}.npz', dict(
                pose=pose[key][i],
                intrinsics=intrinsics[key][i],
                shape=(720,960),
            ))


# for key in rgbs.keys():


# from camviz import Draw
# draw = Draw((1600, 900))
# draw.add3Dworld('wld')

# cvcams = []
# for key in rgbs.keys():
#     n = rgbs[key].shape[0]
#     for i in range(n):
#         cvcams.append(draw.cvcam(camera[key][i], scale=0.1))
#         if key == 0:
#             draw.addBufferf(f'rgb{key}{i}', rgbs[key][i])
# for key in depths.keys():
#     n = depths[key].shape[0]
#     for i in range(n):
#         if key == 0:
#             draw.addBufferf(f'pts{key}{i}', points[key][i])



# t = 0
# while draw.input():
#     if draw.UP:
#         t += 1
#     if draw.DOWN:
#         t -= 1
#     draw.clear()
#     for cvcam in cvcams:
#         draw['wld'].object(cvcam)
#     for key in rgbs.keys():
#         if key != 0: continue
#         # for i in range(n):
#         #     draw['wld'].points(f'pts{key}{i}',f'rgb{key}{i}')
#         #     draw['wld'].object(cvcams[i])
#         draw['wld'].points(f'pts{key}{t}',f'rgb{key}{t}')
#         draw['wld'].object(cvcams[t], color='red')
#     draw.update(30)