from pathlib import Path
import shutil
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import pickle
import cv2
from tqdm.contrib.concurrent import process_map



dust3r_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/DUSt3R/res/scannet_posed_35")
pcds_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/DUSt3R/res_scennet_posed_35_new")

def resize_intrinsics(K, base_shape, new_shape):
    fx = K[0, 0]
    fy = K[1, 1]
    cx = K[0, 2]
    cy = K[1, 2]

    fx_new = fx * new_shape[0] / base_shape[0]
    fy_new = fy * new_shape[1] / base_shape[1]
    cx_new = cx * new_shape[0] / base_shape[0]
    cy_new = cy * new_shape[1] / base_shape[1]
    
    return np.array([[fx_new, 0, cx_new, 0], [0, fy_new, cy_new, 0], [0, 0, 1, 0], [0, 0, 0, 1]])

def prepare_scene(scene_id):

    scannet_processed_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/MaskClustering/data/scannet/processed") / scene_id
    out_path = Path("data/scannet_dust3r_posed_35_bulat/processed") / scene_id
    out_path.mkdir(parents=True, exist_ok=True)

    path_to_scene = dust3r_path / scene_id
    scene_params = torch.load(path_to_scene / 'scene_params.pt')
    pcd_path = pcds_path / (scene_id + '.ply')


    poses = scene_params['poses']
    depths = scene_params['depths']
    Ks = scene_params['Ks']
    image_names = [Path(image_name) for image_name in scene_params['image_files']]
    shutil.copy(pcd_path, out_path / f'{scene_id}_vh_clean_2.ply')


    path_to_out_intrinsic = out_path / 'intrinsic'
    path_to_out_intrinsic.mkdir(parents=True, exist_ok=True)

    shutil.copy(scannet_processed_path / 'intrinsic' / 'intrinsic_color.txt', path_to_out_intrinsic / 'intrinsic_color.txt')
    shutil.copy(scannet_processed_path / 'intrinsic' / 'intrinsic_depth.txt', path_to_out_intrinsic / 'intrinsic_depth.txt')
    shutil.copy(scannet_processed_path / 'intrinsic' / 'extrinsic_color.txt', path_to_out_intrinsic / 'extrinsic_color.txt')
    shutil.copy(scannet_processed_path / 'intrinsic' / 'extrinsic_depth.txt', path_to_out_intrinsic / 'extrinsic_depth.txt')

    # base_intrinsic_path = f'/home/jovyan/users/lemeshko/mmdetection3d/data/scannet/posed_images/{scene_id}/intrinsic.txt'
    # base_intrinsic = np.loadtxt(base_intrinsic_path)

    # np.savetxt(path_to_out_intrinsic / 'intrinsic_color.txt', resize_intrinsics(Ks.mean(0), (512, 384), (1280, 960)))
    # np.savetxt(path_to_out_intrinsic / 'intrinsic_depth.txt', resize_intrinsics(Ks.mean(0), (512, 384), (640, 480)))
    # np.savetxt(path_to_out_intrinsic / 'extrinsic_color.txt', np.eye(4))
    # np.savetxt(path_to_out_intrinsic / 'extrinsic_depth.txt', np.eye(4))

    path_to_out_color = out_path / 'color'
    path_to_out_color.mkdir(parents=True, exist_ok=True)


    path_to_out_depth = out_path / 'depth'
    path_to_out_depth.mkdir(parents=True, exist_ok=True)

    path_to_out_pose = out_path / 'pose'
    path_to_out_pose.mkdir(parents=True, exist_ok=True)


    image_file = Path(image_names[0])
    gt_depth = image_file.parent.parent / 'depth_maps' / (image_file.stem + ".png")
    gt_pose = image_file.parent.parent / 'poses' / (image_file.stem + ".txt")

    print(gt_depth)
    # gt_depth = cv2.imread(str(gt_depth), -1) / 1000.0
    # gt_pose = np.loadtxt(str(gt_pose))


    # h, w = gt_depth.shape
    # depth = depths[0].numpy().astype(np.float32)
    # depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_NEAREST)

    # valid_mask = (gt_depth > 1e-10) & (depth > 1e-10)
    # depth = depth[valid_mask]
    # gt_depth = gt_depth[valid_mask]

    # scale = np.median(gt_depth / depth)
    # first_pose = np.linalg.inv(poses[0].numpy())


    for i, image_name in enumerate(image_names):
        image_name = Path(image_name)
        def get_image_name(image_name):
            return str(int(image_name.stem))
        shutil.copy(Path("/home/jovyan/users/lemeshko/mmdetection3d/data/scannet/posed_images") / scene_id / image_name.name, path_to_out_color / (get_image_name(image_name) + '.jpg'))

        depth_map = depths[i].numpy() * 1000
        depth_map = cv2.resize(depth_map, (640, 480))
        depth_map = depth_map.astype(np.uint16)
        cv2.imwrite(str(path_to_out_depth / (get_image_name(image_name) + '.png')), depth_map)

        pose = poses[i].numpy()
        np.savetxt(str(path_to_out_pose / (get_image_name(image_name) + '.txt')), pose)


if __name__ == "__main__":
    pkl_path = '/home/jovyan/users/lemeshko/mmdetection3d/data/scannetOV20/scannet_oneformer3d_infos_val.pkl'
    with open(pkl_path, 'rb') as file:
        lst = pickle.load(file)['data_list']
        scenes = sorted([val['lidar_points']['lidar_path'][:-4] for val in lst])
    process_map(prepare_scene, scenes, chunksize=10)