# Created by BVH, Apr 2026.
# Camera-related dataset utilities: extrinsic re-anchoring, degenerate intrinsic handling.

import torch

from anydata.geometry.camera import Camera
from imaginaire.utils import log


def reanchor_extrinsics(extrinsics, reference_camera, ref_cam_rel_per_timestep=False,
                        camera_names=None):
    '''Re-anchor extrinsics so that the reference camera becomes identity.
        NOTE(bvh): this happens PRE collate.
    :param extrinsics: dict mapping (time, cam_idx) to (4, 4) Tcw tensor.
    :param reference_camera: camera index/name to anchor to.
    :param ref_cam_rel_per_timestep: if True, anchor each frame t to ref cam at t.
        If False, anchor all frames to ref cam at t=0.
    :param camera_names: optional list mapping int index to str name (from metadata['cameras']).
    :return: new dict of re-anchored extrinsics (input not mutated).
    '''
    if not isinstance(extrinsics, dict) or len(extrinsics) == 0:
        return extrinsics

    if camera_names is None:
        if isinstance(reference_camera, int):
            ref_cam_idx = reference_camera
        else:
            log.error(f'Reference camera {reference_camera!r} is a name but camera_names is None; '
                      f'using first camera (index = 0) instead', rank0_only=False)
            ref_cam_idx = 0
    elif reference_camera not in camera_names:
        log.error(f'Reference camera {reference_camera} not present in loaded cam_names: {camera_names}; '
                    f'using first camera (index = 0) instead', rank0_only=False)
        ref_cam_idx = 0
    else:
        ref_cam_idx = camera_names.index(reference_camera)

    result = {}
    for (t, c) in extrinsics.keys():
        
        ref_t = t if ref_cam_rel_per_timestep else 0
        ref_key = (ref_t, ref_cam_idx)
        if ref_key not in extrinsics:
            log.error(f'Reference camera {ref_cam_idx} not found at t={ref_t}, skipping re-anchor', rank0_only=False)
            result[(t, c)] = extrinsics[(t, c)]
            continue

        # NOTE(bvh): do everything in high precision to minimize numerical errors
        ref_Tcw = extrinsics[ref_key].to(torch.float64)
        ref_Twc = torch.linalg.inv(ref_Tcw)
        cur_Tcw = extrinsics[(t, c)].to(torch.float64)
        new_Tcw = ref_Twc @ cur_Tcw
        result[(t, c)] = new_Tcw.to(torch.float32)

    return result


def reanchor_cameras_to_view(cams_dict, ref_cam_idx, per_frame=True):
    '''Re-anchor Camera objects so that the reference camera becomes identity.
        Used by reference_view to make plucker rays relative to a specific view.
        NOTE(bvh): this happens POST collate.
    :param cams_dict: dict mapping (time, cam_idx) to Camera objects,
        containing batched Tcw extrinsics (B, 4, 4).
    :param ref_cam_idx: camera index to use as reference (from used_cams[reference_view]).
    :param per_frame: if True, anchor each frame t to ref cam at t; if False, anchor all to t=0.
    :return: new dict of re-anchored Camera objects (input not mutated).
    '''
    result = {}
    
    for (t, c), cam in cams_dict.items():
        ref_t = t if per_frame else 0
        ref_key = (ref_t, ref_cam_idx)
        if ref_key not in cams_dict:
            result[(t, c)] = cam
            continue
        
        # NOTE(bvh): do everything in high precision to minimize numerical errors
        ref_Twc = cams_dict[ref_key].Twc.to(torch.float64)  # (B, 4, 4)
        ref_Tcw = torch.linalg.inv(ref_Twc)
        cam_Twc = cam.Twc.to(torch.float64)
        new_Twc = cam_Twc @ ref_Tcw
        new_Twc = new_Twc.to(torch.float32)
        result[(t, c)] = Camera(K=cam.K, Twc=new_Twc, hw=cam.hw)
    
    return result


def sanitize_intrinsics(K, hw, key=None):
    '''Guard against degenerate intrinsics (e.g. all-zero K from DROID episodes
    without calibration). These produce NaN in plucker ray computation.
    Replaces degenerate K (fx=fy=0) with a default pinhole (fx=W/2, fy=H/2).

    :param K: (3, 3) or (B, 3, 3) intrinsics tensor, or None.
    :param hw: (H, W) tuple or similar.
    :param key: optional (t, c) key for logging.
    :return: sanitized K (same shape), or None if input was None.
    '''
    if K is None or not isinstance(K, torch.Tensor):
        return K

    if isinstance(hw, dict):
        hw = next(iter(hw.values()))
    _H, _W = hw if isinstance(hw, tuple) else (hw[0], hw[1])

    if K.dim() == 2 and K.shape == (3, 3):
        if K[0, 0].abs() < 1e-6 and K[1, 1].abs() < 1e-6:
            log.warning(f'Degenerate intrinsics (fx=fy=0) at {key}, using default pinhole instead', rank0_only=False)
            return None  # signals caller to use Camera's built-in default
    elif K.dim() == 3 and K.shape[-2:] == (3, 3):
        degen_mask = (K[:, 0, 0].abs() < 1e-6) & (K[:, 1, 1].abs() < 1e-6)
        if degen_mask.any():
            default_K = torch.tensor([
                [_W / 2, 0, _W / 2],
                [0, _H / 2, _H / 2],
                [0, 0, 1.0],
            ], dtype=K.dtype, device=K.device)
            K = K.clone()
            K[degen_mask] = default_K
            n_degen = degen_mask.sum().item()
            log.warning(f'Degenerate intrinsics (fx=fy=0) in {n_degen}/{K.shape[0]} '
                        f'batch elements at {key}, replaced with default pinhole', rank0_only=False)
    return K