        )
        image_pil, _, intrinsics2 = cropping.crop_image_depthmap(
            image_pil, None, intrinsics, crop_bbox
        )

        # convert back to normalized CHW tensor on original device
        image_arr = np.array(image_pil)
        if image_arr.ndim == 2:
            image_arr = np.repeat(image_arr[..., None], 3, axis=2)
        image_tensor = ImgNorm(image_pil)  # CHW, [-1, 1]
        if had_batch_dim:
            image_tensor = image_tensor.unsqueeze(0)
        if img_device is not None:
            image_tensor = image_tensor.to(img_device)

        # return intrinsics as torch tensor with batch dim like upstream expects
        intrinsics_tensor = torch.from_numpy(intrinsics2).unsqueeze(0)

        return image_tensor, intrinsics_tensor

    def prepare_input(
        img_paths,
        img_mask,
        msk_paths,
        size,
        raymaps=None,
        raymap_mask=None,
        revisit=1,
        update=True,
        crop=True,
        load_func=None,
        annots=None,
        sampled_indices=None,
        reset_interval=100,
        crop_res=None
    ):
        images = load_images(img_paths, size=size, verbose=False, crop=crop)
        images = load_func((img_paths, images, annots, sampled_indices))

        has_msk = len(msk_paths) > 0
        if has_msk:
            msks = load_masks(msk_paths, size=size, verbose=False, crop=crop)

        views = []
        if raymaps is None and raymap_mask is None:
            num_views = len(images)

            for i in range(num_views):
                view = {
                    "img": images[i]["img"],
                    "ray_map": torch.full(
                        (
                            images[i]["img"].shape[0],
                            6,
                            images[i]["img"].shape[-2],
                            images[i]["img"].shape[-1],
                        ),
                        torch.nan,
                    ),
                    "true_shape": torch.from_numpy(images[i]["true_shape"]),
                    "idx": i,
                    "instance": str(i),
                    "camera_pose": torch.from_numpy(
                        images[i]["camera_pose"]
                    ).unsqueeze(0),
                    "camera_intrinsics": torch.from_numpy(
                        images[i]["intrinsics"]
                    ).unsqueeze(0),
                    "img_mask": torch.tensor(True).unsqueeze(0),
                    "ray_mask": torch.tensor(False).unsqueeze(0),
                    "update": torch.tensor(True).unsqueeze(0),
                    "reset": torch.tensor((i+1) % reset_interval == 0).unsqueeze(0),
                }

                if crop_res is not None:
                    view["img"], view["camera_intrinsics"] = _crop_resize(
                        images[i]["img"], images[i]["intrinsics"], crop_res)
                    # update true_shape to reflect cropped image size (H, W)
                    view["true_shape"] = torch.tensor([view["img"].shape[-2:]], dtype=torch.int32)

                if has_msk:
                    view["msk"] = msks[i]
                for key in images[i].keys():
                    if key.startswith(("smpl", "T_w2c")):
                        view[key] = torch.tensor(images[i][key]).unsqueeze(0)
                views.append(view)
                if (i+1) % reset_interval == 0:
                    overlap_view = deepcopy(view)
                    overlap_view["reset"] = torch.tensor(False).unsqueeze(0)
                    views.append(overlap_view)

        else:

            num_views = len(images) + len(raymaps)
            assert len(img_mask) == len(raymap_mask) == num_views
            assert sum(img_mask) == len(images) and sum(raymap_mask) == len(raymaps)

            j = 0
            k = 0
            for i in range(num_views):
                view = {
                    "img": (
                        images[j]["img"]
                        if img_mask[i]
                        else torch.full_like(images[0]["img"], torch.nan)
                    ),
                    "ray_map": (
                        raymaps[k]
                        if raymap_mask[i]
                        else torch.full_like(raymaps[0], torch.nan)
                    ),
                    "true_shape": (
                        torch.from_numpy(images[j]["true_shape"])
                        if img_mask[i]
                        else torch.from_numpy(np.int32([raymaps[k].shape[1:-1][::-1]]))
                    ),
                    "idx": i,
                    "instance": str(i),
                    "camera_pose": torch.from_numpy(
                        images[i]["camera_pose"]
                    ).unsqueeze(0),
                    "camera_intrinsics": torch.from_numpy(
                        images[i]["intrinsics"]
                    ).unsqueeze(0),
                    "img_mask": torch.tensor(img_mask[i]).unsqueeze(0),
                    "ray_mask": torch.tensor(raymap_mask[i]).unsqueeze(0),
                    "update": torch.tensor(img_mask[i]).unsqueeze(0),
                    "reset": torch.tensor((i+1) % reset_interval == 0).unsqueeze(0),
                }

                if crop_res is not None:
                    view["img"], view["camera_intrinsics"] = _crop_resize(
                        images[i]["img"], images[i]["intrinsics"], crop_res)
                    # update true_shape to reflect cropped image size (H, W)
                    view["true_shape"] = torch.tensor([view["img"].shape[-2:]], dtype=torch.int32)

                if has_msk:
                    view["msk"] = (
                        msks[j]
                        if img_mask[i]
                        else torch.full_like(images[0]["img"], torch.nan)
                    )
                for key in images[i].keys():
                    if key.startswith(("smpl", "T_w2c")):
                        view[key] = torch.tensor(images[i][key]).unsqueeze(0)

                if img_mask[i]:
                    j += 1
                if raymap_mask[i]:
                    k += 1
                views.append(view)
                if (i+1) % reset_interval == 0:
                    overlap_view = deepcopy(view)
                    overlap_view["reset"] = torch.tensor(False).unsqueeze(0)
                    views.append(overlap_view)
            assert j == len(images) and k == len(raymaps)

        if revisit > 1:
            # repeat input for 'revisit' times
            new_views = []
            for r in range(revisit):
                for i in range(len(views)):
                    new_view = deepcopy(views[i])
                    new_view["idx"] = r * len(views) + i
                    new_view["instance"] = str(r * len(views) + i)
                    if r > 0:
                        if not update:
                            new_view["update"] = torch.tensor(False).unsqueeze(0)
                    new_views.append(new_view)
            return new_views

        del images
        return views

    def prepare_gt(gts, revisit=1):
        target_out = defaultdict(list)

        valid_length = len(gts) // revisit
        gts = gts[-valid_length:]

        # delet overlaps: reset_mask=True
        gts = [gt for gt in gts if not gt["reset"]]

        intrinsics = [gt["camera_intrinsics"] for gt in gts]
        K_mhmr = [gt["K_mhmr"] for gt in gts]
        camera_pose = [gt["camera_pose"] for gt in gts]
        imgs = [gt["img"] for gt in gts]
        target_out['K'] = torch.cat(intrinsics, 0)
        target_out['K_mhmr'] = torch.cat(K_mhmr, 0)
        target_out['T_c2w'] = torch.cat(camera_pose, 0)
        target_out['img'] = torch.cat(imgs, 0)
        if 'T_w2c' in gts[0]:
            T_w2c_list = [gt["T_w2c"] for gt in gts]


                    views.append(overlap_view)
            assert j == len(images) and k == len(raymaps)

        if revisit > 1:
            # repeat input for 'revisit' times
            new_views = []
            for r in range(revisit):
                for i in range(len(views)):
                    new_view = deepcopy(views[i])
                    new_view["idx"] = r * len(views) + i
                    new_view["instance"] = str(r * len(views) + i)
                    if r > 0:
                        if not update:
                            new_view["update"] = torch.tensor(False).unsqueeze(0)
                    new_views.append(new_view)
            return new_views

        del images
        return views

    def prepare_gt(gts, revisit=1):
        target_out = defaultdict(list)

        valid_length = len(gts) // revisit
        gts = gts[-valid_length:]

        # delet overlaps: reset_mask=True
        gts = [gt for gt in gts if not gt["reset"]]

        intrinsics = [gt["camera_intrinsics"] for gt in gts]
        K_mhmr = [gt["K_mhmr"] for gt in gts]
        camera_pose = [gt["camera_pose"] for gt in gts]
        imgs = [gt["img"] for gt in gts]
        target_out['K'] = torch.cat(intrinsics, 0)
        target_out['K_mhmr'] = torch.cat(K_mhmr, 0)
        target_out['T_c2w'] = torch.cat(camera_pose, 0)
        target_out['img'] = torch.cat(imgs, 0)
        if 'T_w2c' in gts[0]:
            T_w2c_list = [gt["T_w2c"] for gt in gts]
            target_out['T_w2c'] = torch.cat(T_w2c_list, 0)

        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        for gt, smpl_mask in zip(gts, smpl_mask_list):
            target_out['v3d_c'].append(gt["smpl_v3d_c"][smpl_mask])
            target_out['j3d_c'].append(gt["smpl_j3d_c"][smpl_mask])
            target_out['v3d_w'].append(gt["smpl_v3d_w"][smpl_mask])
            target_out['j3d_w'].append(gt["smpl_j3d_w"][smpl_mask])
            target_out['v2d'].append(gt["smpl_v2d"][smpl_mask])
            target_out['j2d'].append(gt["smpl_j2d"][smpl_mask])

        del gts
        return target_out

    def prepare_output(outputs, revisit=1, solve_pose=False, is_save=False):
        pred_out = defaultdict(list)

        valid_length = len(outputs["pred"]) // revisit
        outputs["pred"] = outputs["pred"][-valid_length:]
        outputs["views"] = outputs["views"][-valid_length:]

        # delet overlaps: reset_mask=True outputs["pred"] and outputs["views"]
        reset_mask = torch.cat([view["reset"] for view in outputs["views"]], 0)
        shifted_reset_mask = torch.cat([torch.tensor(False).unsqueeze(0), reset_mask[:-1]], dim=0)
        outputs["pred"] = [
            pred for pred, mask in zip(outputs["pred"], shifted_reset_mask) if not mask]
        reset_mask = reset_mask[~shifted_reset_mask]

        if solve_pose:
            pts3ds_self_to_save = [
                output["pts3d_in_self_view"] for output in outputs["pred"]
            ]
            pts3ds_other = [


        target_out['K'] = torch.cat(intrinsics, 0)
        target_out['K_mhmr'] = torch.cat(K_mhmr, 0)
        target_out['T_c2w'] = torch.cat(camera_pose, 0)
        target_out['img'] = torch.cat(imgs, 0)
        if 'T_w2c' in gts[0]:
            T_w2c_list = [gt["T_w2c"] for gt in gts]
            target_out['T_w2c'] = torch.cat(T_w2c_list, 0)

        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        for gt, smpl_mask in zip(gts, smpl_mask_list):
            target_out['v3d_c'].append(gt["smpl_v3d_c"][smpl_mask])
            target_out['j3d_c'].append(gt["smpl_j3d_c"][smpl_mask])
            target_out['v3d_w'].append(gt["smpl_v3d_w"][smpl_mask])
            target_out['j3d_w'].append(gt["smpl_j3d_w"][smpl_mask])
            target_out['v2d'].append(gt["smpl_v2d"][smpl_mask])
            target_out['j2d'].append(gt["smpl_j2d"][smpl_mask])

        del gts
        return target_out

    def prepare_output(outputs, revisit=1, solve_pose=False, is_save=False):
        pred_out = defaultdict(list)

        valid_length = len(outputs["pred"]) // revisit
        outputs["pred"] = outputs["pred"][-valid_length:]
        outputs["views"] = outputs["views"][-valid_length:]

        # delet overlaps: reset_mask=True outputs["pred"] and outputs["views"]
        reset_mask = torch.cat([view["reset"] for view in outputs["views"]], 0)
        shifted_reset_mask = torch.cat([torch.tensor(False).unsqueeze(0), reset_mask[:-1]], dim=0)
        outputs["pred"] = [
            pred for pred, mask in zip(outputs["pred"], shifted_reset_mask) if not mask]
        reset_mask = reset_mask[~shifted_reset_mask]

        if solve_pose:
            pts3ds_self_to_save = [
                output["pts3d_in_self_view"] for output in outputs["pred"]
            ]
            pts3ds_other = [
                output["pts3d_in_other_view"] for output in outputs["pred"]
            ]
            conf_self = [output["conf_self"] for output in outputs["pred"]]
            conf_other = [output["conf"] for output in outputs["pred"]]
            pr_poses, focal, pp = recover_cam_params(
                torch.cat(pts3ds_self_to_save, 0),
                torch.cat(pts3ds_other, 0),
                torch.cat(conf_self, 0),
                torch.cat(conf_other, 0),
            )
            pts3ds_self = torch.cat(pts3ds_self_to_save, 0)
        else:
            pts3ds_self_to_save = [
                output["pts3d_in_self_view"] for output in outputs["pred"]
            ]
            pts3ds_other = [
                output["pts3d_in_other_view"] for output in outputs["pred"]
            ]
            conf_self = [output["conf_self"] for output in outputs["pred"]]
            conf_other = [output["conf"] for output in outputs["pred"]]
            pts3ds_self = torch.cat(pts3ds_self_to_save, 0)
            pr_poses = [
                pose_encoding_to_camera(pred["camera_pose"].clone())
                for pred in outputs["pred"]
            ]
            pr_poses = torch.cat(pr_poses, 0)

            B, H, W, _ = pts3ds_self.shape
            pp = (
                torch.tensor([W // 2, H // 2], device=pts3ds_self.device)
                .float()
                .repeat(B, 1)
                .reshape(B, 2)
            )
            focal = estimate_focal_knowing_depth(
                pts3ds_self, pp, focal_mode="weiszfeld"
            )

        if is_save:
            has_mask = "msk" in outputs["pred"][0]
            if has_mask:
                msks = [output["msk"][...,0] for output in outputs["pred"]]
                msks = [unpad_image(m, [H, W]) for m in msks]
            else:
                msks = [torch.zeros(1, H, W) for _ in range(B)]

            pred_out["pts3d_self"] = pts3ds_self_to_save
            pred_out["conf_self"] = conf_self
            pred_out["msk"] = msks

        if reset_mask.any():
            identity = torch.eye(4, device=pr_poses.device)
            reset_poses = torch.where(reset_mask.unsqueeze(-1).unsqueeze(-1), pr_poses, identity)
            cumulative_bases = matrix_cumprod(reset_poses)
            shifted_bases = torch.cat([identity.unsqueeze(0), cumulative_bases[:-1]], dim=0)
            pr_poses = torch.einsum('bij,bjk->bik', shifted_bases, pr_poses)
        pred_out['T_c2w'] = pr_poses

        intrinsics = torch.eye(3, device=pp.device).unsqueeze(0).repeat(B, 1, 1)
        intrinsics[:, 0, 0] = focal  # fx
        intrinsics[:, 1, 1] = focal  # fy
        intrinsics[:, [0, 1], 2] = pp
        pred_out['K'] = intrinsics

        # get SMPL parameters from outputs
        pred_out['shape'] = [output.get(
            "smpl_shape", torch.empty(1,0,10))[0] for output in outputs["pred"]]
        pred_out['rotvec'] = [roma.rotmat_to_rotvec(output.get(
                "smpl_rotmat", torch.empty(1,0,53,3,3))[0]) for output in outputs["pred"]]
        pred_out['transl'] = [output.get(
            "smpl_transl", torch.empty(1,0,3))[0] for output in outputs["pred"]]
        pred_out['expression'] = [output.get(
            "smpl_expression", [None])[0] for output in outputs["pred"]]
        pred_out['loc'] = [output.get(
            "smpl_loc", torch.empty(1,0,2))[0] for output in outputs["pred"]]

        del outputs
        return pred_out

    model = ARCroco3DStereo.from_pretrained(args.weights)
    # SMPL model for gt
    smpl_model = SMPLModel(
        "cpu",
        model_args={
            'patch_size': model.croco_args['patch_size'],
            'mhmr_img_res': model.mhmr_img_res,
            'bb_patch_size': model.bb_patch_size
        },
        eval_args={
            'dataset': args.eval_dataset,
            'use_fake_K': args.use_fake_K
        }
    )
    # SMPL layer for pred
    smpl_layer = SMPL_Layer(type='smplx',
                            gender='neutral',
                            num_betas=10,


#!/usr/bin/env python3
"""
Modified from CUT3R [https://github.com/CUT3R/CUT3R].

Process Bedlam scenes by computing camera intrinsics and extrinsics
from extracted data. The script reads per-scene CSV and image/depth files,
computes the necessary camera parameters, and saves the resulting camera
files (as .npz files) in an output directory.
Note: CUT3R filtered out HDRI scenes and closeup scenes.
We also filter out the sequences without SMPLX annotations following Multi-HMR.

Usage:
    python preprocess_bedlam.py --root /path/to/extracted_data \
                             --outdir /path/to/processed_bedlam \
                             --annot_dir /path/to/bedlam/processed_labels
"""

import os
import sys
import cv2
import numpy as np
import pandas as pd
from glob import glob
import shutil
import OpenEXR  # Ensure OpenEXR is installed
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import argparse
import pickle

# Enable OpenEXR support in OpenCV.
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

# Global constants
IMG_FORMAT = ".png"
rotate_flag = False
SENSOR_W = 36
SENSOR_H = 20.25
IMG_W = 1280
IMG_H = 720

test_list = [
    "20221018_1_250_batch01hand_zoom_suburb_b",
    "20221018_3_250_batch01hand_orbit_archVizUI3_time15",
    "20221019_3-8_250_highbmihand_orbit_stadium",
    # "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
]

# -----------------------------------------------------------------------------
# Helper functions for camera parameter conversion
# -----------------------------------------------------------------------------

def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
    return focal_pixel

def get_cam_int(fl, sens_w, sens_h, cx, cy):
    flx = focalLength_mm2px(fl, sens_w, cx)
    fly = focalLength_mm2px(fl, sens_h, cy)
    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat

def convert_translation_to_opencv(x, y, z):
    t_cv = np.array([y, -z, x])
    return t_cv

def rotation_matrix_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    # Yaw (left-handed)
    R_yaw = np.array(
        [
            [np.cos(-yaw_rad), -np.sin(-yaw_rad), 0],
            [np.sin(-yaw_rad), np.cos(-yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    # Pitch (right-handed)
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, np.sin(pitch_rad)],
            [0, 1, 0],
            [-np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    # Roll (right-handed)
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), -np.sin(roll_rad)],
            [0, np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_roll @ R_pitch @ R_yaw
    return R_unreal

def convert_rotation_to_opencv(R_unreal):
    # Transformation matrix from Unreal to OpenCV coordinate system.
    C = np.array([[0, 1, 0], [0, 0, -1], [1, 0, 0]])
    R_cv = C @ R_unreal @ C.T
    return R_cv

def get_rot_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    R_yaw = np.array(
        [
            [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
            [np.sin(yaw_rad), np.cos(yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, -np.sin(pitch_rad)],
            [0, 1, 0],
            [np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), np.sin(roll_rad)],
            [0, -np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_yaw @ R_pitch @ R_roll
    return R_unreal

def get_extrinsics_unreal(R_unreal, t_unreal):
    cam_trans = np.array(t_unreal)
    ext = np.eye(4)
    ext[:3, :3] = R_unreal
    ext[:3, 3] = cam_trans.reshape(1, 3)
    return ext

def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
    R_unreal = get_rot_unreal(yaw, pitch, roll)
    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
    T_opencv2unreal = np.array(
        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    T_wu2ou = np.array(
        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    return np.linalg.inv(T_opencv2unreal @ T_u2wu @ T_wu2ou)

# -----------------------------------------------------------------------------
# Get camera parameters from the extracted images and CSV data.
# -----------------------------------------------------------------------------

def get_params(
    image_folder,
    fl,
    trans_body,
    cam_x,
    cam_y,
    cam_z,
    fps,
    cam_pitch_,
    cam_roll_,
    cam_yaw_,
):
    all_images = sorted(glob(os.path.join(image_folder, "*" + IMG_FORMAT)))
    imgnames, cam_ext, cam_int = [], [], []

    for img_ind, image_path in enumerate(all_images):
        # Process every 5th frame.
        if img_ind % 5 != 0:
            continue
        cam_ind = img_ind

        cam_pitch_ind = cam_pitch_[cam_ind]
        cam_yaw_ind = cam_yaw_[cam_ind]
        cam_roll_ind = cam_roll_[cam_ind]

        CAM_INT = get_cam_int(fl[cam_ind], SENSOR_W, SENSOR_H, IMG_W / 2.0, IMG_H / 2.0)

        rot_unreal = rotation_matrix_unreal(cam_yaw_ind, cam_pitch_ind, cam_roll_ind)
        rot_cv = convert_rotation_to_opencv(rot_unreal)
        trans_cv = convert_translation_to_opencv(
            cam_x[cam_ind] / 100.0, cam_y[cam_ind] / 100.0, cam_z[cam_ind] / 100.0
        )
        cam_ext_ = np.eye(4)
        cam_ext_[:3, :3] = rot_cv
        # The camera pose is computed as the inverse of the transformed translation.
        cam_ext_[:3, 3] = -rot_cv @ trans_cv

        imgnames.append(
            os.path.join(image_path.split("/")[-2], image_path.split("/")[-1])
        )
        cam_ext.append(cam_ext_)  # camera_pose: c2w
        cam_int.append(CAM_INT)
    return imgnames, cam_ext, cam_int

# -----------------------------------------------------------------------------
# Processing per sequence.
# -----------------------------------------------------------------------------

def process_seq(args):
    """
    Process a single sequence task. For each image, load the corresponding
    depth and image files, and save the computed camera intrinsics and the inverse
    of the extrinsic matrix (i.e. the camera pose in world coordinates) as an NPZ file.
    """
    (
        scene,
        seq_name,
        outdir,
        image_folder_base,
        depth_folder_base,
        mask_folder_base,
        imgnames,
        cam_ext,
        cam_int,
        humans,
        imgname_array,
    ) = args

    split = 'Test' if scene in test_list else 'Training'

    out_rgb_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'rgb')
    out_depth_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'depth')
    out_cam_dir = os.path.join(outdir, split, "_".join([scene, seq_name]), "cam")
    out_smpl_dir = os.path.join(outdir, split, "_".join([scene, seq_name]), "smpl")
    out_mask_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'mask')
    os.makedirs(out_rgb_dir, exist_ok=True)
    os.makedirs(out_depth_dir, exist_ok=True)
    os.makedirs(out_cam_dir, exist_ok=True)
    os.makedirs(out_smpl_dir, exist_ok=True)
    os.makedirs(out_mask_dir, exist_ok=True)

    assert (
        len(imgnames) == len(cam_ext) == len(cam_int)
    ), f"Inconsistent lengths for {scene}_{seq_name}"

    invalid_images = []
    for imgname, ext, intr, human in zip(imgnames, cam_ext, cam_int, humans):
        if imgname not in imgname_array:
            invalid_images.append(f"Invalid image: {scene}/{imgname}")
            continue

        depthname = imgname.replace(".png", "_depth.exr")
        maskname = imgname.replace(".png", "_env.png")
        imgpath = os.path.join(image_folder_base, imgname)
        depthpath = os.path.join(depth_folder_base, depthname)
        maskpath = os.path.join(mask_folder_base, maskname)
        depth= OpenEXR.File(depthpath).parts[0].channels['Depth'].pixels
        depth = depth.astype(np.float32)/100.0

        mask = cv2.imread(maskpath, cv2.IMREAD_GRAYSCALE)
        if mask is None:
            # if mask file not exist, create a zero mask
            mask = np.zeros((IMG_H, IMG_W), dtype=np.uint8)
        else:
            # invert the mask
            mask = 255 - mask

        outimg_path = os.path.join(out_rgb_dir, os.path.basename(imgpath))
        outmask_path = os.path.join(out_mask_dir, os.path.basename(imgpath))
        outdepth_path = os.path.join(out_depth_dir, os.path.basename(imgpath).replace('.png','.npy'))
        outcam_path = os.path.join(
            out_cam_dir, os.path.basename(imgpath).replace(".png", ".npz")
        )

        out_smpl_path = os.path.join(out_smpl_dir, os.path.basename(imgpath).replace(".png", ".pkl"))

        shutil.copy(imgpath, outimg_path)
        cv2.imwrite(outmask_path, mask)
        np.save(outdepth_path, depth)
        np.savez(outcam_path, intrinsics=intr, pose=np.linalg.inv(ext))  # pose: w2c
        with open(out_smpl_path, 'wb') as f:
            pickle.dump(human, f, protocol=pickle.HIGHEST_PROTOCOL)
    return invalid_images

# -----------------------------------------------------------------------------
# Main entry point.
# -----------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Process Bedlam scenes: compute camera intrinsics and extrinsics, "
        "and save processed camera files."
    )
    parser.add_argument(
        "--root",
        type=str,
        required=True,
        help="Root directory of the extracted data (scenes).",
    )
    parser.add_argument(
        "--outdir", type=str, required=True, help="Output directory for processed data."
    )
    parser.add_argument(
        "--annot_dir", type=str, required=True, help="Annotation directory."
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=None,
        help="Number of worker processes (default: os.cpu_count()//2).",
    )
    args = parser.parse_args()

    root = args.root
    outdir = args.outdir
    annot_dir = args.annot_dir
    num_workers = (
        args.num_workers if args.num_workers is not None else (os.cpu_count() or 4) // 2
    )

    invalid_list = []

    # Get scene directories from the root folder.
    scenes = sorted(
        [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
    )
    # Exclude HDRI scenes.
    hdri_scenes = [
        "20221010_3_1000_batch01hand",
        "20221017_3_1000_batch01hand",
        "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
        "20221019_3_250_highbmihand",
    ]
    scenes = np.setdiff1d(scenes, hdri_scenes)

    max_human = 0

    tasks = []
    for scene in tqdm(scenes, desc="Collecting tasks"):
        # Skip closeup scenes.
        if "closeup" in scene:
            continue
        base_folder = os.path.join(root, scene)
        image_folder_base = os.path.join(root, scene, "png")
        depth_folder_base = os.path.join(root, scene, "depth")
        mask_folder_base = os.path.join(root, scene, "masks")
        csv_path = os.path.join(base_folder, "be_seq.csv")
        annot_path = os.path.join(annot_dir, scene + "_6fps.npz")
        if not os.path.exists(annot_path):
            annot_path = os.path.join(annot_dir, scene + "_30fps.npz")
        if not os.path.exists(csv_path):
            continue
        csv_data = pd.read_csv(csv_path)
        csv_data = csv_data.to_dict("list")
        cam_csv_base = os.path.join(base_folder, "ground_truth", "camera")
        annot_x = np.load(annot_path)

        # Retrieving SMPL parameters once
        pose_cam_array = annot_x['pose_cam']
        H_array = annot_x['cam_ext']
        shape_array = annot_x['shape']
        imgname_array = annot_x['imgname']
        trans_cam_array = annot_x['trans_cam']

        seq_count = 0
        max_seq_per_scene = 1500000000000

        # Look for a row in the CSV with a "sequence_name" comment.
        for idx, comment in enumerate(csv_data.get("Comment", [])):
            if "sequence_name" in comment:
                if seq_count >= max_seq_per_scene:
                    break

                seq_name = comment.split(";")[0].split("=")[-1]
                if not np.any(np.char.startswith(imgname_array, seq_name + '/')):
                    invalid_list.append(f"Invalid sequence: {scene}/{seq_name}")
                    continue

                cam_csv_path = os.path.join(cam_csv_base, seq_name + "_camera.csv")
                if not os.path.exists(cam_csv_path):
                    continue
                cam_csv_data = pd.read_csv(cam_csv_path)
                cam_csv_data = cam_csv_data.to_dict("list")
                cam_x = cam_csv_data["x"]
                cam_y = cam_csv_data["y"]
                cam_z = cam_csv_data["z"]
                cam_yaw_ = cam_csv_data["yaw"]
                cam_pitch_ = cam_csv_data["pitch"]
                cam_roll_ = cam_csv_data["roll"]
                fl = cam_csv_data["focal_length"]
                image_folder = os.path.join(image_folder_base, seq_name)
                trans_body = None  # Not used here.
                imgnames, cam_ext, cam_int = get_params(
                    image_folder,
                    fl,
                    trans_body,
                    cam_x,
                    cam_y,
                    cam_z,
                    6,
                    cam_pitch_=cam_pitch_,
                    cam_roll_=cam_roll_,
                    cam_yaw_=cam_yaw_,
                )
                humans = [] # humans for each sequence
                for imgname in imgnames:
                    idxs = np.where(imgname == imgname_array)[0]

                    if len(idxs) > max_human:
                        max_human = len(idxs)

                    persons_per_img = [] # persons for each image
                    for i in idxs:
                        sys.stdout.flush()

                        # SMPLX params
                        pose = pose_cam_array[i]
                        root_pose = pose[:3]
                        body_pose=pose[3:66]
                        jaw_pose=pose[66:69]
                        leye_pose=pose[69:72]
                        reye_pose=pose[72:75]
                        left_hand_pose=pose[75:120]
                        right_hand_pose=pose[120:165]
                        betas=shape_array[i]
                        transl = trans_cam_array[i] + H_array[i][:, 3][:3]

                        person = {
                            # SMPL GT in camera coordinates system
                            'smplx_root_pose': root_pose.reshape(1,3), # axis-angle
                            'smplx_body_pose': body_pose.reshape(21,3), # axis-angle
                            'smplx_jaw_pose': jaw_pose.reshape(1,3), # axis-angle
                            'smplx_leye_pose': leye_pose.reshape(1,3), # axis-angle
                            'smplx_reye_pose': reye_pose.reshape(1,3), # axis-angle
                            'smplx_left_hand_pose': left_hand_pose.reshape(15,3), # axis-angle
                            'smplx_right_hand_pose': right_hand_pose.reshape(15,3), # axis-angle
                            'smplx_shape': betas.reshape(11),
                            'smplx_gender': 'neutral',
                            'smplx_transl': transl.reshape(3),
                        }
                        persons_per_img.append(person)
                    humans.append(persons_per_img)

                tasks.append(
                    (
                        scene,
                        seq_name,
                        outdir,
                        image_folder_base,
                        depth_folder_base,
                        mask_folder_base,
                        imgnames,
                        cam_ext,
                        cam_int,
                        humans,
                        imgname_array,
                    )
                )

                seq_count += 1

    print(f"max_human: {max_human}")

    # Process each task in parallel.
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = {executor.submit(process_seq, task): task for task in tasks}
        for future in tqdm(
            as_completed(futures), total=len(futures), desc="Processing sequences"
        ):
            # error = future.result()
            # if error:
            #     print(error)
            invalid_images = future.result()
            if invalid_images:
                invalid_list.extend(invalid_images)

    # print invalid items
    if invalid_list:
        print("\nInvalid sequences and images:")
        for item in invalid_list:
            print(item)

if __name__ == "__main__":
    main()


# Global constants
IMG_FORMAT = ".png"
rotate_flag = False
SENSOR_W = 36
SENSOR_H = 20.25
IMG_W = 1280
IMG_H = 720

test_list = [
    "20221018_1_250_batch01hand_zoom_suburb_b",
    "20221018_3_250_batch01hand_orbit_archVizUI3_time15",
    "20221019_3-8_250_highbmihand_orbit_stadium",
    # "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
]

# -----------------------------------------------------------------------------
# Helper functions for camera parameter conversion
# -----------------------------------------------------------------------------

def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
    return focal_pixel

def get_cam_int(fl, sens_w, sens_h, cx, cy):
    flx = focalLength_mm2px(fl, sens_w, cx)
    fly = focalLength_mm2px(fl, sens_h, cy)
    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))


IMG_W = 1280
IMG_H = 720

test_list = [
    "20221018_1_250_batch01hand_zoom_suburb_b",
    "20221018_3_250_batch01hand_orbit_archVizUI3_time15",
    "20221019_3-8_250_highbmihand_orbit_stadium",
    # "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
]

# -----------------------------------------------------------------------------
# Helper functions for camera parameter conversion
# -----------------------------------------------------------------------------

def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
    return focal_pixel

def get_cam_int(fl, sens_w, sens_h, cx, cy):
    flx = focalLength_mm2px(fl, sens_w, cx)
    fly = focalLength_mm2px(fl, sens_h, cy)
    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):


    # "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
]

# -----------------------------------------------------------------------------
# Helper functions for camera parameter conversion
# -----------------------------------------------------------------------------

def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
    return focal_pixel

def get_cam_int(fl, sens_w, sens_h, cx, cy):
    flx = focalLength_mm2px(fl, sens_w, cx)
    fly = focalLength_mm2px(fl, sens_h, cy)
    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):


def focalLength_mm2px(focalLength, dslr_sens, focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint * 2
    return focal_pixel

def get_cam_int(fl, sens_w, sens_h, cx, cy):
    flx = focalLength_mm2px(fl, sens_w, cx)
    fly = focalLength_mm2px(fl, sens_h, cy)
    cam_mat = np.array([[flx, 0, cx], [0, fly, cy], [0, 0, 1]])
    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat


    return cam_mat

def unreal2cv2(points):
    # Permute coordinates: x --> y, y --> z, z --> x
    points = np.roll(points, 2, axis=1)
    # Invert the y-axis
    points = points * np.array([1.0, -1.0, 1.0])
    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat

def convert_translation_to_opencv(x, y, z):
    t_cv = np.array([y, -z, x])
    return t_cv

def rotation_matrix_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)


    return points

def get_cam_trans(body_trans, cam_trans):
    cam_trans = np.array(cam_trans) / 100
    cam_trans = unreal2cv2(np.reshape(cam_trans, (1, 3)))
    body_trans = np.array(body_trans) / 100
    body_trans = unreal2cv2(np.reshape(body_trans, (1, 3)))
    trans = body_trans - cam_trans
    return trans

def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat

def convert_translation_to_opencv(x, y, z):
    t_cv = np.array([y, -z, x])
    return t_cv

def rotation_matrix_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    # Yaw (left-handed)
    R_yaw = np.array(
        [
            [np.cos(-yaw_rad), -np.sin(-yaw_rad), 0],
            [np.sin(-yaw_rad), np.cos(-yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    # Pitch (right-handed)


def get_cam_rotmat(pitch, yaw, roll):
    rotmat_yaw, _ = cv2.Rodrigues(np.array([[0, (yaw / 180) * np.pi, 0]], dtype=float))
    rotmat_pitch, _ = cv2.Rodrigues(np.array([pitch / 180 * np.pi, 0, 0]).reshape(3, 1))
    rotmat_roll, _ = cv2.Rodrigues(np.array([0, 0, roll / 180 * np.pi]).reshape(3, 1))
    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat

def convert_translation_to_opencv(x, y, z):
    t_cv = np.array([y, -z, x])
    return t_cv

def rotation_matrix_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    # Yaw (left-handed)
    R_yaw = np.array(
        [
            [np.cos(-yaw_rad), -np.sin(-yaw_rad), 0],
            [np.sin(-yaw_rad), np.cos(-yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    # Pitch (right-handed)
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, np.sin(pitch_rad)],
            [0, 1, 0],
            [-np.sin(pitch_rad), 0, np.cos(pitch_rad)],


    final_rotmat = rotmat_roll @ (rotmat_pitch @ rotmat_yaw)
    return final_rotmat

def get_global_orient(cam_pitch, cam_yaw, cam_roll):
    pitch_rotmat, _ = cv2.Rodrigues(
        np.array([cam_pitch / 180 * np.pi, 0, 0]).reshape(3, 1)
    )
    roll_rotmat, _ = cv2.Rodrigues(
        np.array([0, 0, cam_roll / 180 * np.pi]).reshape(3, 1)
    )
    final_rotmat = roll_rotmat @ pitch_rotmat
    return final_rotmat

def convert_translation_to_opencv(x, y, z):
    t_cv = np.array([y, -z, x])
    return t_cv

def rotation_matrix_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    # Yaw (left-handed)
    R_yaw = np.array(
        [
            [np.cos(-yaw_rad), -np.sin(-yaw_rad), 0],
            [np.sin(-yaw_rad), np.cos(-yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    # Pitch (right-handed)
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, np.sin(pitch_rad)],
            [0, 1, 0],
            [-np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    # Roll (right-handed)
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), -np.sin(roll_rad)],
            [0, np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_roll @ R_pitch @ R_yaw
    return R_unreal

def convert_rotation_to_opencv(R_unreal):
    # Transformation matrix from Unreal to OpenCV coordinate system.
    C = np.array([[0, 1, 0], [0, 0, -1], [1, 0, 0]])
    R_cv = C @ R_unreal @ C.T
    return R_cv

def get_rot_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    R_yaw = np.array(
        [
            [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
            [np.sin(yaw_rad), np.cos(yaw_rad), 0],
            [0, 0, 1],
        ]
    )


    # Pitch (right-handed)
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, np.sin(pitch_rad)],
            [0, 1, 0],
            [-np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    # Roll (right-handed)
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), -np.sin(roll_rad)],
            [0, np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_roll @ R_pitch @ R_yaw
    return R_unreal

def convert_rotation_to_opencv(R_unreal):
    # Transformation matrix from Unreal to OpenCV coordinate system.
    C = np.array([[0, 1, 0], [0, 0, -1], [1, 0, 0]])
    R_cv = C @ R_unreal @ C.T
    return R_cv

def get_rot_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    R_yaw = np.array(
        [
            [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
            [np.sin(yaw_rad), np.cos(yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, -np.sin(pitch_rad)],
            [0, 1, 0],
            [np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )


    )
    # Roll (right-handed)
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), -np.sin(roll_rad)],
            [0, np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_roll @ R_pitch @ R_yaw
    return R_unreal

def convert_rotation_to_opencv(R_unreal):
    # Transformation matrix from Unreal to OpenCV coordinate system.
    C = np.array([[0, 1, 0], [0, 0, -1], [1, 0, 0]])
    R_cv = C @ R_unreal @ C.T
    return R_cv

def get_rot_unreal(yaw, pitch, roll):
    yaw_rad = np.deg2rad(yaw)
    pitch_rad = np.deg2rad(pitch)
    roll_rad = np.deg2rad(roll)
    R_yaw = np.array(
        [
            [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
            [np.sin(yaw_rad), np.cos(yaw_rad), 0],
            [0, 0, 1],
        ]
    )
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, -np.sin(pitch_rad)],
            [0, 1, 0],
            [np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), np.sin(roll_rad)],
            [0, -np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_yaw @ R_pitch @ R_roll
    return R_unreal

def get_extrinsics_unreal(R_unreal, t_unreal):
    cam_trans = np.array(t_unreal)
    ext = np.eye(4)
    ext[:3, :3] = R_unreal
    ext[:3, 3] = cam_trans.reshape(1, 3)
    return ext

def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
    R_unreal = get_rot_unreal(yaw, pitch, roll)
    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
    T_opencv2unreal = np.array(
        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    T_wu2ou = np.array(
        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
    )


        ]
    )
    R_pitch = np.array(
        [
            [np.cos(pitch_rad), 0, -np.sin(pitch_rad)],
            [0, 1, 0],
            [np.sin(pitch_rad), 0, np.cos(pitch_rad)],
        ]
    )
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), np.sin(roll_rad)],
            [0, -np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_yaw @ R_pitch @ R_roll
    return R_unreal

def get_extrinsics_unreal(R_unreal, t_unreal):
    cam_trans = np.array(t_unreal)
    ext = np.eye(4)
    ext[:3, :3] = R_unreal
    ext[:3, 3] = cam_trans.reshape(1, 3)
    return ext

def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
    R_unreal = get_rot_unreal(yaw, pitch, roll)
    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
    T_opencv2unreal = np.array(
        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    T_wu2ou = np.array(
        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    return np.linalg.inv(T_opencv2unreal @ T_u2wu @ T_wu2ou)

# -----------------------------------------------------------------------------
# Get camera parameters from the extracted images and CSV data.
# -----------------------------------------------------------------------------


    )
    R_roll = np.array(
        [
            [1, 0, 0],
            [0, np.cos(roll_rad), np.sin(roll_rad)],
            [0, -np.sin(roll_rad), np.cos(roll_rad)],
        ]
    )
    R_unreal = R_yaw @ R_pitch @ R_roll
    return R_unreal

def get_extrinsics_unreal(R_unreal, t_unreal):
    cam_trans = np.array(t_unreal)
    ext = np.eye(4)
    ext[:3, :3] = R_unreal
    ext[:3, 3] = cam_trans.reshape(1, 3)
    return ext

def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
    R_unreal = get_rot_unreal(yaw, pitch, roll)
    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
    T_opencv2unreal = np.array(
        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    T_wu2ou = np.array(
        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    return np.linalg.inv(T_opencv2unreal @ T_u2wu @ T_wu2ou)

# -----------------------------------------------------------------------------
# Get camera parameters from the extracted images and CSV data.
# -----------------------------------------------------------------------------

def get_params(
    image_folder,
    fl,
    trans_body,
    cam_x,
    cam_y,
    cam_z,
    fps,
    cam_pitch_,
    cam_roll_,
    cam_yaw_,
):
    all_images = sorted(glob(os.path.join(image_folder, "*" + IMG_FORMAT)))


def get_extrinsics_opencv(yaw, pitch, roll, x, y, z):
    R_unreal = get_rot_unreal(yaw, pitch, roll)
    t_unreal = np.array([x / 100.0, y / 100.0, z / 100.0])
    T_u2wu = get_extrinsics_unreal(R_unreal, t_unreal)
    T_opencv2unreal = np.array(
        [[0, 0, -1, 0], [1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    T_wu2ou = np.array(
        [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32
    )
    return np.linalg.inv(T_opencv2unreal @ T_u2wu @ T_wu2ou)

# -----------------------------------------------------------------------------
# Get camera parameters from the extracted images and CSV data.
# -----------------------------------------------------------------------------

def get_params(
    image_folder,
    fl,
    trans_body,
    cam_x,
    cam_y,
    cam_z,
    fps,
    cam_pitch_,
    cam_roll_,
    cam_yaw_,
):
    all_images = sorted(glob(os.path.join(image_folder, "*" + IMG_FORMAT)))
    imgnames, cam_ext, cam_int = [], [], []

    for img_ind, image_path in enumerate(all_images):
        # Process every 5th frame.
        if img_ind % 5 != 0:
            continue
        cam_ind = img_ind

        cam_pitch_ind = cam_pitch_[cam_ind]
        cam_yaw_ind = cam_yaw_[cam_ind]
        cam_roll_ind = cam_roll_[cam_ind]

        CAM_INT = get_cam_int(fl[cam_ind], SENSOR_W, SENSOR_H, IMG_W / 2.0, IMG_H / 2.0)

        rot_unreal = rotation_matrix_unreal(cam_yaw_ind, cam_pitch_ind, cam_roll_ind)
        rot_cv = convert_rotation_to_opencv(rot_unreal)
        trans_cv = convert_translation_to_opencv(
            cam_x[cam_ind] / 100.0, cam_y[cam_ind] / 100.0, cam_z[cam_ind] / 100.0
        )
        cam_ext_ = np.eye(4)
        cam_ext_[:3, :3] = rot_cv
        # The camera pose is computed as the inverse of the transformed translation.
        cam_ext_[:3, 3] = -rot_cv @ trans_cv

        imgnames.append(
            os.path.join(image_path.split("/")[-2], image_path.split("/")[-1])
        )
        cam_ext.append(cam_ext_)  # camera_pose: c2w
        cam_int.append(CAM_INT)
    return imgnames, cam_ext, cam_int

# -----------------------------------------------------------------------------
# Processing per sequence.
# -----------------------------------------------------------------------------

def process_seq(args):
    """
    Process a single sequence task. For each image, load the corresponding
    depth and image files, and save the computed camera intrinsics and the inverse
    of the extrinsic matrix (i.e. the camera pose in world coordinates) as an NPZ file.
    """
    (
        scene,
        seq_name,
        outdir,
        image_folder_base,
        depth_folder_base,
        mask_folder_base,


            cam_x[cam_ind] / 100.0, cam_y[cam_ind] / 100.0, cam_z[cam_ind] / 100.0
        )
        cam_ext_ = np.eye(4)
        cam_ext_[:3, :3] = rot_cv
        # The camera pose is computed as the inverse of the transformed translation.
        cam_ext_[:3, 3] = -rot_cv @ trans_cv

        imgnames.append(
            os.path.join(image_path.split("/")[-2], image_path.split("/")[-1])
        )
        cam_ext.append(cam_ext_)  # camera_pose: c2w
        cam_int.append(CAM_INT)
    return imgnames, cam_ext, cam_int

# -----------------------------------------------------------------------------
# Processing per sequence.
# -----------------------------------------------------------------------------

def process_seq(args):
    """
    Process a single sequence task. For each image, load the corresponding
    depth and image files, and save the computed camera intrinsics and the inverse
    of the extrinsic matrix (i.e. the camera pose in world coordinates) as an NPZ file.
    """
    (
        scene,
        seq_name,
        outdir,
        image_folder_base,
        depth_folder_base,
        mask_folder_base,
        imgnames,
        cam_ext,
        cam_int,
        humans,
        imgname_array,
    ) = args

    split = 'Test' if scene in test_list else 'Training'

    out_rgb_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'rgb')
    out_depth_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'depth')
    out_cam_dir = os.path.join(outdir, split, "_".join([scene, seq_name]), "cam")
    out_smpl_dir = os.path.join(outdir, split, "_".join([scene, seq_name]), "smpl")
    out_mask_dir = os.path.join(outdir, split, '_'.join([scene, seq_name]), 'mask')
    os.makedirs(out_rgb_dir, exist_ok=True)
    os.makedirs(out_depth_dir, exist_ok=True)
    os.makedirs(out_cam_dir, exist_ok=True)
    os.makedirs(out_smpl_dir, exist_ok=True)
    os.makedirs(out_mask_dir, exist_ok=True)

    assert (
        len(imgnames) == len(cam_ext) == len(cam_int)
    ), f"Inconsistent lengths for {scene}_{seq_name}"

    invalid_images = []
    for imgname, ext, intr, human in zip(imgnames, cam_ext, cam_int, humans):
        if imgname not in imgname_array:
            invalid_images.append(f"Invalid image: {scene}/{imgname}")
            continue

        depthname = imgname.replace(".png", "_depth.exr")
        maskname = imgname.replace(".png", "_env.png")
        imgpath = os.path.join(image_folder_base, imgname)
        depthpath = os.path.join(depth_folder_base, depthname)
        maskpath = os.path.join(mask_folder_base, maskname)
        depth= OpenEXR.File(depthpath).parts[0].channels['Depth'].pixels
        depth = depth.astype(np.float32)/100.0

        mask = cv2.imread(maskpath, cv2.IMREAD_GRAYSCALE)
        if mask is None:
            # if mask file not exist, create a zero mask
            mask = np.zeros((IMG_H, IMG_W), dtype=np.uint8)
        else:
            # invert the mask
            mask = 255 - mask

        outimg_path = os.path.join(out_rgb_dir, os.path.basename(imgpath))
        outmask_path = os.path.join(out_mask_dir, os.path.basename(imgpath))
        outdepth_path = os.path.join(out_depth_dir, os.path.basename(imgpath).replace('.png','.npy'))
        outcam_path = os.path.join(
            out_cam_dir, os.path.basename(imgpath).replace(".png", ".npz")
        )

        out_smpl_path = os.path.join(out_smpl_dir, os.path.basename(imgpath).replace(".png", ".pkl"))

        shutil.copy(imgpath, outimg_path)
        cv2.imwrite(outmask_path, mask)
        np.save(outdepth_path, depth)
        np.savez(outcam_path, intrinsics=intr, pose=np.linalg.inv(ext))  # pose: w2c
        with open(out_smpl_path, 'wb') as f:
            pickle.dump(human, f, protocol=pickle.HIGHEST_PROTOCOL)
    return invalid_images

# -----------------------------------------------------------------------------
# Main entry point.
# -----------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Process Bedlam scenes: compute camera intrinsics and extrinsics, "
        "and save processed camera files."
    )
    parser.add_argument(
        "--root",
        type=str,
        required=True,
        help="Root directory of the extracted data (scenes).",
    )
    parser.add_argument(
        "--outdir", type=str, required=True, help="Output directory for processed data."


main

        outcam_path = os.path.join(
            out_cam_dir, os.path.basename(imgpath).replace(".png", ".npz")
        )

        out_smpl_path = os.path.join(out_smpl_dir, os.path.basename(imgpath).replace(".png", ".pkl"))

        shutil.copy(imgpath, outimg_path)
        cv2.imwrite(outmask_path, mask)
        np.save(outdepth_path, depth)
        np.savez(outcam_path, intrinsics=intr, pose=np.linalg.inv(ext))  # pose: w2c
        with open(out_smpl_path, 'wb') as f:
            pickle.dump(human, f, protocol=pickle.HIGHEST_PROTOCOL)
    return invalid_images

# -----------------------------------------------------------------------------
# Main entry point.
# -----------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Process Bedlam scenes: compute camera intrinsics and extrinsics, "
        "and save processed camera files."
    )
    parser.add_argument(
        "--root",
        type=str,
        required=True,
        help="Root directory of the extracted data (scenes).",
    )
    parser.add_argument(
        "--outdir", type=str, required=True, help="Output directory for processed data."
    )
    parser.add_argument(
        "--annot_dir", type=str, required=True, help="Annotation directory."
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=None,
        help="Number of worker processes (default: os.cpu_count()//2).",
    )
    args = parser.parse_args()

    root = args.root
    outdir = args.outdir
    annot_dir = args.annot_dir
    num_workers = (
        args.num_workers if args.num_workers is not None else (os.cpu_count() or 4) // 2
    )

    invalid_list = []

    # Get scene directories from the root folder.
    scenes = sorted(
        [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
    )
    # Exclude HDRI scenes.
    hdri_scenes = [
        "20221010_3_1000_batch01hand",
        "20221017_3_1000_batch01hand",
        "20221018_3-8_250_batch01hand", # filtered out by CUT3R, but Multi-HMR involves it
        "20221019_3_250_highbmihand",
    ]
    scenes = np.setdiff1d(scenes, hdri_scenes)

    max_human = 0

    tasks = []
    for scene in tqdm(scenes, desc="Collecting tasks"):
        # Skip closeup scenes.
        if "closeup" in scene:
            continue
        base_folder = os.path.join(root, scene)
        image_folder_base = os.path.join(root, scene, "png")
        depth_folder_base = os.path.join(root, scene, "depth")
        mask_folder_base = os.path.join(root, scene, "masks")
        csv_path = os.path.join(base_folder, "be_seq.csv")
        annot_path = os.path.join(annot_dir, scene + "_6fps.npz")
        if not os.path.exists(annot_path):
            annot_path = os.path.join(annot_dir, scene + "_30fps.npz")
        if not os.path.exists(csv_path):
            continue
        csv_data = pd.read_csv(csv_path)
        csv_data = csv_data.to_dict("list")
        cam_csv_base = os.path.join(base_folder, "ground_truth", "camera")
        annot_x = np.load(annot_path)

        # Retrieving SMPL parameters once
        pose_cam_array = annot_x['pose_cam']
        H_array = annot_x['cam_ext']
        shape_array = annot_x['shape']
        imgname_array = annot_x['imgname']
        trans_cam_array = annot_x['trans_cam']

        seq_count = 0
        max_seq_per_scene = 1500000000000

        # Look for a row in the CSV with a "sequence_name" comment.
        for idx, comment in enumerate(csv_data.get("Comment", [])):
            if "sequence_name" in comment:
                if seq_count >= max_seq_per_scene:
                    break

                seq_name = comment.split(";")[0].split("=")[-1]
                if not np.any(np.char.startswith(imgname_array, seq_name + '/')):
                    invalid_list.append(f"Invalid sequence: {scene}/{seq_name}")
                    continue

                cam_csv_path = os.path.join(cam_csv_base, seq_name + "_camera.csv")
                if not os.path.exists(cam_csv_path):
                    continue
                cam_csv_data = pd.read_csv(cam_csv_path)
                cam_csv_data = cam_csv_data.to_dict("list")
                cam_x = cam_csv_data["x"]
                cam_y = cam_csv_data["y"]
                cam_z = cam_csv_data["z"]
                cam_yaw_ = cam_csv_data["yaw"]
                cam_pitch_ = cam_csv_data["pitch"]
                cam_roll_ = cam_csv_data["roll"]
                fl = cam_csv_data["focal_length"]
                image_folder = os.path.join(image_folder_base, seq_name)
                trans_body = None  # Not used here.
                imgnames, cam_ext, cam_int = get_params(
                    image_folder,
                    fl,
                    trans_body,
                    cam_x,
                    cam_y,
                    cam_z,
                    6,
                    cam_pitch_=cam_pitch_,
                    cam_roll_=cam_roll_,
                    cam_yaw_=cam_yaw_,
                )
                humans = [] # humans for each sequence
                for imgname in imgnames:
                    idxs = np.where(imgname == imgname_array)[0]

                    if len(idxs) > max_human:
                        max_human = len(idxs)

                    persons_per_img = [] # persons for each image
                    for i in idxs:
                        sys.stdout.flush()

                        # SMPLX params
                        pose = pose_cam_array[i]
                        root_pose = pose[:3]
                        body_pose=pose[3:66]
                        jaw_pose=pose[66:69]
                        leye_pose=pose[69:72]
                        reye_pose=pose[72:75]
                        left_hand_pose=pose[75:120]
                        right_hand_pose=pose[120:165]
                        betas=shape_array[i]
                        transl = trans_cam_array[i] + H_array[i][:, 3][:3]

                        person = {
                            # SMPL GT in camera coordinates system
                            'smplx_root_pose': root_pose.reshape(1,3), # axis-angle
                            'smplx_body_pose': body_pose.reshape(21,3), # axis-angle
                            'smplx_jaw_pose': jaw_pose.reshape(1,3), # axis-angle
                            'smplx_leye_pose': leye_pose.reshape(1,3), # axis-angle
                            'smplx_reye_pose': reye_pose.reshape(1,3), # axis-angle
                            'smplx_left_hand_pose': left_hand_pose.reshape(15,3), # axis-angle
                            'smplx_right_hand_pose': right_hand_pose.reshape(15,3), # axis-angle
                            'smplx_shape': betas.reshape(11),
                            'smplx_gender': 'neutral',
                            'smplx_transl': transl.reshape(3),
                        }
                        persons_per_img.append(person)
                    humans.append(persons_per_img)

                tasks.append(
                    (
                        scene,
                        seq_name,
                        outdir,
                        image_folder_base,
                        depth_folder_base,
                        mask_folder_base,
                        imgnames,
                        cam_ext,
                        cam_int,
                        humans,
                        imgname_array,
                    )
                )

                seq_count += 1

    print(f"max_human: {max_human}")

    # Process each task in parallel.
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = {executor.submit(process_seq, task): task for task in tasks}
        for future in tqdm(
            as_completed(futures), total=len(futures), desc="Processing sequences"
        ):
            # error = future.result()
            # if error:
            #     print(error)
            invalid_images = future.result()
            if invalid_images:
                invalid_list.extend(invalid_images)

    # print invalid items
    if invalid_list:
        print("\nInvalid sequences and images:")
        for item in invalid_list:
            print(item)

if __name__ == "__main__":
    main()


src.train

# --------------------------------------------------------
# training code for Human3R
# --------------------------------------------------------
# References:
# CUT3R: https://github.com/CUT3R/CUT3R
# DUSt3R: https://github.com/naver/dust3r
# --------------------------------------------------------
import argparse
import datetime
import json
import numpy as np
import os
import sys
import time
import math
from collections import defaultdict
from pathlib import Path
from typing import Sized

import torch
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

from dust3r.utils.device import todevice

torch.backends.cuda.matmul.allow_tf32 = True  # for gpu >= Ampere and pytorch >= 1.12

from dust3r.model import (
    PreTrainedModel,
    ARCroco3DStereo,
    ARCroco3DStereoConfig,
    inf,
    strip_module,
    strip_module_mhmr,
)  # noqa: F401, needed when loading the model
from dust3r.smpl_model import SMPLModel
from dust3r.datasets import get_data_loader
from dust3r.losses import *  # noqa: F401, needed when loading the model
from dust3r.inference import loss_of_one_batch  # noqa
from dust3r.viz import colorize
from dust3r.utils.render import get_render_results, get_render_smpl
import dust3r.utils.path_to_croco  # noqa: F401
import croco.utils.misc as misc  # noqa
from croco.utils.misc import NativeScalerWithGradNormCount as NativeScaler  # noqa

import hydra
from omegaconf import OmegaConf
import logging
import pathlib
from tqdm import tqdm
import random
import builtins
import shutil

from accelerate import Accelerator
from accelerate import DistributedDataParallelKwargs, InitProcessGroupKwargs
from accelerate.logging import get_logger
from datetime import timedelta
import torch.multiprocessing

torch.multiprocessing.set_sharing_strategy("file_system")

printer = get_logger(__name__, log_level="DEBUG")

def setup_for_distributed(accelerator: Accelerator):
    """
    This function disables printing when not in master process
    """
    builtin_print = builtins.print

    def print(*args, **kwargs):
        force = kwargs.pop("force", False)
        force = force or (accelerator.num_processes > 8)
        if accelerator.is_main_process or force:
            now = datetime.datetime.now().time()
            builtin_print("[{}] ".format(now), end="")  # print with time stamp
            builtin_print(*args, **kwargs)

    builtins.print = print

def save_current_code(outdir):
    now = datetime.datetime.now()  # current date and time
    date_time = now.strftime("%m_%d-%H:%M:%S")
    src_dir = "."
    dst_dir = os.path.join(outdir, "code", "{}".format(date_time))
    shutil.copytree(
        src_dir,
        dst_dir,
        ignore=shutil.ignore_patterns(
            ".vscode*",
            "assets*",
            "example*",
            "checkpoints*",
            "OLD*",
            "logs*",
            "out*",
            "runs*",
            "*.png",
            "*.mp4",
            "*__pycache__*",
            "*.git*",
            "*.idea*",
            "*.zip",
            "*.jpg",
            "*.pth",
            "*.pt",
            "*.npy",
            "*.npz",
            "*.pkl",
        ),
        dirs_exist_ok=True,
    )
    return dst_dir

def train(args):

    accelerator = Accelerator(
        gradient_accumulation_steps=args.accum_iter,
        mixed_precision="bf16",
        kwargs_handlers=[
            DistributedDataParallelKwargs(find_unused_parameters=True),
            InitProcessGroupKwargs(timeout=timedelta(seconds=6000)),
        ],
    )
    device = accelerator.device

    setup_for_distributed(accelerator)

    printer.info("output_dir: " + args.output_dir)
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    if accelerator.is_main_process:
        dst_dir = save_current_code(outdir=args.output_dir)
        printer.info(f"Saving current code to {dst_dir}")

    # auto resume
    if not args.resume:
        last_ckpt_fname = os.path.join(args.output_dir, f"checkpoint-last.pth")
        args.resume = last_ckpt_fname if os.path.isfile(last_ckpt_fname) else None

    printer.info("job dir: {}".format(os.path.dirname(os.path.realpath(__file__))))

    # fix the seed
    seed = args.seed + accelerator.state.process_index
    printer.info(
        f"Setting seed to {seed} for process {accelerator.state.process_index}"
    )
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = args.benchmark

    # training dataset and loader
    printer.info("Building train dataset %s", args.train_dataset)
    #  dataset and loader
    data_loader_train = build_dataset(
        args.train_dataset,
        args.batch_size,
        args.num_workers,
        accelerator=accelerator,
        test=False,
        fixed_length=args.fixed_length
    )
    printer.info("Building test dataset %s", args.test_dataset)
    data_loader_test = {
        dataset.split("(")[0]: build_dataset(
            dataset,
            args.batch_size,
            args.num_workers,
            accelerator=accelerator,
            test=True,
            fixed_length=True
        )
        for dataset in args.test_dataset.split("+")
    }

    # model
    printer.info("Loading model: %s", args.model)
    model: PreTrainedModel = eval(args.model)
    smpl_model: SMPLModel = SMPLModel(
        device,
        model_args={
            'patch_size': model.croco_args['patch_size'],
            'mhmr_img_res': model.mhmr_img_res,
            'bb_patch_size': model.bb_patch_size
        })
    printer.info(f"All model parameters: {sum(p.numel() for p in model.parameters())}")
    printer.info(
        f"Encoder parameters: {sum(p.numel() for p in model.enc_blocks.parameters())}"
    )
    printer.info(
        f"Decoder parameters: {sum(p.numel() for p in model.dec_blocks.parameters())}"
    )

    printer.info(f">> Creating train criterion = {args.train_criterion}")
    train_criterion = eval(args.train_criterion).to(device)
    printer.info(
        f">> Creating test criterion = {args.test_criterion or args.train_criterion}"
    )
    test_criterion = eval(args.test_criterion or args.criterion).to(device)

    model.to(device)

    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
    if args.long_context:
        model.fixed_input_length = False

    if args.pretrained and not args.resume:
        printer.info(f"Loading pretrained: {args.pretrained}")
        ckpt = torch.load(args.pretrained, map_location=device)
        load_only_encoder = getattr(args, "load_only_encoder", False)
        if load_only_encoder:
            filtered_state_dict = {
                k: v
                for k, v in ckpt["model"].items()
                if "enc_blocks" in k or "patch_embed" in k
            }
            merge_state_dict = strip_module(filtered_state_dict)
        else:
            merge_state_dict = strip_module(ckpt["model"])
        del ckpt  # in case it occupies memory

        if args.pretrained_mhmr:
            printer.info(f"Loading Multi-HMR pretrained: {args.pretrained_mhmr}")
            ckpt_mhmr = torch.load(args.pretrained_mhmr, map_location=device)
            merge_state_dict.update(strip_module_mhmr(ckpt_mhmr["model_state_dict"]))
            del ckpt_mhmr  # in case it occupies memory

        printer.info(
            model.load_state_dict(merge_state_dict, strict=False)
        )

    # # following timm: set wd as 0 for bias and norm layers
    param_groups = misc.get_parameter_groups(model, args.weight_decay)
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
    # print(optimizer)
    loss_scaler = NativeScaler(accelerator=accelerator)

    accelerator.even_batches = False
    optimizer, model, data_loader_train = accelerator.prepare(
        optimizer, model, data_loader_train
    )

    def write_log_stats(epoch, train_stats, test_stats):
        if accelerator.is_main_process:
            if log_writer is not None:
                log_writer.flush()

            log_stats = dict(
                epoch=epoch, **{f"train_{k}": v for k, v in train_stats.items()}
            )
            for test_name in data_loader_test:
                if test_name not in test_stats:
                    continue
                log_stats.update(
                    {test_name + "_" + k: v for k, v in test_stats[test_name].items()}
                )

            with open(
                os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8"
            ) as f:
                f.write(json.dumps(log_stats) + "\n")

    def save_model(epoch, fname, best_so_far):
        misc.save_model(
            accelerator=accelerator,
            args=args,
            model_without_ddp=model,
            optimizer=optimizer,
            loss_scaler=loss_scaler,
            epoch=epoch,
            fname=fname,
            best_so_far=best_so_far,
        )

    best_so_far = misc.load_model(
        args=args, model_without_ddp=model, optimizer=optimizer, loss_scaler=loss_scaler
    )
    if best_so_far is None:
        best_so_far = float("inf")
    log_writer = (
        SummaryWriter(log_dir=args.output_dir) if accelerator.is_main_process else None
    )

    printer.info(f"Start training for {args.epochs} epochs")
    start_time = time.time()
    train_stats = test_stats = {}

    for epoch in range(args.start_epoch, args.epochs + 1):

        # Save immediately the last checkpoint
        if epoch > args.start_epoch:
            if (
                args.save_freq
                and np.allclose(epoch / args.save_freq, int(epoch / args.save_freq))
                or epoch == args.epochs
            ):
                save_model(epoch - 1, "last", best_so_far)

        # Test on multiple datasets
        new_best = False
        if epoch >= 0 and args.eval_freq > 0 and epoch % args.eval_freq == 0:
            test_stats = {}
            for test_name, testset in data_loader_test.items():
                stats = test_one_epoch(
                    model,
                    test_criterion,
                    testset,
                    accelerator,
                    device,
                    epoch,
                    log_writer=log_writer,
                    args=args,
                    prefix=test_name,
                    smpl_model=smpl_model,
                )
                test_stats[test_name] = stats

                # Save best of all
                if stats["loss_med"] < best_so_far:
                    best_so_far = stats["loss_med"]
                    new_best = True
        # Save more stuff
        write_log_stats(epoch, train_stats, test_stats)

        if epoch > args.start_epoch:
            if args.keep_freq and epoch % args.keep_freq == 0:
                save_model(epoch - 1, str(epoch), best_so_far)
            if new_best:
                save_model(epoch - 1, "best", best_so_far)
        if epoch >= args.epochs:
            break  # exit after writing last test to disk

        # Train
        train_stats = train_one_epoch(
            model,
            train_criterion,
            data_loader_train,
            optimizer,
            accelerator,
            epoch,
            loss_scaler,
            log_writer=log_writer,
            args=args,
            smpl_model=smpl_model,
        )

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    printer.info("Training time {}".format(total_time_str))

    save_final_model(accelerator, args, args.epochs, model, best_so_far=best_so_far)

def save_final_model(accelerator, args, epoch, model_without_ddp, best_so_far=None):
    output_dir = Path(args.output_dir)
    checkpoint_path = output_dir / "checkpoint-final.pth"
    to_save = {
        "args": args,
        "model": (
            model_without_ddp
            if isinstance(model_without_ddp, dict)
            else model_without_ddp.cpu().state_dict()
        ),
        "epoch": epoch,
    }
    if best_so_far is not None:
        to_save["best_so_far"] = best_so_far
    printer.info(f">> Saving model to {checkpoint_path} ...")
    misc.save_on_master(accelerator, to_save, checkpoint_path)

def build_dataset(dataset, batch_size, num_workers, accelerator, test=False, fixed_length=False):
    split = ["Train", "Test"][test]
    printer.info(f"Building {split} Data loader for dataset: {dataset}")
    loader = get_data_loader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        pin_mem=True,
        shuffle=not (test),
        drop_last=not (test),
        accelerator=accelerator,
        fixed_length=fixed_length
    )
    return loader

def train_one_epoch(
    model: torch.nn.Module,
    criterion: torch.nn.Module,
    data_loader: Sized,
    optimizer: torch.optim.Optimizer,
    accelerator: Accelerator,
    epoch: int,
    loss_scaler,
    args,
    log_writer=None,
    smpl_model: SMPLModel = None
):
    assert torch.backends.cuda.matmul.allow_tf32 == True

    model.train(True)
    metric_logger = misc.MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
    header = "Epoch: [{}]".format(epoch)
    accum_iter = args.accum_iter

    def save_model(epoch, fname, best_so_far):
        misc.save_model(
            accelerator=accelerator,
            args=args,
            model_without_ddp=model,
            optimizer=optimizer,
            loss_scaler=loss_scaler,
            epoch=epoch,
            fname=fname,
            best_so_far=best_so_far,
        )

    if log_writer is not None:
        printer.info("log_dir: {}".format(log_writer.log_dir))

    if hasattr(data_loader, "dataset") and hasattr(data_loader.dataset, "set_epoch"):
        data_loader.dataset.set_epoch(epoch)
    if (
        hasattr(data_loader, "batch_sampler")
        and hasattr(data_loader.batch_sampler, "batch_sampler")
        and hasattr(data_loader.batch_sampler.batch_sampler, "set_epoch")
    ):
        data_loader.batch_sampler.batch_sampler.set_epoch(epoch)

    optimizer.zero_grad()

    for data_iter_step, batch in enumerate(
        metric_logger.log_every(data_loader, args.print_freq, accelerator, header)
    ):
        with accelerator.accumulate(model):
            epoch_f = epoch + data_iter_step / len(data_loader)
            step = int(epoch_f * len(data_loader))
            # we use a per iteration (instead of per epoch) lr scheduler
            if data_iter_step % accum_iter == 0:
                misc.adjust_learning_rate(optimizer, epoch_f, args)
            if not args.long_context:
                result = loss_of_one_batch(
                    batch,
                    model,
                    criterion,
                    accelerator,
                    symmetrize_batch=False,
                    use_amp=bool(args.amp),
                    smpl_model=smpl_model
                )
            else:
                NotImplementedError("Long context is not supported")
            has_msk = "msk" in result["pred"][0]
            loss, loss_details = result["loss"]  # criterion returns two values
            loss_value = float(loss)

            if not math.isfinite(loss_value):
                print(
                    f"Loss is {loss_value}, stopping training, loss details: {loss_details}"
                )
                sys.exit(1)
            if not result.get("already_backprop", False):
                loss_scaler(
                    loss,
                    optimizer,
                    parameters=model.parameters(),
                    update_grad=True,
                    clip_grad=1.0,
                )
                optimizer.zero_grad()

            is_metric = batch[0]["is_metric"]
            curr_num_view = len(batch)

            del loss
            tb_vis_img = (data_iter_step + 1) % accum_iter == 0 and (
                (step + 1) % (args.print_img_freq)
            ) == 0
            if not tb_vis_img:
                del batch
            else:
                torch.cuda.empty_cache()

            lr = optimizer.param_groups[0]["lr"]
            metric_logger.update(epoch=epoch_f)
            metric_logger.update(lr=lr)
            metric_logger.update(step=step)

            metric_logger.update(loss=loss_value, **loss_details)

            if (data_iter_step + 1) % accum_iter == 0 and (
                (data_iter_step + 1) % (accum_iter * args.print_freq)
            ) == 0:
                loss_value_reduce = accelerator.gather(
                    torch.tensor(loss_value).to(accelerator.device)
                ).mean()  # MUST BE EXECUTED BY ALL NODES

                if log_writer is None:
                    continue
                """ We use epoch_1000x as the x-axis in tensorboard.
                This calibrates different curves when batch size changes.
                """
                epoch_1000x = int(epoch_f * 1000)
                log_writer.add_scalar("train_loss", loss_value_reduce, step)
                log_writer.add_scalar("train_lr", lr, step)
                log_writer.add_scalar("train_iter", epoch_1000x, step)
                for name, val in loss_details.items():
                    if isinstance(val, torch.Tensor):
                        if val.ndim > 0:
                            continue
                    if isinstance(val, dict):
                        continue
                    log_writer.add_scalar("train_" + name, val, step)

            if tb_vis_img:
                if log_writer is None:
                    continue
                with torch.no_grad():
                    depths_self, gt_depths_self = get_render_results(
                        batch, result["pred"], self_view=True
                    )
                    depths_cross, gt_depths_cross = get_render_results(
                        batch, result["pred"], self_view=False
                    )
                    gt_msks, pr_msks, gt_hms, pr_hms, gt_smpls, pr_smpls = get_render_smpl(
                        batch, result["pred"], smpl_model, loss_details, has_msk=has_msk
                    )
                    for k in range(len(batch)):
                        loss_details[f"self_pred_depth_{k+1}"] = depths_self[k].detach().cpu()
                        loss_details[f"self_gt_depth_{k+1}"] = gt_depths_self[k].detach().cpu()
                        loss_details[f"pred_depth_{k+1}"] = depths_cross[k].detach().cpu()
                        loss_details[f"gt_depth_{k+1}"] = gt_depths_cross[k].detach().cpu()
                        loss_details[f"pred_hm_{k+1}"] = pr_hms[k].detach().cpu()
                        loss_details[f"gt_hm_{k+1}"] = gt_hms[k].detach().cpu()
                        loss_details[f"pred_smpl_rend_{k+1}"] = pr_smpls[k].detach().cpu()
                        loss_details[f"gt_smpl_rend_{k+1}"] = gt_smpls[k].detach().cpu()
                        if has_msk:
                            loss_details[f"pred_msk_{k+1}"] = pr_msks[k].detach().cpu()
                            loss_details[f"gt_msk_{k+1}"] = gt_msks[k].detach().cpu()

                imgs_stacked_dict = get_vis_imgs_new(
                    loss_details,
                    args.num_imgs_vis,
                    curr_num_view,
                    is_metric=is_metric,
                    has_msk=has_msk)
                for name, imgs_stacked in imgs_stacked_dict.items():
                    log_writer.add_images(
                        "train" + "/" + name, imgs_stacked, step, dataformats="HWC"
                    )
                del batch

        if (
            data_iter_step % int(args.save_freq * len(data_loader)) == 0
            and data_iter_step != 0
            and data_iter_step != len(data_loader) - 1
        ):
            print("saving at step", data_iter_step)
            sa
# ... truncated ...


import hydra
from omegaconf import OmegaConf
import logging
import pathlib
from tqdm import tqdm
import random
import builtins
import shutil

from accelerate import Accelerator
from accelerate import DistributedDataParallelKwargs, InitProcessGroupKwargs
from accelerate.logging import get_logger
from datetime import timedelta
import torch.multiprocessing

torch.multiprocessing.set_sharing_strategy("file_system")

printer = get_logger(__name__, log_level="DEBUG")

def setup_for_distributed(accelerator: Accelerator):
    """
    This function disables printing when not in master process
    """
    builtin_print = builtins.print

    def print(*args, **kwargs):
        force = kwargs.pop("force", False)
        force = force or (accelerator.num_processes > 8)
        if accelerator.is_main_process or force:
            now = datetime.datetime.now().time()
            builtin_print("[{}] ".format(now), end="")  # print with time stamp
            builtin_print(*args, **kwargs)

    builtins.print = print

def save_current_code(outdir):
    now = datetime.datetime.now()  # current date and time
    date_time = now.strftime("%m_%d-%H:%M:%S")
    src_dir = "."
    dst_dir = os.path.join(outdir, "code", "{}".format(date_time))
    shutil.copytree(
        src_dir,
        dst_dir,
        ignore=shutil.ignore_patterns(
            ".vscode*",
            "assets*",
            "example*",
            "checkpoints*",
            "OLD*",
            "logs*",
            "out*",
            "runs*",
            "*.png",


printer = get_logger(__name__, log_level="DEBUG")

def setup_for_distributed(accelerator: Accelerator):
    """
    This function disables printing when not in master process
    """
    builtin_print = builtins.print

    def print(*args, **kwargs):
        force = kwargs.pop("force", False)
        force = force or (accelerator.num_processes > 8)
        if accelerator.is_main_process or force:
            now = datetime.datetime.now().time()
            builtin_print("[{}] ".format(now), end="")  # print with time stamp
            builtin_print(*args, **kwargs)

    builtins.print = print

def save_current_code(outdir):
    now = datetime.datetime.now()  # current date and time
    date_time = now.strftime("%m_%d-%H:%M:%S")
    src_dir = "."
    dst_dir = os.path.join(outdir, "code", "{}".format(date_time))
    shutil.copytree(
        src_dir,
        dst_dir,
        ignore=shutil.ignore_patterns(
            ".vscode*",
            "assets*",
            "example*",
            "checkpoints*",
            "OLD*",
            "logs*",
            "out*",
            "runs*",
            "*.png",
            "*.mp4",
            "*__pycache__*",
            "*.git*",
            "*.idea*",
            "*.zip",
            "*.jpg",
            "*.pth",
            "*.pt",
            "*.npy",
            "*.npz",
            "*.pkl",
        ),
        dirs_exist_ok=True,
    )
    return dst_dir

def train(args):

    accelerator = Accelerator(
        gradient_accumulation_steps=args.accum_iter,
        mixed_precision="bf16",
        kwargs_handlers=[
            DistributedDataParallelKwargs(find_unused_parameters=True),
            InitProcessGroupKwargs(timeout=timedelta(seconds=6000)),
        ],
    )
    device = accelerator.device

    setup_for_distributed(accelerator)

    printer.info("output_dir: " + args.output_dir)
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)