# Created by BVH, Apr 2026.
# Dataloader for DRIVING datasets in the UNIFIED (anydata) pipeline.
# Based on unified_flex (trajectory) and unified_anyact (simplified task logic).
# Supports: forecast, world model (trajectory conditioning).
# Does NOT support: cross_modal, pose_estimation, inverse_dynamics, policy, robot action.
# Handles different driving datasets with different # cameras via dataset-specific cam mapping.

import numpy as np
import torch
from rich import print

from custom.dataloader.utils import prepare_dense, prepare_traj_anydata, random_choice, random_value, reanchor_cameras_to_view
from custom.dataloader.wm_utils import compute_trajectory_from_extrinsics, perturb_traj_basile1, perturb_traj_basile2, perturb_traj_basile3
from custom.utils.constants import VAE_RATIO_THW
from imaginaire.utils import log


# Each driving dataset camera rig maps to 4 directional groups: front (=ego), left, right, back.
# Each group is a priority-ordered list of camera names.
# View assignment: [rgb0, rgb1, rgb2, ...] =
# [front[0] (ego), left[0], right[0], back[0], left[1], right[1], back[1], ...]
# depending on availability (skip missing) + V_used in this batch.
DRIVING_CAM_LAYOUT = {
    # ArgoVerse2sync: 7 cameras
    # TODO verify names
    'argoverse2sync': {
        'front': ['ring_front_center'],
        'left':  ['ring_front_left', 'ring_side_left', 'ring_rear_left'],
        'right': ['ring_front_right', 'ring_side_right', 'ring_rear_right'],
        'back':  [],
    },
    # DDAD: 6 cameras: CAMERA_01=front, 05=front_left, 06=front_right,
    # 07=rear_left, 08=rear_right, 09=rear.
    'ddad': {
        'front': ['CAMERA_01'],
        'left':  ['CAMERA_05', 'CAMERA_07'],
        'right': ['CAMERA_06', 'CAMERA_08'],
        'back':  ['CAMERA_09'],
    },
    # GAIA: 11 cameras total (parking scenarios, fisheye).
    # Both standard and PVM configs use this single layout; _order_cameras
    # matches whichever cameras are present in the batch.
    # Front priority: front_wide (standard ego), front_pvm (PVM ego), front_tele (aux).
    'gaia': {
        'front': ['front_wide', 'front_pvm', 'front_tele'],
        'left':  ['left_side_forward', 'left_side_pvm', 'left_side_rearward'],
        'right': ['right_side_forward', 'right_side_pvm', 'right_side_rearward'],
        'back':  ['rear_tele', 'rear_pvm'],
    },
    # LyftL5: 7 cameras: 0=front, 5=front_left, 1=front_right
    # TODO check others
    'lyftl5': {
        'front': ['0'],
        'left':  ['5', '3'],
        'right': ['1', '2'],
        'back':  ['4', '6'],
    },
    # Waymo: 5 cameras
    'waymo': {
        'front': ['0'],
        'left':  ['1', '3'],   # FRONT_LEFT, SIDE_LEFT
        'right': ['2', '4'],   # FRONT_RIGHT, SIDE_RIGHT
        'back':  [],
    },
    # PD4D: 19 cameras total, but only 3 egocentric cameras are useful/realistic,
    # the others are exo views (inward facing from the air)
    'pd4d': {
        'front': ['yaw-0'],
        'left':  ['yaw-neg-60'],
        'right': ['yaw-60'],
        'back':  [],
    },
    # PDv2: 6 cameras: same as DDAD but lowercase
    'pdv2': {
        'front': ['camera_01'],
        'left':  ['camera_05', 'camera_07'],
        'right': ['camera_06', 'camera_08'],
        'back':  ['camera_09'],
    },
}


def _get_ego_camera(dset_name, cam_names):
    """Return integer index of the ego (front) camera for this dataset.
    :param dset_name: str, dataset name (e.g. 'DDAD', 'LyftL5').
    :param cam_names: list of str, camera names in index order, e.g. ['yaw-0', 'yaw-60', 'yaw-neg-60'].
    :return: int, camera index of the ego camera.
    """
    # TODO(bvh): debug/verify this more rigorously

    dset_key = dset_name.lower().split('_')[0]
    layout = DRIVING_CAM_LAYOUT.get(dset_key, None)
    assert layout is not None, f'No camera layout for dataset {dset_name!r} (key={dset_key!r})'

    name_to_idx = {name: idx for idx, name in enumerate(cam_names)}
    for ego_name in layout['front']:
        if ego_name in name_to_idx:
            return name_to_idx[ego_name]

    # NOTE(bvh): temporary fallback -- some webbed variants only contain non-front
    # cameras (e.g. IKE_24 DDAD has CAMERA_07/08/09 only). Warn and use first available.
    log.error(
        f'Ego camera not found: layout front={layout["front"]}, '
        f'available names={cam_names} (dataset={dset_name!r}). '
        f'Falling back to first available: {cam_names[0]!r}.', rank0_only=False)
    return 0


