# Created by BVH & VG, Jul 2025.

import numpy as np
import torch
from rich import print

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


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

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

    # load_modals represents M_max, but always use all
    V_max = config.num_views  # Maximum number of views supported by the current model config.
    # underlying data represents T_max, but we always use all

    # 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_cams = ('cams' in config.load_modals)
    load_action = ('action' in config.load_modals)

    # TODO apply lbm_real_modals, or generalize this setting

    # 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_cams = vidar_dict.get('cams', None) if load_cams else None        # dict mapping (time, camera) to CameraPinhole object.
    vidar_action = vidar_dict.get('action', None) if load_action else None  # dict mapping (time, 0) to {'action', 'proprio': (B, 16, 20) tensor of float}.

    # 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_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_cams or \
                   latent_rgb or latent_depth or latent_cams)
    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, action)
    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 []
    M_data = len(avail_modals)

    # Get task probabilities
    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)
    p_id = config.task_probs.get('inverse_dynamics', -0.1)
    p_p = config.task_probs.get('policy', -0.1)
    p_wm = config.task_probs.get('world_model', -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 config.use_views == 'rand':
        V_used = rnd_choice_fn(list(range(1, min(V_data, V_max) + 1)))  # Sometimes ignore some
    elif config.use_views == 'all':
        V_used = min(V_data, V_max)
    T_used = T_data  # Always use all available

    # NOTE(bvh): Even if V_used = V_data, this step is important to ensure debiasing
    # NOTE(bvh): Make sure to respect these indices when we visualize stuff later!
    used_modals = random_choice_k_unique(avail_modals, M_used, device, generator)
    used_cams = random_choice_k_unique(avail_cams, V_used, device, generator)
    # ^ This will randomize the ordering when going from cameras to viewpoints as well, which is desired.

    # 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

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

    if is_cross_modal:
        # Select strict, non-empty subset of modalities as input, the complement is output.
        cond_modals = []
        while len(cond_modals) in [0, M_used]:
            cond_modals = [modality for modality in used_modals if bernoulli_fn(0.5)]
    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): Output means either cams is still input but predict all modalities (DVS),
        # or some modalities are input but predict cams (PE).
        cond_choices = list(range(1, V_used))
        cond_views = random_choice(cond_choices, device, generator)

    if is_forecast:
        # 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
        cond_choices = list(range(1, (T_used * 2) // 3, VAE_RATIO_THW[0]))
        cond_frames_raw = random_choice(cond_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

    # This action part works quite differently from the above:
    # Only one of these can be active at a time, but nothing is also possible!
    is_inverse_dynamics = bernoulli_fn(p_id)
    is_policy = bernoulli_fn(p_p / (1.0 - p_id) if p_id < 1.0 else 0.0)
    is_world_model = bernoulli_fn(p_wm / (1.0 - p_id - p_p) if p_id + p_p < 1.0 else 0.0)
    
    if is_inverse_dynamics:
        cond_frames_action = 0
        pred_frames_action = T_used
    elif is_policy:
        cond_frames_action = cond_frames_raw
        pred_frames_action = T_used
    elif is_world_model:  # Aka forward dynamics
        cond_frames_action = T_used
        pred_frames_action = 0
    else:
        cond_frames_action = 0
        pred_frames_action = 0

    # Populate Any4D dictionaries with entries (either raw or latent).
    # 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)
            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.

            # DEBUG / TEMP: apply clear visual binary indicators
            IS = 16
            H = a4d_raw[f'rgb{v}'].shape[-2]
            a4d_raw[f'rgb{v}'][:, :, :, H - V_max*IS:, :] = 0.0
            for b in range(B):
                for (i, v2) in enumerate(range(V_used)):
                    dset_cam_idx = int(vidar_dict['time_cam'][(0, used_cams[v2])][b, 1].item())
                    a4d_raw[f'rgb{v}'][b, :, :, H - (i+1)*IS:H - i*IS,
                                                dset_cam_idx*IS:(dset_cam_idx+1)*IS] = 1.0

        if 'depth' in used_modals:
            cur_input = ('depth' in cond_modals and v >= V_used - cond_views)
            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.

            # DEBUG / TEMP: apply clear visual binary indicators
            IS = 16
            H = a4d_raw[f'depth{v}'].shape[-2]
            a4d_raw[f'depth{v}'][:, :, :, H - V_max*IS:, :] = 0.0
            for b in range(B):
                for (i, v2) in enumerate(range(V_used)):
                    dset_cam_idx = int(vidar_dict['time_cam'][(0, used_cams[v2])][b, 1].item())
                    a4d_raw[f'depth{v}'][b, :, :, H - (i+1)*IS:H - i*IS,
                                                dset_cam_idx*IS:(dset_cam_idx+1)*IS] = 1.0

    if vidar_action is not None:
        entry_input = 'proprio' if 'proprio' in config.all_lowdim_entries else 'action'
        vidar_input = 'proprio' if 'proprio' in vidar_action[keys[0]] else 'action'
        (a4d_raw, a4d_latent) = prepare_action_vidar(
            a4d_raw, a4d_latent, vidar_action, cond_frames_action, pred_frames=pred_frames_action,
            entry_input=entry_input, entry_output='action',
            vidar_input=vidar_input, vidar_output='action')
        # ^ 2x (B, Tp, 20) tensor of float + 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:
    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 []
    meta_tasks += ['inverse_dynamics'] if is_inverse_dynamics else []
    meta_tasks += ['policy'] if is_policy else []
    meta_tasks += ['world_model'] if is_world_model 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': cond_views,
        'cond_frames_raw': cond_frames_raw,
        'pred_frames_raw': pred_frames_raw,
        'cond_frames_latent': cond_frames_latent,
        'pred_frames_latent': pred_frames_latent,
        'cond_frames_action': cond_frames_action,
        'pred_frames_action': pred_frames_action,
    }

    return data_dict