# Created by BVH, Aug 2025.
# Tutorial / reference dataloader showing how data-to-Any4D conversion works internally.
# Multi-view, RGB-only, forecast-only. Manually implements what prepare_dense() does
# to serve as a readable example of the raw mechanics.
#
# Supported tasks: forecast only (no cross-modal, DVS, pose estimation, action, or trajectory).
# Limitations:
#   - RGB modality only (no depth, points, or cams).
#   - No task dice rolling; always performs forecasting.
#   - No directives support.
#   - Hardcoded conditioning frame choices: [1, 5, 9, 13, 17] (train), 9 (val).
#   - Reads from data_dict['anydata'] (unified pipeline).
# For full multi-modal / multi-task support, use unified_flex.py or vidar_flex.py.

import numpy as np
import torch
from rich import print

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


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

    V_max = config.num_views  # what is supported by the model
    
    # Initialize dictionaries to populate & return.
    a4d_raw = dict()
    a4d_latent = dict()

    load_rgb = ('rgb' in config.load_modals)

    # Get data dictionary (this already includes some model-side transforms)
    anydata_dict = data_dict['anydata']

    # Extract relevant raw / pixel modalities.
    anydata_rgb = anydata_dict.get('rgb', None) if load_rgb else None           # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1].

    # Extract relevant latent modalities.
    latent_rgb = anydata_dict.get('lat_rgb', None) if load_rgb else None           # dict mapping (time, camera) to (B, 16, H, W) tensor of float.
    # Available keys & device
    first_valid = (anydata_rgb or latent_rgb)
    assert first_valid is not None, 'Insufficient input information'
    keys = list(first_valid.keys())
    device = first_valid[keys[0]].device
    tensor_kwargs = dict(device=device, dtype=torch.bfloat16)

    # Available timesteps & cameras
    avail_times = sorted(list(set([key[0] for key in anydata_dict['timestep'].keys()])))
    avail_cams = sorted(list(set([key[1] for key in anydata_dict['timestep'].keys()])))
    T_data = len(avail_times)
    V_data = len(avail_cams)
    
    V_used = min(V_data, V_max)  # we could also randomly sample this
    used_cams = [avail_cams[v] for v in range(V_used)]  # we could also shuffle this

    # Perform forecasting for all viewpoints with variable number of input frames
    if phase == 'train':
        cond_frames_raw = random_choice([1, 5, 9, 13, 17], device, generator)
    elif phase == 'val':
        cond_frames_raw = 9
    else:
        raise ValueError(f'Invalid phase {phase}')
    
    pred_frames_raw = T_data
    cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1
    pred_frames_latent = (pred_frames_raw - 1) // VAE_RATIO_THW[0] + 1

    # Populate Any4D dictionaries with entries (either raw or latent)
    for v in range(V_used):
        cam_idx = used_cams[v]
        
        # In other dataloaders, this part can get very boilerplate and is therefore handled by
        # methods like prepare_dense(), but in this example we implement the equivalent
        # specific case here for simplicity. Note that we support both raw (pixel) & latent
        # (cached / pre-encoded) webdataset formats as input.
        if 0:
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'rgb{v}', anydata_rgb, latent_rgb, cam_idx, cond_frames_latent)
        
        # ^ this method call is equivalent to the following:
        if latent_rgb is not None:
            # We received cached latent data only
            a4d_latent[f'rgb{v}'] = latent_rgb[(0, cam_idx)]  # contains input and ground truth together
            a4d_latent[f'rgb{v}'] = a4d_latent[f'rgb{v}'].detach().clone().to(**tensor_kwargs)
            
            (B, Cl, Tl, Hl, Wl) = a4d_latent[f'rgb{v}'].shape
            assert Cl == 16
            Tp = (Tl - 1) * VAE_RATIO_THW[0] + 1
            Hp = Hl * VAE_RATIO_THW[1]
            Wp = Wl * VAE_RATIO_THW[2]
            
        else:
            # We received raw pixel data only
            a4d_raw[f'rgb{v}'] = [anydata_rgb[(t, cam_idx)] for t in avail_times]
            a4d_raw[f'rgb{v}'] = torch.stack(a4d_raw[f'rgb{v}'], dim=2)
            a4d_raw[f'rgb{v}'] = a4d_raw[f'rgb{v}'].detach().clone().to(**tensor_kwargs)
            
            (B, Cp, Tp, Hp, Wp) = a4d_raw[f'rgb{v}'].shape
            assert Cp == 3
            Tl = (Tp - 1) // VAE_RATIO_THW[0] + 1
            Hl = Hp // VAE_RATIO_THW[1]
            Wl = Wp // VAE_RATIO_THW[2]

        # In both cases, we must also define the per-token masks
        # (in DiT / latent space only to keep VRAM overhead negligible):
        a4d_latent[f'rgb{v}_input_mask'] = torch.zeros(B, 1, Tl, Hl, Wl, **tensor_kwargs)
        a4d_latent[f'rgb{v}_output_mask'] = torch.zeros(B, 1, Tl, Hl, Wl, **tensor_kwargs)
        a4d_latent[f'rgb{v}_input_mask'][:, :, 0:cond_frames_latent] = 1.0
        a4d_latent[f'rgb{v}_output_mask'][:, :, cond_frames_latent:pred_frames_latent] = 1.0

        # Optionally, we can also assign supervise_mask,
        # but by default they are set to be copies of output_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:
    data_dict['meta'] = {
        'used_cams': ','.join(str(cam) for cam in used_cams),  # str because possibly variable length
        ####
        'tasks': 'forecast',
        ####
        'cond_frames_raw': cond_frames_raw,
        'pred_frames_raw': pred_frames_raw,
        'cond_frames_latent': cond_frames_latent,
        'pred_frames_latent': pred_frames_latent,
    }

    # output print example with batch_size = 1:
    # data_dict = {
    #   'dl_idx': tensor[1] i64 cuda:0 [0],
    #   'anydata': {...},
    #   'dset_name': ['pdnvs25'], 
    #   'sample_id': ['25_camera__scene_000005/seq-000000-0_50.tar_3'], 
    #   'fps': tensor[1] bf16 cuda:0 [10.000], 
    #   'prompt': ['drive the car around'], 
    #   'scale_factor': tensor[1] bf16 cuda:0 [0.008], 
    #   '...',
    #   'a4d_raw': {
    #       'rgb0': tensor[1, 3, 25, 288, 384] bf16 n=8294400 (16Mb) x∈[0., 1.000] μ=0.139 σ=0.038 cuda:0, 
    #       'rgb1': tensor[1, 3, 25, 288, 384] bf16 n=8294400 (16Mb) x∈[0., 1.000] μ=0.141 σ=0.052 cuda:0
    #   }, 
    #   'a4d_latent': {
    #       'rgb0_input_mask': tensor[1, 1, 7, 36, 48] bf16 n=12096 (24Kb) x∈[0., 1.000] μ=0.428 σ=0.494 cuda:0, 
    #       'rgb0_output_mask': tensor[1, 1, 7, 36, 48] bf16 n=12096 (24Kb) x∈[0., 1.000] μ=0.570 σ=0.494 cuda:0, 
    #       'rgb1_input_mask': tensor[1, 1, 7, 36, 48] bf16 n=12096 (24Kb) x∈[0., 1.000] μ=0.428 σ=0.494 cuda:0, 
    #       'rgb1_output_mask': tensor[1, 1, 7, 36, 48] bf16 n=12096 (24Kb) x∈[0., 1.000] μ=0.570 σ=0.494 cuda:0
    #   }, 
    #   'meta': {
    #       'used_cams': '0,1', 
    #       'tasks': 'forecast', 
    #       'cond_frames_raw': 9, 
    #       'pred_frames_raw': 25, 
    #       'cond_frames_latent': 3, 
    #       'pred_frames_latent': 7
    #   }
    # }
    
    return data_dict