def _order_cameras(dset_name, cam_names):
    """Order camera indices: front (ego) first, then interleaved left/right/back. Clamp to V_max.
    Returns a list of integer camera indices (matching (t, c) dict keys).

    :param dset_name: str, dataset name.
    :param cam_names: list of str, camera names in index order.
    :return: list of int, camera indices in priority order.
    """
    # TODO(bvh): debug/verify this more rigorously

    dset_key = dset_name.lower().split('_')[0]
    layout = DRIVING_CAM_LAYOUT.get(dset_key, None)
    all_indices = list(range(len(cam_names)))
    if layout is None:
        return all_indices

    name_to_idx = {name: idx for idx, name in enumerate(cam_names)}
    ordered = []

    # Front (ego) first
    for cam_name in layout['front']:
        if cam_name in name_to_idx:
            idx = name_to_idx[cam_name]
            if idx not in ordered:
                ordered.append(idx)
    if not ordered:
        ordered.append(0)  # fallback: index 0

    # Interleave left, right, back by priority
    groups = [layout['left'], layout['right'], layout['back']]
    max_depth = max((len(g) for g in groups), default=0)
    for depth in range(max_depth):
        for group in groups:
            if depth < len(group):
                cam_name = group[depth]
                if cam_name in name_to_idx:
                    idx = name_to_idx[cam_name]
                    if idx not in ordered:
                        ordered.append(idx)

    # Append any remaining indices not covered by layout
    for idx in all_indices:
        if idx not in ordered:
            ordered.append(idx)

    return ordered


