# Created by BVH & VG, Jul 2025.
# Simple single-view forecast-only dataloader for the VIDAR (legacy) pipeline.
# Reads from data_dict['vidar'] and assigns Any4D entries for a single camera.
#
# Supported tasks: forecast only (no cross-modal, DVS, pose estimation, action, or trajectory).
# Limitations:
#   - Single viewpoint only (V_max == 1, always uses the first available camera).
#   - No task dice rolling; always performs forecasting.
#   - No directives support.
#   - Hardcoded conditioning frame choices: [1, 5, 9, 13, 17] (train), 9 (val).
# For multi-view / multi-task support, use vidar_flex.py instead.
# NOTE: This file is deprecated in favor of unified_flex.py for new datasets.

import numpy as np
import torch
from rich import print

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


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

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

    V_max = config.num_views if hasattr(config, 'num_views') else 1  # Number of viewpoints.
    assert V_max == 1, 'Only one viewpoint is supported in vidar_default. Use vidar_flex for multi-view.'

    # 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.
    latent_cams = vidar_dict.get('lat_cams', None) if load_cams else None        # dict mapping (time, camera) to (B, 32/48, H, W) tensor of float.
    
    # Available keys & device
    first_valid = (vidar_rgb or vidar_depth or vidar_points or vidar_cams or \
                   latent_rgb or latent_depth or latent_points or latent_cams)
    assert first_valid is not None, 'Insufficient input information'
    keys = list(first_valid.keys())
    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)
    
    # Simply select first viewpoint
    cam_idx = avail_cams[0]

    # Perform forecasting 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)
    if vidar_rgb is not None or latent_rgb is not None:
        (a4d_raw, a4d_latent) = prepare_dense(
            a4d_raw, a4d_latent, 'rgb0', vidar_rgb, latent_rgb, cam_idx, cond_frames_latent)
        # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [0, 1] + masks.

    if vidar_depth is not None or latent_depth is not None:
        (a4d_raw, a4d_latent) = prepare_dense(
            a4d_raw, a4d_latent, 'depth0', vidar_depth, latent_depth, cam_idx, cond_frames_latent,
            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 vidar_points is not None or latent_points is not None:
        (a4d_raw, a4d_latent) = prepare_dense(
            a4d_raw, a4d_latent, 'points0', vidar_points, latent_points, cam_idx, cond_frames_latent,
            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.

    if vidar_cams is not None or latent_cams is not None:
        (a4d_raw, a4d_latent) = prepare_dense(
            a4d_raw, a4d_latent, 'cams0', vidar_cams, latent_cams, cam_idx, cond_frames_latent,
            scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0)
        # ^ (B, 6/9/32/48, T, H, W) tensor of bfloat16 in [-1, 1] + masks.

    # 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': str(cam_idx),
        ####
        '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,
    }
    
    return data_dict