# Created by BVH, Mar 2026.

# DataTransforms: standalone replacement for vidar's GenericModel.prepare_batch.
# Applies temporal stride, camera building, depth clipping/scaling.
# All parameters come from per-dataset configs and overrides, passed per-call.
# Only used by the anydata pipeline (vidar uses legacy GenericModel transforms).

from __future__ import annotations

from typing import Optional

import torch

from anydata.geometry.camera import Camera
from custom.dataset.camera_utils import sanitize_intrinsics
from custom.utils.constants import VAE_RATIO_THW
from imaginaire.utils import log


def _safe_inv(mat):
    '''
    Invert a batch of matrices, returning zeros for singular/zero-padded entries.
    '''
    # mat: (..., 4, 4)
    try:
        return torch.linalg.inv(mat)
    except torch.linalg.LinAlgError:
        pass
    # Fallback: per-element inversion for mixed valid/zero-padded batches.
    result = torch.zeros_like(mat)
    flat = mat.reshape(-1, *mat.shape[-2:])
    out = result.reshape(-1, *mat.shape[-2:])
    for i in range(flat.shape[0]):
        if flat[i].any():
            out[i] = torch.linalg.inv(flat[i])
    return result


class DataTransforms:
    '''
    Stateless transform engine for the anydata pipeline.
    All parameters come from per-dataset configs and overrides, passed per-call.
    NOTE(bvh): this whole thing happens POST-collate, but PRE-dataloader.
    '''

    def prepare_batch(
        self,
        batch: dict,
        fps=None,
        frame_stride: int = 1,
        frame_trim: int = 0,
        frame_offset: float = 0.0,
        clip_depth: Optional[float] = None,
        scale_depth: Optional[float] = None,
    ):
        '''
        Main entry point, replacing GenericModel.prepare_batch.
        :return (batch, fps): fps is scaled by frame_stride alongside the frame subsampling, so the
            two stay coupled.
        '''
        if batch.get('prepared', False):
            return batch, fps

        # 1. Temporal stride + trim: subsample/cut frames and scale fps to match.
        batch, fps = self._apply_temporal_stride_and_trim(
            batch, fps, frame_stride, frame_trim, frame_offset)

        # 2. Build cameras from intrinsics + extrinsics (if not already present).
        if 'cams' not in batch:
            batch = self._build_cameras(batch)

        # 3. Clip depth if configured
        if clip_depth is not None and 'depth' in batch:
            batch = self._clip_depth(batch, clip_depth)

        # 4. Scale depth if configured
        if scale_depth is not None and 'depth' in batch:
            batch = self._scale_depth(batch, scale_depth)

        batch['prepared'] = True
        return batch, fps

    def _apply_temporal_stride_and_trim(
        self,
        batch: dict,
        fps,
        frame_stride: int,
        frame_trim: int = 0,
        frame_offset: float = 0.0,
    ):
        '''
        Subsample frames by frame_stride, then cap the count to frame_trim (cut at the end), re-index
        contiguous, and scale fps by the stride. Operates on every temporal-keyed sub-dict, whether
        keyed by (time, cam) tuples (high-dim rgb/depth/cams) or by int time (low-dim actions), so it
        stays general. Kept length is VAE-trimmed to T = k * VAE_RATIO_THW[0] + 1.
        NOTE(bvh): must be POST-collate to keep length (# frames) consistent within batches.
        :param frame_stride: int >= 1; <= 1 means no striding.
        :param frame_trim: upper bound on # frames (cut at end); <= 0 means no trim.
        :param frame_offset: pre-sampled fraction in [0,1) of the available window slack; 0.0 = front.
        Also records the kept original frame indices in batch['frame_indices'] (for the visuals footer).
        :return (batch, fps): stride + trim applied together; fps scaled by frame_stride.
        '''
        frame_stride = frame_stride or 1
        frame_trim = frame_trim or 0

        def _time(k):  # time component of a (time, cam) tuple key or a bare int time key
            return k[0] if isinstance(k, tuple) else k

        def _is_temporal(d):  # keyed by (time, cam) tuples (high dim) or bare int time (low dim)
            sk = next(iter(d), None)
            return sk is not None and (
                (isinstance(sk, tuple) and len(sk) == 2) or isinstance(sk, int))

        if frame_stride <= 1 and frame_trim <= 0:
            # No striding/trimming, but still record the (already contiguous) kept indices for the footer.
            for key in list(batch.keys()):
                if isinstance(batch[key], dict) and _is_temporal(batch[key]):
                    batch['frame_indices'] = sorted(set(_time(k) for k in batch[key]))
                    break
            return batch, fps

        # apply the pre-sampled window offset: fraction in [0,1) of the slack [0, post_stride_len -
        # frame_trim], shared across all modalities/views (0.0 = front).
        offset = 0
        if frame_offset > 0.0 and frame_trim > 0:
            for key in list(batch.keys()):
                if isinstance(batch[key], dict) and _is_temporal(batch[key]):
                    n_post = len(sorted(set(_time(k) for k in batch[key]))[0::frame_stride])
                    max_offset = n_post - frame_trim
                    if max_offset > 0:
                        offset = min(int(frame_offset * (max_offset + 1)), max_offset)
                    break

        frame_indices = None  # canonical kept original indices (same across temporal entries)
        for key in list(batch.keys()):
            if not isinstance(batch[key], dict):
                continue
            if not _is_temporal(batch[key]):
                continue

            # Actually apply selected stride & trim
            orig_times = sorted(set(_time(k) for k in batch[key]))
            kept = orig_times[0::frame_stride]            # stride: [0,1,..,12] -> [0,2,..,12]
            if frame_trim > 0:
                kept = kept[offset:offset + frame_trim]   # trim window
            
            # VAE tokenizer length constraint: latent T = (T-1)/VAE_RATIO_THW[0] + 1, so T = k*ratio+1.
            while len(kept) > 0 and (len(kept) - 1) % VAE_RATIO_THW[0] != 0:
                kept = kept[:-1]
            if len(kept) == 0:
                log.warning(f'No valid times left: frame_stride={frame_stride}, frame_trim={frame_trim}, '
                            f'key={key}, sample_key={next(iter(batch[key]), None)}', rank0_only=False)
                continue
            if frame_indices is None:
                frame_indices = list(kept)  # original indices that survived stride+trim

            # Re-index kept times to contiguous 0..N-1 (downstream assumes contiguous increasing).
            remap = {t: i for i, t in enumerate(kept)}
            batch[key] = {
                ((remap[_time(k)], k[1]) if isinstance(k, tuple) else remap[k]): val
                for k, val in batch[key].items()
                if _time(k) in remap
            }

        if frame_indices is not None:
            batch['frame_indices'] = frame_indices

        if fps is not None and frame_stride > 1:
            fps = fps / frame_stride

        return batch, fps

    def _build_cameras(self, batch: dict) -> dict:
        '''
        Create Camera objects from intrinsics and (absolute cam-to-world) extrinsics. Produces:
            batch['cams_orig'] - global cameras (for action visualization / projection).
            batch['cams']      - cameras for plucker rays (already re-anchored pre-collate in
                                 AnyDataset.__getitem__ if reference_camera was set).
        '''
        intrinsics = batch.get('intrinsics', None)
        extrinsics = batch.get('extrinsics', None)
        extrinsics_orig = batch.get('extrinsics_orig', None)

        # Skip if missing or if collate returned a list (mixed None/dict from heterogeneous batches).
        if intrinsics is None or extrinsics is None:
            return batch
        if not isinstance(extrinsics, dict) or not isinstance(intrinsics, dict):
            return batch

        # Determine image resolution
        if 'rgb' in batch:
            hw = {k: val.shape[-2:] for k, val in batch['rgb'].items()}
        elif 'lat_rgb' in batch:
            base_key = next(iter(batch['lat_rgb'].keys()))
            hw_lat = list(batch['lat_rgb'][base_key].shape[-2:])
            hw = tuple(hw_lat[i] * 8 for i in range(len(hw_lat)))
        else:
            return batch

        # Build global cameras from Tcw (cam-to-world) extrinsics.
        # cams: possibly re-anchored (from extrinsics, for plucker rays).
        # NOTE(bvh): reference_camera re-anchoring now happens pre-collate in anydata_dataset.py.
        batch['cams'] = self._build_cams_dict(extrinsics, intrinsics, hw)

        # cams_orig: always absolute poses (from extrinsics_orig, for visualization / action projection).
        if extrinsics_orig is not None and isinstance(extrinsics_orig, dict):
            batch['cams_orig'] = self._build_cams_dict(extrinsics_orig, intrinsics, hw)

        # Drop raw dicts -- everything goes through Camera objects from here on
        batch.pop('extrinsics', None)
        batch.pop('intrinsics', None)

        return batch

    @staticmethod
    def _build_cams_dict(extrinsics: dict, intrinsics: dict, hw) -> dict:
        '''
        Build a {(t, cam): Camera} dict from extrinsics, intrinsics, and hw.
        '''
        # Ensure [B,4,4] shape
        sample_key = next(iter(extrinsics))
        if extrinsics[sample_key].dim() == 2:
            extrinsics = {k: v.unsqueeze(0) for k, v in extrinsics.items()}

        cams = {}
        for key in extrinsics.keys():  # loop over (time, cam) tuples
            
            if key[1] == 'ego':
                # DEPRECATED(bvh): Ego camera system. Use reference_camera=<cam_name> in the
                # dataset YAML instead (e.g. reference_camera: 'CAMERA_01' for driving datasets).
                # This makes the named camera the identity reference pre-collate, replacing the
                # need for a separate ego camera entry.
                K_val = None
                hw_val = (1, 1)
            
            else:
                K_val = intrinsics.get(key, intrinsics.get((0, key[1]), None))
                hw_val = hw[key] if isinstance(hw, dict) else hw
            
            K_val = sanitize_intrinsics(K_val, hw_val, key=key)
            # Guard NaN/inf poses with an identity fallback (a valid pose; zeroing would leave a
            # degenerate rotation). Mirrors sanitize_intrinsics' default-pinhole fallback.
            Tcw = extrinsics[key].float()
            ok = torch.isfinite(Tcw).flatten(-2).all(-1)  # (B,) per-pose
            if not ok.all():
                eye = torch.eye(4, dtype=Tcw.dtype, device=Tcw.device).expand_as(Tcw)
                Tcw = torch.where(ok[..., None, None], Tcw, eye)
                log.warning(f'Non-finite extrinsics at {key} ({int((~ok).sum())}/{len(ok)} poses), '
                            f'using identity fallback', rank0_only=False)
            cams[key] = Camera(K=K_val, Tcw=Tcw, hw=hw_val)
        
        return cams

    @staticmethod
    def _clip_depth(batch: dict, max_depth: float) -> dict:
        '''
        Clip depth values to max_depth.
        '''
        for key in batch['depth']:
            batch['depth'][key] = torch.clamp(batch['depth'][key], max=max_depth)
        return batch

    @staticmethod
    def _scale_depth(batch: dict, scale: float) -> dict:
        '''
        Scale depth values by a factor.
        '''
        for key in batch['depth']:
            batch['depth'][key] = batch['depth'][key] * scale
        return batch