def anydata_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None,
                     train_step=-1, val_step=-1, **leftover):

    # TODO(bvh): use train_step and/or val_step to orchestrate
    # synchronized video length variations

    assert phase in ['train', 'val', 'test']

    if directives is None:
        directives = dict()

    # V_max from entries (not config.num_views, since concatenation might be involved)
    _entries_no_mask = [x for x in config.all_highdim_entries if '_mask' not in x]
    V_max = 1 + max([int(x[3:]) for x in _entries_no_mask if x.startswith('rgb')])
    has_rgb_entries = any(x.startswith('rgb') for x in _entries_no_mask)
    has_cams_entries = any(x.startswith('cams') for x in _entries_no_mask)

    a4d_raw = dict()
    a4d_latent = dict()

    anydata_dict = data_dict['anydata']
    scale_factors = data_dict['scale_factor']  # (B) tensor of float
    dset_names = data_dict['dset_name']  # list (B) of str
    dset_name = dset_names[0]

    # Extract modalities
    anydata_rgb = anydata_dict.get('rgb', None)
    anydata_cams = anydata_dict.get('cams', None)
    latent_rgb = anydata_dict.get('lat_rgb', None)
    latent_cams = anydata_dict.get('lat_cams', None)

    first_valid = anydata_rgb or latent_rgb
    assert first_valid is not None, 'Need at least rgb or lat_rgb'
    keys = list(first_valid.keys())
    B = first_valid[keys[0]].shape[0]
    device = first_valid[keys[0]].device

    # Available timesteps & camera indices.
    avail_times = sorted(set(key[0] for key in anydata_dict['timestep'].keys()))
    avail_cam_inds = sorted(set(key[1] for key in anydata_dict['timestep'].keys()))
    T_data = len(avail_times)
    V_data = len(avail_cam_inds)

    # NOTE(bvh): using first batch element for now, but ordering may differ across batch elements.
    # This means that camera-to-view ordering will be "correct" for first sample,
    # but may be only approximately correct for remaining samples.
    # This is only an issue if cameras are being subsampled randomly.
    first_cam_names = anydata_dict['metadata']['cameras'][0]

    # Get task probabilities
    # Forecast must always happen, world model (trajectory conditioning) is optional
    p_fc = 1.0
    p_wm = config.task_probs.get('world_model', -0.1)

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

    # Determine tokens used either for input, or output, or neither
    # NOTE(bvh): We always use available cameras up to V_used, which is up to V_max,
    # in a procedural order per dataset (ego first).
    # This replaces the shuffle_cams2views / use_views logic from robotics dataloaders.

    V_used = len(first_cam_names)
    T_used = T_data

    # TODO(bvh): debug/verify this
    # _get_ego_camera and _order_cameras return integer indices for views.
    ego_cam = _get_ego_camera(dset_name, first_cam_names)  # e.g. 0
    used_cams = _order_cameras(dset_name, first_cam_names)  # e.g. [0, 2, 3, 1, 4]
    used_cams = used_cams[0:V_max]

    # Re-anchor cameras to a reference view if configured.
    # This makes plucker rays relative to whichever camera was assigned to the reference view.
    ref_view = getattr(config, 'reference_view', False)
    if ref_view not in (False, None) and anydata_cams is not None:
        ref_cam_idx = used_cams[ref_view]
        per_timestep = getattr(config, 'ref_view_rel_per_timestep', True)
        anydata_cams = reanchor_cameras_to_view(anydata_cams, ref_cam_idx, per_frame=per_timestep)

    # ---- Task selection: forecast + world model only ----

    if 'tasks' in directives:
        is_forecast = 'forecast' in directives['tasks']
        is_world_model = 'world_model' in directives['tasks']
    else:
        is_forecast = bernoulli_fn(p_fc)
        if not is_forecast:
            is_forecast = True  # Fallback: always forecast
        # World model (trajectory conditioning) sampled independently
        is_world_model = bernoulli_fn(p_wm)

    # Forecast: condition on first N frames, predict the rest
    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':
                    max_cond_frames = (T_used * 2) // 3
                else:
                    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

    total_frames_raw = T_used
    cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1
    total_frames_latent = (total_frames_raw - 1) // VAE_RATIO_THW[0] + 1

    log.debug(f"cond_frames_raw: {cond_frames_raw}, total_frames_raw: {total_frames_raw}, "
              f"cond_frames_latent: {cond_frames_latent}, total_frames_latent: {total_frames_latent}")

    # World model: trajectory is fully conditioned (input), not predicted
    if is_world_model:
        cond_frames_traj = T_used
        total_frames_traj = 0
    else:
        cond_frames_traj = 0
        total_frames_traj = 0

    assert is_forecast, 'Forecast task must be active for driving'

    # ---- Populate Any4D entries ----

    for v in range(V_used):
        # RGB
        if has_rgb_entries and (anydata_rgb is not None or latent_rgb is not None):
            cond_frames_rgb = cond_frames_latent
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'rgb{v}', anydata_rgb, latent_rgb,
                used_cams[v], cond_frames_rgb)

        # Cameras (plucker)
        if has_cams_entries and (anydata_cams is not None or latent_cams is not None):
            cond_frames_cams = total_frames_latent
            (a4d_raw, a4d_latent) = prepare_dense(
                a4d_raw, a4d_latent, f'cams{v}', anydata_cams, latent_cams,
                used_cams[v], cond_frames_cams, plucker_orig=False,
                scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0)

    # ---- Trajectory from ego camera extrinsics ----

    orig_traj = None
    if anydata_cams is not None:
        # TODO(bvh): debug/verify this

        # DEBUG: check ego camera position at t=0 to verify reference_camera re-anchoring
        if 0:
            ego_pos_t0 = anydata_cams[(0, ego_cam)].get_center()  # (B, 3)
            log.debug(f'ego_cam={ego_cam} ({first_cam_names[ego_cam]}), pos@t0={ego_pos_t0[0].tolist()}')  # DEBUG

        traj_multiplier = getattr(config, 'traj_multiplier', 1.0)
        trajectory = compute_trajectory_from_extrinsics(
            anydata_dict, avail_times, ego_cam, traj_multiplier)
        # trajectory: (B, T_used, 2)

        # Get expected token count from adaln config to pad shorter clips
        adaln_traj = config.lowdim_adaln_entries.get('traj', None)
        target_tokens = (adaln_traj[1] - adaln_traj[0]) if adaln_traj else None

        if 'perturb_traj' in directives:
            perturb_traj = directives['perturb_traj']
            perturb_seed = directives.get('perturb_seed', None)
            orig_traj = trajectory.clone()
            if perturb_traj == 'basile1':
                trajectory = perturb_traj_basile1(orig_traj, cond_frames_raw, seed=perturb_seed)
            elif perturb_traj == 'basile2':
                trajectory = perturb_traj_basile2(orig_traj, cond_frames_raw, seed=perturb_seed)
            elif perturb_traj == 'basile3':
                trajectory = perturb_traj_basile3(orig_traj, cond_frames_raw, seed=perturb_seed)

        # NOTE(bvh): trajectory is input-only for now (pred_frames=0).
        (a4d_raw, a4d_latent) = prepare_traj_anydata(
            a4d_raw, a4d_latent, trajectory, cond_frames_traj,
            pred_frames=total_frames_traj, target_tokens=target_tokens)

        if orig_traj is None:
            orig_traj = a4d_raw['traj'].clone()

    # Store results
    data_dict['a4d_raw'] = a4d_raw
    data_dict['a4d_latent'] = a4d_latent

    # Metadata
    meta_tasks = []
    meta_tasks += ['forecast'] if is_forecast else []
    meta_tasks += ['world_model'] if is_world_model else []
    data_dict['meta'] = {
        'used_modals': ','.join(
            m for m, v in [('rgb', anydata_rgb or latent_rgb),
                           ('cams', anydata_cams or latent_cams),
                           ('traj', orig_traj)] if v is not None),
        'used_cams': ','.join(str(cam) for cam in used_cams),  # a bit ugly but we should avoid changing standards
        'tasks': ','.join(meta_tasks),
        'cond_modals': 'rgb',  # TODO update
        'cond_views_cams': V_used,
        'cond_views_notcams': V_used,
        'cond_frames_raw': cond_frames_raw,
        'total_frames_raw': total_frames_raw,
        'cond_frames_latent': cond_frames_latent,
        'total_frames_latent': total_frames_latent,
        'cond_frames_traj': cond_frames_traj,
        'total_frames_traj': total_frames_traj,
    }

    if orig_traj is not None:
        data_dict['meta']['orig_traj'] = orig_traj

    return data_dict