from glob import glob
import numpy as np 
import torch 

from vidar.utils.read import read_image, read_numpy
from externals.vidar.vidar.datasets.augmentations.tensor import to_tensor_image, to_tensor


def read_txt(path):
    with open(path, 'r') as file:
        data = file.readlines()
    for i in range(len(data)):
        data[i] = data[i].replace('\n', '').split(' ')
        data[i] = [float(v) for v in data[i]]
    out = to_tensor(np.array(data))
    if out.shape[0] == 4: 
        new = torch.eye(4)
        new[:3, :3] = np.transpose(out[1:, :3])
        new[:3, -1] = out[:1, :3]
        out = new
    return out


src = '/data/cv_datasets/raw/NVIDIA_DynamicScenes/extracted'
dst = '/data/cv_datasets/processed/DynamicScenes'


paths = sorted(glob(f'{src}/*'))
for path in paths:

    print(path)

    name = path.split('/')[-1]
    dst_i = f'{dst}/{name}'

    sequences = glob(f'{path}/multiview_GT/*')
    sequences = sorted([s.split('/')[-1] for s in sequences])

    rgbs = []
    depths = []

    for seq in sequences:
        rgb = sorted(glob(f'{path}/multiview_GT/{seq}/*'))
        rgb = [to_tensor_image(read_image(v)) for v in rgb]
        rgbs.append(torch.stack(rgb, 0))

    depth = sorted(glob(f'{path}/depth_GT/*'))
    depth = [to_tensor(read_numpy(v)) for v in depth]
    depth = torch.stack(depth, 0).unsqueeze(1)

    intrinsics = sorted(glob(f'{path}/calibration/*/intrinsic.txt'))
    intrinsics = torch.stack([read_txt(v) for v in intrinsics], 0)

    extrinsics = sorted(glob(f'{path}/calibration/*/extrinsic.txt'))
    extrinsics = torch.stack([read_txt(v) for v in extrinsics], 0)

    from vidar.utils.write import write_image, write_npz
    from vidar.geometry.pose_utils import invert_pose

    extrinsics = invert_pose(extrinsics)

    for j in range(extrinsics.shape[0]):
        write_npz(f'{dst_i}/depth/{j}/{sequences[0]}.npz', dict(data=depth[j][0]))
        for i in range(len(rgbs)):
            write_npz(f'{dst_i}/depth/{j}/{sequences[i]}.npz', dict(data=depth[j][0]))
            write_image(f'{dst_i}/rgb/{j}/{sequences[i]}.jpg', rgbs[i][j])
            write_npz(f'{dst_i}/lowdim/{j}/{sequences[i]}.npz', dict(
                shape=(rgbs[i][j].shape[-2], rgbs[i][j].shape[-1]), 
                intrinsics=intrinsics[j], 
                pose=extrinsics[j],
            ))

import sys
sys.exit()
aa = torch.tensor([[[0,0,1]]], dtype=extrinsics.dtype)
print(extrinsics.shape, aa.shape)
extrinsics = torch.cat([extrinsics, aa], 2)

from vidar.geometry.camera import Camera

cam = Camera(K=intrinsics, Twc=extrinsics, hw=rgbs[0])
pts = cam.reconstruct_depth_map(depths, to_world=True)

from camviz import Draw
draw = Draw((1600, 900))
draw.add3Dworld('wld')
cvcams = [draw.cvcam(cam, i) for i in range(12)]
for i in range(pts.shape[0]):
    draw.addBufferf(f'pts{i}', pts[i])
    draw.addBufferf(f'rgb{i}', rgbs[0][i])

while draw.input():
    draw.clear()
    for c in cvcams:
        draw['wld'].object(c)
    for i in range(12):
        draw['wld'].size(2).points(f'pts{i}', f'rgb{i}')
    draw.update(30)