import pickle
from pathlib import Path
import torch
import cv2
import numpy as np
import os
import open3d as o3d
from tqdm import tqdm
from tqdm.contrib.concurrent import process_map, thread_map
from rec_utils.datasets import ARKitDataset, VGGTDataset
from rec_utils.aligner import build_aligner_1p1d
from rec_utils.datasets.arkit.utils import rotate_image

def process_scene(scene_params, target_depth_shape=None):
    images = scene_params["image_files"]
    poses = scene_params["poses"]
    depths = scene_params["depths"]
    Ks = scene_params["Ks"]
    pts3d = scene_params["pts3d"]
    im_confs = scene_params["im_conf"]
    print(scene_params.keys())
    im_shapes = scene_params["imshapes"]
    im_shape = im_shapes[0]

    image_hw = cv2.imread(images[0]).shape[:2]
    image_scale = np.ones((3, 3))
    image_scale[0] *= image_hw[1] / im_shape[1]
    image_scale[1] *= image_hw[0] / im_shape[0]

    if target_depth_shape is None:
        target_depth_shape = image_hw


    depth_scale = np.ones((3, 3))
    depth_scale[0] *= target_depth_shape[1] / im_shape[1]
    depth_scale[1] *= target_depth_shape[0] / im_shape[0]

    data = [
        {
            "image_path": image,
            "pose": pose,
            "depth": cv2.resize(depth.numpy(), target_depth_shape[::-1], interpolation=cv2.INTER_LINEAR),
            "source_K": K,
            "image_K": K * image_scale,
            "depth_K": K * depth_scale,
            "pts3d": pts,
            "im_conf": im_conf,
            "im_shape_target": image_hw,
            "depth_shape_target": target_depth_shape,
            "shape_original": im_shape,
        } for image, pose, depth, K, pts, im_conf in zip(images, poses, depths, Ks, pts3d, im_confs)
    ]

    return data

def es_wrap(data):
    try:
        return export_scene(data)
    except Exception as e:
        print(e)
        return None

def export_scene(data):
    vggt_dataset, arkit_dataset, i, out_dir, processing_args = data
    vggt_scene = vggt_dataset[i]
    scene_id = vggt_scene.id
    arkit_scene = arkit_dataset[scene_id]
    out_path = out_dir / scene_id
    out_path.mkdir(parents=True, exist_ok=True)
    # if os.path.exists(out_path / f'{vggt_scene.id}_vh_clean_2.ply'):
    #     return
    arkit_scene.frames = arkit_scene.frames[-100:]
    aligner = build_aligner_1p1d(source_scene=vggt_scene, target_scene=arkit_scene)
    scene = aligner.align(vggt_scene, inplace=True)
    # scene = arkit_scene
    
    K_color = arkit_scene[0].image_intrinsics
    K_depth = arkit_scene[0].depth_intrinsics
    intrinsics_path = out_path / "intrinsic"
    intrinsics_path.mkdir(parents=True, exist_ok=True)
    color_path = out_path / "color"
    color_path.mkdir(parents=True, exist_ok=True)
    depth_path = out_path / "depth"
    depth_path.mkdir(parents=True, exist_ok=True)
    pose_path = out_path / "pose"
    pose_path.mkdir(parents=True, exist_ok=True)
    np.savetxt(intrinsics_path / "intrinsic_color.txt", K_color)
    np.savetxt(intrinsics_path / "intrinsic_depth.txt", K_depth)

    np.savetxt(intrinsics_path / "extrinsic_color.txt", np.eye(4))
    np.savetxt(intrinsics_path / "extrinsic_depth.txt", np.eye(4))
    tsdffusion = o3d.pipelines.integration.ScalableTSDFVolume(
        voxel_length=0.025,
        sdf_trunc=0.1,
        color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8
    )
    print("rotation angle", arkit_scene.rotation_angle)
    for frame in tqdm(scene.frames):
        # image = rotate_image(frame.image, arkit_scene.rotation_angle)
        # # image = frame.image
        # h, w = image.shape[:2]
        # depth = frame.depth.astype(np.float32)
        # depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_LINEAR)
        # color = o3d.geometry.Image(image)
        np.savetxt(str(pose_path / f'{frame.frame_id}.txt'), frame.pose)
    #     cv2.imwrite(str(color_path / f'{frame.frame_id}.jpg'), image[..., ::-1])
    #     cv2.imwrite(str(depth_path / f'{frame.frame_id}.png'), (depth * 1000.).astype(np.uint16))

    #     fx, fy = K_color[0, 0], K_color[1, 1]
    #     cx, cy = K_color[0, 2], K_color[1, 2]
        
    #     dh, dw = depth.shape
    #     depth_o3d = o3d.geometry.Image(depth)

    #     rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
    #         color, depth_o3d, depth_trunc=10., convert_rgb_to_intensity=False, depth_scale=1.0
    #     )
    #     camera_o3d = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy)
    #     tsdffusion.integrate(
    #         rgbd, camera_o3d,
    #         np.linalg.inv(frame.pose),
    #     )
    # pc = tsdffusion.extract_point_cloud()
    # pc.voxel_down_sample(voxel_size=0.025)
    # o3d.io.write_point_cloud(str(out_path / f'{scene.id}_vh_clean_2.ply'), pc)


    
    
    
    

vggt_dataset = VGGTDataset("/home/jovyan/users/bulat/workspace/3drec/vggt/output/arkit_new/")
arkit_dataset = ARKitDataset("/workspace-SR006.nfs2/datasets/arkitscenes/offline_prepared_data/posed_images/")
processing_args = {
    "voxel_size": 0.025,
    "out_dir": Path("data/arkit_dust3r_posed/processed"),
}


val_path = Path("../") / "OKNO/data/arkitscenes/arkitscenes_offline_infos_val.pkl"
out_dir = Path("data/arkit_vggt/processed")
with open(val_path, "rb") as f:
    data = pickle.load(f)

data_list = data["data_list"]
val_scenes = [scene["lidar_points"]["lidar_path"] for scene in data_list]
def extract_name(item):
    return item.split("_")[0]
val_scenes = [extract_name(scene) for scene in val_scenes]
data = [(vggt_dataset, arkit_dataset, i, out_dir, processing_args) for i in range(len(vggt_dataset))]
thread_map(es_wrap, data, chunksize=128)
    # break