# Created by BVH, May 2026.
# Modernized from custom/dataloader/old/debug/dvs_adaln.py to match the
# directives / harmonization / metadata conventions of vidar_flex.py.
# Encodes camera extrinsics as a low-dim adaln entry called 'relcams'
# (per-timestep view 0 + view 1 Tcw matrices, both relative to (view 0, time 0)).
# DVS-only flow: no action modality, no pose_estim task (camera is adaln-conditioned,
# never predicted).

import numpy as np
import torch
from rich import print

from custom.dataloader.utils import prepare_dense, random_choice, random_choice_k_unique, random_value
from custom.utils.constants import VAE_RATIO_THW
from imaginaire.utils import log


def _compute_relcams(vidar_cams, used_cams, V_used, T_used, T_max, B, device):
    """Build per-timestep relative extrinsics for view 0 and view 1.

    For each timestep t in [0, T_used), stores 24 channels:
        ch  0:12 = top 3 rows of view 0 Tcw at time t, relative to (view 0, time 0)
        ch 12:24 = top 3 rows of view 1 Tcw at time t, relative to (view 0, time 0)
    Both relative to the same reference (view 0, time 0); view 0's first 12 ch
    at t=0 is identity. Timesteps in [T_used, T_max) are zero-padded; the
    accompanying mask marks which timestep tokens are populated.

    Requires V_used >= 2 (DVS task). For V_used > 2, only the first two views
    are encoded -- 2-view is the supported setting here.

    :return (relcams, relcams_mask): (B, T_max, 24) and (B, T_max, 1) bf16 tensors.
    """
    relcams = torch.zeros(B, T_max, 24, device=device, dtype=torch.bfloat16)
    relcams_mask = torch.zeros(B, T_max, 1, device=device, dtype=torch.bfloat16)

    if V_used < 2:
        return relcams, relcams_mask

    ref_cam = used_cams[0]
    other_cam = used_cams[1]

    ref_Tcw_t0 = vidar_cams[(0, ref_cam)].Tcw.T  # (B, 4, 4) view 0 at time 0
    ref_inv = torch.linalg.inv(ref_Tcw_t0.float()).to(ref_Tcw_t0.dtype)

    for t in range(T_used):
        Tcw_v0 = vidar_cams[(t, ref_cam)].Tcw.T  # (B, 4, 4)
        Tcw_v1 = vidar_cams[(t, other_cam)].Tcw.T  # (B, 4, 4)
        rel_v0 = (Tcw_v0 @ ref_inv)[:, 0:3, :].reshape(B, -1)  # (B, 12)
        rel_v1 = (Tcw_v1 @ ref_inv)[:, 0:3, :].reshape(B, -1)  # (B, 12)
        relcams[:, t, 0:12] = rel_v0.to(torch.bfloat16)
        relcams[:, t, 12:24] = rel_v1.to(torch.bfloat16)
        relcams_mask[:, t, :] = 1.0

    return relcams, relcams_mask


