#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Preprocessing code for the StaticThings3D dataset
# dataset at https://github.com/lmb-freiburg/robustmvd/blob/master/rmvd/data/README.md#staticthings3d
# 1) Download StaticThings3D in /path/to/StaticThings3D/
#    with the script at https://github.com/lmb-freiburg/robustmvd/blob/master/rmvd/data/scripts/download_staticthings3d.sh
#    --> depths.tar.bz2 frames_finalpass.tar.bz2 poses.tar.bz2 frames_cleanpass.tar.bz2 intrinsics.tar.bz2
# 2) unzip everything in the same /path/to/StaticThings3D/ directory
# 5) python datasets_preprocess/preprocess_staticthings3d.py --StaticThings3D_dir /path/to/tmp/StaticThings3D/
# --------------------------------------------------------
import os
import os.path as osp
import re

import numpy as np
from tqdm import tqdm

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
import path_to_root  # noqa
from dust3r.datasets.utils import cropping  # noqa
from dust3r.utils.parallel import parallel_threads


def get_parser():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--StaticThings3D_dir", required=True)
    parser.add_argument("--precomputed_pairs", required=True)
    parser.add_argument("--output_dir", default="data/staticthings3d_processed")
    return parser


def main(db_root, pairs_path, output_dir):
    all_scenes = _list_all_scenes(db_root)

    # crop images
    args = [
        (db_root, osp.join(split, subsplit, seq), camera, f"{n:04d}", output_dir)
        for split, subsplit, seq in all_scenes
        for camera in ["left", "right"]
        for n in range(6, 16)
    ]
    parallel_threads(load_crop_and_save, args, star_args=True, front_num=1)

    # verify that all images are there
    CAM = {b"l": "left", b"r": "right"}
    pairs = np.load(pairs_path)
    for scene, seq, cam1, im1, cam2, im2 in tqdm(pairs):
        seq_path = osp.join("TRAIN", scene.decode("ascii"), f"{seq:04d}")
        for cam, idx in [(CAM[cam1], im1), (CAM[cam2], im2)]:
            for ext in ["clean", "final"]:
                impath = osp.join(output_dir, seq_path, cam, f"{idx:04n}_{ext}.jpg")
                assert osp.isfile(impath), f"missing an image at {impath=}"

    print(f">> Saved all data to {output_dir}!")


def load_crop_and_save(db_root, relpath_, camera, num, out_dir):
    relpath = osp.join(relpath_, camera, num)
    if osp.isfile(osp.join(out_dir, relpath + ".npz")):
        return
    os.makedirs(osp.join(out_dir, relpath_, camera), exist_ok=True)

    # load everything
    intrinsics_in = readFloat(
        osp.join(db_root, "intrinsics", relpath_, num + ".float3")
    )
    cam2world = np.linalg.inv(
        readFloat(osp.join(db_root, "poses", relpath + ".float3"))
    )
    depthmap_in = readFloat(osp.join(db_root, "depths", relpath + ".float3"))
    img_clean = cv2.cvtColor(
        cv2.imread(
            osp.join(db_root, "frames_cleanpass", relpath + ".png"), cv2.IMREAD_COLOR
        ),
        cv2.COLOR_BGR2RGB,
    )
    img_final = cv2.cvtColor(
        cv2.imread(
            osp.join(db_root, "frames_finalpass", relpath + ".png"), cv2.IMREAD_COLOR
        ),
        cv2.COLOR_BGR2RGB,
    )

    # do the crop
    assert img_clean.shape[:2] == (540, 960)
    assert img_final.shape[:2] == (540, 960)
    (clean_out, final_out), depthmap, intrinsics_out, R_in2out = _crop_image(
        intrinsics_in, (img_clean, img_final), depthmap_in, (512, 384)
    )

    # write everything
    clean_out.save(osp.join(out_dir, relpath + "_clean.jpg"), quality=80)
    final_out.save(osp.join(out_dir, relpath + "_final.jpg"), quality=80)
    cv2.imwrite(osp.join(out_dir, relpath + ".exr"), depthmap)

    # New camera parameters
    cam2world[:3, :3] = cam2world[:3, :3] @ R_in2out.T
    np.savez(
        osp.join(out_dir, relpath + ".npz"),
        intrinsics=intrinsics_out,
        cam2world=cam2world,
    )


def _crop_image(intrinsics_in, color_image_in, depthmap_in, resolution_out=(512, 512)):
    image, depthmap, intrinsics_out = cropping.rescale_image_depthmap(
        color_image_in, depthmap_in, intrinsics_in, resolution_out
    )
    R_in2out = np.eye(3)
    return image, depthmap, intrinsics_out, R_in2out


def _list_all_scenes(path):
    print(">> Listing all scenes")

    res = []
    for split in ["TRAIN"]:
        for subsplit in "ABC":
            for seq in os.listdir(osp.join(path, "intrinsics", split, subsplit)):
                res.append((split, subsplit, seq))
    print(f"   (found ({len(res)}) scenes)")
    assert res, f"Did not find anything at {path=}"
    return res


def readFloat(name):
    with open(name, "rb") as f:
        if (f.readline().decode("utf-8")) != "float\n":
            raise Exception("float file %s did not contain <float> keyword" % name)

        dim = int(f.readline())

        dims = []
        count = 1
        for i in range(0, dim):
            d = int(f.readline())
            dims.append(d)
            count *= d

        dims = list(reversed(dims))
        data = np.fromfile(f, np.float32, count).reshape(dims)
    return data  # Hxw or CxHxW NxCxHxW


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args.StaticThings3D_dir, args.precomputed_pairs, args.output_dir)