def vidar_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None, **leftover):

    # Valid phases
    assert phase in ['train', 'val', 'test']

    if directives is None:
        directives = dict()  # to avoid crashes later

    # load_modals represents M_max, but always use all
    # underlying data represents T_max, but we always use all
    # Redefine V_max based on ENTRIES rather than STREAMS since concatenation might be involved.
    V_max = 1 + max([int(x[3:]) for x in config.all_highdim_entries
                     if x.startswith('rgb') and '_mask' not in x])

    # Initialize dictionaries to populate & return.
    a4d_raw = dict()
    a4d_latent = dict()

    load_rgb = ('rgb' in config.load_modals)
    load_depth = ('depth' in config.load_modals)
    load_points = ('points' in config.load_modals)
    load_cams = ('cams' in config.load_modals)

    # Get vidar dictionary (this already includes some model-side transforms)
    vidar_dict = data_dict['vidar']
    scale_factors = data_dict['scale_factor']  # (B) tensor of float

    # Extract relevant raw / pixel modalities.
    vidar_rgb = vidar_dict.get('rgb', None) if load_rgb else None           # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1].
    vidar_depth = vidar_dict.get('depth', None) if load_depth else None     # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, +inf).
    vidar_points = vidar_dict.get('points', None) if load_points else None  # dict mapping (time, camera) to (B, 3, H, W) tensor of float in (-inf, +inf).
    vidar_cams = vidar_dict.get('cams', None) if load_cams else None        # dict mapping (time, camera) to CameraPinhole object.

    # Extract relevant latent modalities.
    latent_rgb = vidar_dict.get('lat_rgb', None) if load_rgb else None           # dict mapping (time, camera) to (B, 16, H, W) tensor of float.
    latent_depth = vidar_dict.get('lat_depth', None) if load_depth else None     # dict mapping (time, camera) to (B, 16, H, W) tensor of float.
    latent_points = vidar_dict.get('lat_points', None) if load_points else None  # dict mapping (time, camera) to (B, 16, H, W) tensor of float.

    # Available keys & device
    first_valid = (vidar_rgb or vidar_depth or vidar_points or \
                   latent_rgb or latent_depth or latent_points)
    assert first_valid is not None, 'Insufficient input information'
    keys = list(first_valid.keys())
    B = first_valid[keys[0]].shape[0]
    device = first_valid[keys[0]].device

    # Available timesteps & cameras
    avail_times = sorted(list(set([key[0] for key in vidar_dict['timestep'].keys()])))
    avail_cams = sorted(list(set([key[1] for key in vidar_dict['timestep'].keys()])))
    T_data = len(avail_times)
    V_data = len(avail_cams)

    # Available modalities (except special cases: camera, language)
    avail_modals = []
    avail_modals += ['rgb'] if vidar_rgb is not None or latent_rgb is not None else []
    avail_modals += ['depth'] if vidar_depth is not None or latent_depth is not None else []
    avail_modals += ['points'] if vidar_points is not None or latent_points is not None else []
    M_data = len(avail_modals)

    # Get task probabilities (DVS-only flow: no pose_estim, no action tasks)
    p_xm = config.task_probs.get('cross_modal', -0.1)
    p_dvs = config.task_probs.get('dyn_view_synth', -0.1)
    p_fc = config.task_probs.get('forecast', -0.1)

    # Prepare to roll the dice
    bernoulli_fn = lambda p: random_value(device, generator) < p
    rnd_choice_fn = lambda choices: random_choice(choices, device, generator)

    # Determine tokens used either for input or output
    M_used = M_data  # Always use all available
    if phase == 'train':
        if config.use_views == 'rand1':
            V_used = rnd_choice_fn(list(range(1, min(V_data, V_max) + 1)))  # Sometimes ignore some
        elif config.use_views == 'rand2':
            V_used = rnd_choice_fn(list(range(2, min(V_data, V_max) + 1)))  # Sometimes ignore some
        elif config.use_views == 'all':
            V_used = min(V_data, V_max)
    else:
        V_used = min(V_data, V_max)

    T_used = T_data  # Always use all available

    # Determine active modalities in this batch (either for input or output)
    if 'used_modals' in directives:
        used_modals = directives['used_modals']
    else:
        used_modals = random_choice_k_unique(avail_modals, M_used, device, generator)

    # Determine active cameras in this batch (either for input or output)
    # NOTE(bvh): Make sure to respect these indices when we visualize stuff later!
    if 'used_cams' in directives:
        used_cams = directives['used_cams']
    elif config.shuffle_cams2views:
        # NOTE(bvh): Even if V_used == V_data, this random subsampling step is important to ensure
        # debiasing between Any4D viewpoints and original dataset cameras.
        used_cams = random_choice_k_unique(avail_cams, V_used, device, generator)
    else:
        # NOTE(bvh): This is the deterministic variant, useful if we are using all cameras
        # in the webdataset and want to link them to Any4D views in a consistent way.
        used_cams = avail_cams[0:V_used]

    if 'tasks' in directives:

        # supports both str or list of str
        is_cross_modal = 'cross_modal' in directives['tasks']
        is_dyn_view_synth = 'dyn_view_synth' in directives['tasks']
        is_forecast = 'forecast' in directives['tasks']

    else:

        # Imagine everything is given as input as a starting point.
        # Then we start gradually masking out random things.
        # We must perform at least one task, so if none are selected, we keep retrying.
        is_cross_modal = False
        is_dyn_view_synth = False
        is_forecast = False

        for dice_roll in range(1010):
            # Remove one or more modalities (except camera)?
            is_cross_modal = (bernoulli_fn(p_xm) and M_used >= 2)
            # Remove one entire viewpoint (except camera)?
            is_dyn_view_synth = (bernoulli_fn(p_dvs) and vidar_cams is not None and V_used >= 2)
            # Remove most future frames (except camera)?
            is_forecast = bernoulli_fn(p_fc)

            valid_tasks = is_cross_modal or is_dyn_view_synth or is_forecast
            if valid_tasks:
                break

            if dice_roll >= 1000:
                raise ValueError('Failed to select a valid set of tasks after 1000 dice rolls')

    # Then determine the INPUT modalities and viewpoints specifically
    if is_cross_modal:
        # Select strict, non-empty subset of modalities as input, the complement is output.
        if 'cond_modals' in directives:
            cond_modals = directives['cond_modals']  # list of str

        else:
            cond_modals = []

            for dice_roll in range(1010):
                cond_modals = [modality for modality in used_modals if bernoulli_fn(0.5)]

                if 0 < len(cond_modals) < M_used:
                    break

                if dice_roll >= 1000:
                    raise ValueError('Failed to select a valid set of modalities after 1000 dice rolls')

    else:
        cond_modals = used_modals

    if is_dyn_view_synth:
        # Select strict, non-empty subset of viewpoints as input, the complement is output.
        # NOTE(bvh): For DVS, cams is always input via adaln; only the non-cams modalities
        # of the output view get masked out (cond_views_notcams = cond_views).
        if 'num_pred_views' in directives:
            cond_views = V_used - directives.get('num_pred_views', 1)
        else:
            cond_choices_views = directives.get('cond_choices_views', list(range(1, V_used)))
            cond_views = random_choice(cond_choices_views, device, generator)

        cond_views_notcams = cond_views
    else:
        cond_views_notcams = V_used

    if is_forecast:
        if 'cond_frames_raw' in directives:
            cond_frames_raw = directives['cond_frames_raw']

        else:
            if 'cond_frames_choices' in directives:
                cond_frames_choices = directives['cond_frames_choices']

            else:
                if phase == 'train':
                    # Select strict, non-empty subset of frames as input, but at most ~66%.
                    # e.g. [1, 5, 9, 13, 17, 21, 25] if Tp = 41
                    max_cond_frames = (T_used * 2) // 3
                else:
                    # Select strict, non-empty subset of frames as input, but at most ~33%.
                    # e.g. [1, 5, 9, 13] if Tp = 41
                    # Use fewer conditioning frames during validation/test to better stress the model's forecasting capability.
                    max_cond_frames = (T_used + 1) // 3

                cond_frames_choices = list(range(1, max_cond_frames, VAE_RATIO_THW[0]))

            cond_frames_raw = random_choice(cond_frames_choices, device, generator)

    else:
        cond_frames_raw = T_used

    pred_frames_raw = T_used
    cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1
    pred_frames_latent = (pred_frames_raw - 1) // VAE_RATIO_THW[0] + 1

    # Sanity check: at least one task must be active
    assert is_cross_modal or is_dyn_view_synth or is_forecast, \
        'No tasks are selected (or no tokens are output), probably a bug / mistake in directives?'

    # Populate Any4D dictionaries with high-dim entries (rgb / depth / points only -- NO cams here).
    # Loop over all video viewpoints & modalities to process them.
    # NOTE(bvh): Conditional (input) views intentionally start from the LAST one.
    for v in range(V_used):

        if 'rgb' in used_modals:
            cur_input = ('rgb' in cond_modals and v >= V_used - cond_views_notcams)
            cond_frames_rgb = cond_frames_latent if cur_input else 0
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'rgb{v}', vidar_rgb, latent_rgb,
                used_cams[v], cond_frames_rgb)
            # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [0, 1] + masks.

        if 'depth' in used_modals:
            cur_input = ('depth' in cond_modals and v >= V_used - cond_views_notcams)
            cond_frames_depth = cond_frames_latent if cur_input else 0
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'depth{v}', vidar_depth, latent_depth,
                used_cams[v], cond_frames_depth, scale_factor=scale_factors, clip_min=0.0, clip_max=1.0)
            # ^ (B, 1/16, T, H, W) tensor of bfloat16 in [0, 1] + masks.

        if 'points' in used_modals:
            cur_input = ('points' in cond_modals and v >= V_used - cond_views_notcams)
            cond_frames_points = cond_frames_latent if cur_input else 0
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'points{v}', vidar_points, latent_points,
                used_cams[v], cond_frames_points, scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0)
            # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [-1, 1] + masks.

    # Camera extrinsics as a low-dim adaln entry (replaces high-dim cams{v} entries).
    # T_max tokens (= raw frame count), 24 channels per token = view 0 + view 1 Tcw 3x4
    # matrices, both relative to (view 0, time 0). Tokens beyond T_used are zero-padded.
    if vidar_cams is not None and 'relcams' in config.lowdim_adaln_entries:
        relcams_meta = config.lowdim_adaln_entries['relcams']  # (seq_start, seq_end, ch_start, ch_end, init)
        T_max = relcams_meta[1] - relcams_meta[0]  # max adaln timesteps
        assert T_used <= T_max, f'T_used = {T_used} must be <= T_max = {T_max} for relcams adaln entry'

        relcams, relcams_mask = _compute_relcams(
            vidar_cams, used_cams, V_used, T_used, T_max, B, device)

        a4d_latent['relcams'] = relcams
        a4d_latent['relcams_input_mask'] = relcams_mask
        # ^ (B, T_max, 24) tensor of bfloat16 + (B, T_max, 1) mask.

    # This will be processed by Any4D VAE as needed later:
    data_dict['a4d_raw'] = a4d_raw
    data_dict['a4d_latent'] = a4d_latent

    # Keep track of metadata, including selected modalities, viewpoints, tasks, frames, etc.
    # This is informative only (not used by the model).
    meta_tasks = []
    meta_tasks += ['cross_modal'] if is_cross_modal else []
    meta_tasks += ['dyn_view_synth'] if is_dyn_view_synth else []
    meta_tasks += ['forecast'] if is_forecast else []
    data_dict['meta'] = {
        'used_modals': ','.join(used_modals),  # str because variable length
        'used_cams': ','.join(str(cam) for cam in used_cams),  # str because variable length
        ####
        'tasks': ','.join(meta_tasks),
        ####
        'cond_modals': ','.join(cond_modals),
        'cond_views_cams': cond_views,
        'cond_views_notcams': cond_views_notcams,
        'cond_frames_raw': cond_frames_raw,
        'pred_frames_raw': pred_frames_raw,
        'cond_frames_latent': cond_frames_latent,
        'pred_frames_latent': pred_frames_latent,
    }

    return data_dict