# Created by BVH, Jul - Aug 2025.
# Implements conversion between data units ("entries", such as the video tensor of one particular modality from one particular viewpoint) and architectural units ("streams", such as all video tokens from one viewpoint, or all cross-attention tokens).

from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import torch

from custom.utils.constants import VAE_RATIO_THW
from imaginaire.utils import log


def rgb_latent_span(view_entries):
    """Channel span (start, end) of the generated RGB latent within a per-view stream.

    A per-view stream packs the target rgb latent plus conditioning channels (input/output
    masks, reference latent, control latents like canny). Structured noise should touch only
    the generated rgb latent, so return its [start:end). Picks the entry named like 'rgb<N>'
    with no mask/ref/control suffix; falls back to the first entry (rgb is packed first).
    """
    for nm, c in view_entries.items():
        if nm.startswith('rgb') and not (nm.endswith(('_mask', '_ref')) or 'control' in nm):
            return int(c[0]), int(c[1])
    first = next(iter(view_entries.values()))
    return int(first[0]), int(first[1])


def sanitize_nan_inf_streams(stream_dicts, masks=None, iteration=None, sample_info=None):
    '''
    :param stream_dicts (list): List of dicts with {'v0': tensor, 'v1': tensor, ...}.
    Replace NaN & infinity in stream tensors in place with zeros,
    and optionally zero the entire supervise mask for affected batch elements.
    NaN can come from corrupted data in the loading/encoding pipeline.
    All ranks must still run the same forward/backward path (FSDP).
    '''
    for streams_dict in stream_dicts:
        for k, v in streams_dict.items():
            nan_inf_mask = v.isnan() | v.isinf()
            nan_inf_count = nan_inf_mask.sum()
            if not nan_inf_mask.any():
                continue
            
            nan_batch = nan_inf_mask.flatten(1).any(dim=1)  # (B,) per-element flag
            affected_batch_inds = nan_batch.nonzero().flatten().tolist()
            
            my_rank = torch.distributed.get_rank()
            log.error(
                f'[RANK{my_rank}] NaN/Inf in stream [{k}]! '
                f'nan_inf_count={nan_inf_count}/{v.numel()}, '
                f'affected batch indices: {affected_batch_inds}, '
                f'iteration: {iteration}, '
                f'sample_info: {sample_info}, '
                f'replacing with zeros',
                rank0_only=False)
            
            streams_dict[k] = v.nan_to_num(0.0)  # also captures inf
            
            if masks is not None:
                for mk, mv in masks['supervise'].items():
                    if mv[nan_batch].any():
                        masks['supervise'][mk] = mv.clone()
                        masks['supervise'][mk][nan_batch] = 0.0


def dict_to_cpu_recursive(d):
    if isinstance(d, torch.Tensor):
        return d.to(device='cpu', non_blocking=True)
    elif isinstance(d, dict):
        for k, v in d.items():
            d[k] = dict_to_cpu_recursive(v)
        return d
    elif isinstance(d, list):
        return [dict_to_cpu_recursive(x) for x in d]
    else:
        return d


def dict_to_cuda_recursive(d):
    if isinstance(d, torch.Tensor):
        return d.to(device='cuda', non_blocking=True)
    elif isinstance(d, dict):
        for k, v in d.items():
            d[k] = dict_to_cuda_recursive(v)
        return d
    elif isinstance(d, list):
        return [dict_to_cuda_recursive(x) for x in d]
    else:
        return d


def dict_unsqueeze_batch_recursive(d):
    if isinstance(d, torch.Tensor) or isinstance(d, np.ndarray):
        return d.unsqueeze(0)
    elif isinstance(d, dict):
        for k, v in d.items():
            d[k] = dict_unsqueeze_batch_recursive(v)
        return d
    elif isinstance(d, list):
        return [dict_unsqueeze_batch_recursive(x) for x in d]
    else:
        return [d]


def dict_squeeze_batch_recursive(d):
    if isinstance(d, torch.Tensor):
        return d.squeeze(0)
    elif isinstance(d, dict):
        for k, v in d.items():
            d[k] = dict_squeeze_batch_recursive(v)
        return d
    elif isinstance(d, list):
        if len(d) != 1:
            return [dict_squeeze_batch_recursive(x) for x in d]
        else:
            return d[0]
    else:
        return d


def dict_prune_memory_recursive(d, threshold=1024):
    if isinstance(d, torch.Tensor):
        if d.numel() > threshold:
            return 'pruned'
        else:
            return d
    elif isinstance(d, dict):
        for k, v in list(d.items()):
            d[k] = dict_prune_memory_recursive(v, threshold=threshold)
        return d
    elif isinstance(d, list):
        return [dict_prune_memory_recursive(x, threshold=threshold) for x in d]
    else:
        return d


@torch.no_grad()
def detach_dict_recursive(d):
    if isinstance(d, torch.Tensor):
        return d.detach()
    elif isinstance(d, dict):
        for k, v in d.items():
            d[k] = detach_dict_recursive(v)
        return d
    elif isinstance(d, list):
        return [detach_dict_recursive(x) for x in d]
    else:
        return d


@torch.no_grad()
def impute_data_batch_masks(
    data_batch: dict[str, torch.Tensor],
    config,
) -> dict[str, torch.Tensor]:
    '''
    Make the training pipeline robust to partial data availability by adding missing mask keys.
    Also store video dimensions per viewpoint as additional metadata.
    Default input mask is 1;
    Default output mask is 1 - input mask;
    Default supervise mask is output mask.
    NOTE(bvh): It is the responsibility of the caller to process raw and/or latent data as needed.
    :param a4d_raw (dict): Input data batch containing entry values.
    :param a4d_latent (dict): Input data batch containing entry masks.
    :return (a4d_raw, a4d_latent): Updated entries with imputed masks.
    '''
    if config.legacy_logistics_behavior == 2:
        from custom.legacy.a4d_logistics_vpad import impute_data_batch_masks as _legacy
        return _legacy(data_batch, config)
    
    V_max = config.num_views

    # Create shallow copies to avoid modifying the original
    batch_imp = dict(data_batch)
    a4d_raw = dict(data_batch['a4d_raw'])
    a4d_latent = dict(data_batch['a4d_latent'])
    
    video_dims_raw = [None] * V_max
    video_dims_latent = [None] * V_max
    
    # Ensure no pixel space masks exist
    for k in sorted(a4d_raw.keys()):
        assert not(k.endswith('_mask')), \
            'Masks should not be stored in the raw dict anymore for efficiency reasons.'

    # Process input & output channel mappings for all video viewpoints
    for v in range(V_max):
        for (k, c) in config.video_entries[v].items():
            if k.endswith('_mask'):
                continue
            
            if k in a4d_latent:
                device = a4d_latent[k].device
                (B, Cl, Tl, Hl, Wl) = a4d_latent[k].shape
                Tp = (Tl - 1) * VAE_RATIO_THW[0] + 1
                Hp = Hl * VAE_RATIO_THW[1]
                Wp = Wl * VAE_RATIO_THW[2]
            
            elif k in a4d_raw:
                device = a4d_raw[k].device
                (B, Cp, Tp, Hp, Wp) = a4d_raw[k].shape
                Tl = (Tp - 1) // VAE_RATIO_THW[0] + 1
                Hl = Hp // VAE_RATIO_THW[1]
                Wl = Wp // VAE_RATIO_THW[2]
            
            else:
                # This entry is skipped (does not exist in the current data batch)
                # If this true for all entries in this viewpoint, then video_dims remains None,
                # which means the entire stream will also be skipped.
                continue

            tensor_kwargs = dict(device=device, dtype=torch.bfloat16)

            # Keep track of resolutions per viewpoint, and verify consistency
            if video_dims_raw[v] is None:
                video_dims_raw[v] = (Tp, Hp, Wp)
            else:
                assert (Tp, Hp, Wp) == video_dims_raw[v], \
                    (f'All entries (modalities) within a stream (viewpoint) must have the same '
                     f'dimensions (Tp, Hp, Wp), but {(Tp, Hp, Wp)} != {video_dims_raw[v]}.')
            
            if video_dims_latent[v] is None:
                video_dims_latent[v] = (Tl, Hl, Wl)
            else:
                assert (Tl, Hl, Wl) == video_dims_latent[v], \
                    (f'All entries (modalities) within a stream (viewpoint) must have the same '
                     f'dimensions (Tl, Hl, Wl), but {(Tl, Hl, Wl)} != {video_dims_latent[v]}.')
            
            # Add input mask if missing
            if k + '_input_mask' not in a4d_latent:
                # Channel dimension must be reduced to 1
                a4d_latent[k + '_input_mask'] = torch.ones(B, 1, Tl, Hl, Wl, **tensor_kwargs)
            
            # Add output mask if missing
            if k + '_output_mask' not in a4d_latent:
                a4d_latent[k + '_output_mask'] = 1.0 - a4d_latent[k + '_input_mask']
            
            # Add supervise mask if missing
            if k + '_supervise_mask' not in a4d_latent:
                a4d_latent[k + '_supervise_mask'] = a4d_latent[k + '_output_mask']

    # Save info for later usage
    a4d_raw['video_dims'] = video_dims_raw
    a4d_latent['video_dims'] = video_dims_latent
    # used_views_raw = len([v for v in range(V_max) if video_dims_raw[v] is not None])
    # used_views_latent = len([v for v in range(V_max) if video_dims_latent[v] is not None])
    # assert used_views_raw == used_views_latent, f'{used_views_raw} != {used_views_latent}'
    # a4d_raw['used_views'] = used_views_raw
    # a4d_latent['used_views'] = used_views_latent

    # Process self-attention input & output tokens
    for (k, c) in config.lowdim_sattn_entries.items():
        if k not in a4d_raw or k.endswith('_mask'):
            continue

        # Add input mask if missing
        if k + '_input_mask' not in a4d_latent:
            a4d_latent[k + '_input_mask'] = torch.ones_like(a4d_raw[k][:, :, 0:1])
        
        # Add output mask if missing
        if k + '_output_mask' not in a4d_latent:
            a4d_latent[k + '_output_mask'] = 1.0 - a4d_latent[k + '_input_mask']

        # Add supervise mask if missing
        if k + '_supervise_mask' not in a4d_latent:
            a4d_latent[k + '_supervise_mask'] = a4d_latent[k + '_output_mask']

    # Process cross-attention input tokens
    for (k, c) in config.lowdim_xattn_entries.items():
        if k not in a4d_raw or k.endswith('_mask'):
            continue

        # Add input mask if missing
        if k + '_input_mask' not in a4d_latent:
            a4d_latent[k + '_input_mask'] = torch.ones_like(a4d_raw[k][:, :, 0:1])
    
    # Process adaptive layernorm input tokens
    for (k, c) in config.lowdim_adaln_entries.items():
        if k not in a4d_raw or k.endswith('_mask'):
            continue

        # Add input mask if missing
        if k + '_input_mask' not in a4d_latent:
            a4d_latent[k + '_input_mask'] = torch.ones_like(a4d_raw[k][:, 0:1])

    batch_imp['a4d_raw'] = a4d_raw
    batch_imp['a4d_latent'] = a4d_latent
    
    return batch_imp


@torch.no_grad()
def pack_streams_from_entries(
    entries: dict[str, torch.Tensor],
    config,
) -> tuple[dict[str, torch.Tensor], dict[str, dict[str, torch.Tensor]]]:
    '''
    Assemble input and output sequences using the data batch and masks.
    This function organizes Any4D latent data into streams (v0, v1, ..., sattn, xattn, adaln)
    and creates the necessary tensors and mask dictionaries.
    :param entries (dict): Typically data_batch['a4d_latent'].
    :return streams (dict), masks (dict).
    '''
    if config.legacy_logistics_behavior == 2:
        from custom.legacy.a4d_logistics_vpad import pack_streams_from_entries as _legacy
        return _legacy(entries, config)
    
    V_max = config.num_views
    # V_used = entries['used_views']
    # V_used = V_max

    B = entries['rgb0'].shape[0]
    device = entries['rgb0'].device
    dtype = entries['rgb0'].dtype
    tensor_kwargs = dict(device=device, dtype=dtype)

    streams = dict()
    masks = dict()
    masks['is_mask'] = dict()
    masks['input'] = dict()
    masks['output'] = dict()
    masks['supervise'] = dict()
    
    # bookmark(bvh): FSDP sharding
    # NOTE(bvh): V_used may differ across ranks. To avoid NCCL problems / timeouts with FSDP,
    # always create streams for ALL V_max views, creating zero-padded stubs for missing ones
    # (which contribute no loss, and as of 05/19/2026, do not change the forward pass in any way).
    if config.legacy_logistics_behavior == 0:
        ref_dims = (1, 2, 2)  # (T, H, W)
        # NOTE(bvh): ^ stub streams are skipped in a4d_network as decided by numel() <= 4
    elif config.legacy_logistics_behavior == 3:
        ref_dims = entries['video_dims'][0]  # # (T, H, W); view 0 always exists

    for v in range(V_max):
        Cv = config.num_video_in_channels[v]
        cur_dims = entries['video_dims'][v]
        if cur_dims is None:
            if config.legacy_logistics_behavior == 1:
                continue
            else:
                cur_dims = ref_dims

        (Tv, Hv, Wv) = cur_dims
        streams[f'v{v}'] = torch.zeros((B, Cv, Tv, Hv, Wv), **tensor_kwargs)
        masks['is_mask'][f'v{v}'] = torch.zeros((B, Cv, Tv, Hv, Wv), **tensor_kwargs)
        masks['input'][f'v{v}'] = torch.zeros((B, Cv, Tv, Hv, Wv), **tensor_kwargs)
        masks['output'][f'v{v}'] = torch.zeros((B, Cv, Tv, Hv, Wv), **tensor_kwargs)
        masks['supervise'][f'v{v}'] = torch.zeros((B, Cv, Tv, Hv, Wv), **tensor_kwargs)

    if config.has_sattn_stream:
        (Ts, Cs) = (config.num_lowdim_sattn_tokens, config.num_lowdim_sattn_channels)
        streams['sattn'] = torch.zeros((B, Ts, Cs), **tensor_kwargs)
        masks['is_mask']['sattn'] = torch.zeros((B, Ts, Cs), **tensor_kwargs)
        masks['input']['sattn'] = torch.zeros((B, Ts, Cs), **tensor_kwargs)
        masks['output']['sattn'] = torch.zeros((B, Ts, Cs), **tensor_kwargs)
        masks['supervise']['sattn'] = torch.zeros((B, Ts, Cs), **tensor_kwargs)
    
    if config.has_xattn_stream:
        (Tx, Cx) = (config.num_lowdim_xattn_tokens, config.num_lowdim_xattn_channels)
        streams['xattn'] = torch.zeros((B, Tx, Cx), **tensor_kwargs)
        masks['is_mask']['xattn'] = torch.zeros((B, Tx, Cx), **tensor_kwargs)
        masks['input']['xattn'] = torch.zeros((B, Tx, Cx), **tensor_kwargs)
    
    if config.has_adaln_stream:
        # (Ta, Ca) = (Tv0, config.num_lowdim_adaln_channels)
        (Ta, Ca) = (config.num_lowdim_adaln_tokens, config.num_lowdim_adaln_channels)
        streams['adaln'] = torch.zeros((B, Ta, Ca), **tensor_kwargs)
        masks['is_mask']['adaln'] = torch.zeros((B, Ta, Ca), **tensor_kwargs)
        masks['input']['adaln'] = torch.zeros((B, Ta, Ca), **tensor_kwargs)

    # ==========================================================
    # ======== Handle high-dimensional inputs & outputs ========
    # ==========================================================

    for v in range(V_max):
        for (k, c) in config.video_entries[v].items():
            if k.endswith('_mask') and f'v{v}' in streams:
                masks['is_mask'][f'v{v}'][:, c[0]:c[1]] = 1.0
                masks['input'][f'v{v}'][:, c[0]:c[1]] = 1.0
                masks['output'][f'v{v}'][:, c[0]:c[1]] = 0.0
                masks['supervise'][f'v{v}'][:, c[0]:c[1]] = 0.0
            
            if k not in entries:
                if config.strict_populate_data:
                    raise ValueError(f'Missing video entry: {k}')
                else:
                    continue
            
            streams[f'v{v}'][:, c[0]:c[1]] = entries[k]
            if not(k.endswith('_mask')):
                masks['is_mask'][f'v{v}'][:, c[0]:c[1]] = 0.0
                masks['input'][f'v{v}'][:, c[0]:c[1]] = entries[k + '_input_mask']
                masks['output'][f'v{v}'][:, c[0]:c[1]] = entries[k + '_output_mask']
                masks['supervise'][f'v{v}'][:, c[0]:c[1]] = entries[k + '_supervise_mask']
    
    # =========================================================
    # ======== Handle low-dimensional inputs & outputs ========
    # =========================================================

    for (k, c) in config.lowdim_sattn_entries.items():
        if k.endswith('_mask'):
            masks['is_mask']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 1.0
            masks['input']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 1.0
            masks['output']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0
            masks['supervise']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0
        
        if k not in entries:
            if config.strict_populate_data:
                raise ValueError(f'Missing sattn entry: {k}')
            else:
                continue
            
        streams['sattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k]
        if not(k.endswith('_mask')):
            masks['is_mask']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0
            masks['input']['sattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k + '_input_mask']
            masks['output']['sattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k + '_output_mask']
            masks['supervise']['sattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k + '_supervise_mask']

    for (k, c) in config.lowdim_xattn_entries.items():
        if k.endswith('_mask'):
            masks['is_mask']['xattn'][:, c[0]:c[1], c[2]:c[3]] = 1.0
            masks['input']['xattn'][:, c[0]:c[1], c[2]:c[3]] = 1.0

        if k not in entries:
            if config.strict_populate_data:
                raise ValueError(f'Missing xattn entry: {k}')
            else:
                continue
            
        streams['xattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k]
        if not(k.endswith('_mask')):
            masks['is_mask']['xattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0
            masks['input']['xattn'][:, c[0]:c[1], c[2]:c[3]] = entries[k + '_input_mask']

    for (k, c) in config.lowdim_adaln_entries.items():
        if k.endswith('_mask'):
            masks['is_mask']['adaln'][:, c[0]:c[1], c[2]:c[3]] = 1.0
            masks['input']['adaln'][:, c[0]:c[1], c[2]:c[3]] = 1.0

        if k not in entries:
            if config.strict_populate_data:
                raise ValueError(f'Missing adaln entry: {k}')
            else:
                continue

        streams['adaln'][:, c[0]:c[1], c[2]:c[3]] = entries[k]
        if not(k.endswith('_mask')):
            masks['is_mask']['adaln'][:, c[0]:c[1], c[2]:c[3]] = 0.0
            masks['input']['adaln'][:, c[0]:c[1], c[2]:c[3]] = entries[k + '_input_mask']

    return (streams, masks)


def unpack_entries_from_streams(
    streams: dict[str, torch.Tensor],
    config,
    detach: bool = False,
    ignore_masks: bool = False,
) -> dict[str, torch.Tensor]:
    '''
    NOTE: Masks are completely ignored for this purpose.
    :param streams (dict).
    :return entries (dict).
    '''
    if config.legacy_logistics_behavior == 2:
        from custom.legacy.a4d_logistics_vpad import unpack_entries_from_streams as _legacy
        return _legacy(streams, config, detach=detach, ignore_masks=ignore_masks)
    
    V_max = config.num_views
    # V_used = streams['used_views']
    # V_used = V_max
    
    # B = streams['v0'].shape[0]
    # device = streams['v0'].device
    # dtype = streams['v0'].dtype
    # tensor_kwargs = dict(device=device, dtype=dtype)

    entries = dict()

    for v in range(V_max):
        if f'v{v}' in streams:
            assert streams[f'v{v}'].ndim == 5, \
                f'video stream v{v} must be (B, Cl, Tl, Hl, Wl), but got {streams[f"v{v}"].shape} instead.'
            for (k, c) in config.video_entries[v].items():
                if (not(ignore_masks) or not(k.endswith('_mask'))):
                    entries[k] = streams[f'v{v}'][:, c[0]:c[1]]
                    if detach:
                        entries[k] = entries[k].detach()
        
    if 'sattn' in streams:
        assert streams['sattn'].ndim == 3, \
            f'sattn stream must be (B, T, C), but got {streams["sattn"].shape} instead.'
        for (k, c) in config.lowdim_sattn_entries.items():
            if (not(ignore_masks) or not(k.endswith('_mask'))):
                entries[k] = streams['sattn'][:, c[0]:c[1], c[2]:c[3]]
                if detach:
                    entries[k] = entries[k].detach()
    
    if 'xattn' in streams:
        assert streams['xattn'].ndim == 3, \
            f'xattn stream must be (B, T, C), but got {streams["xattn"].shape} instead.'
        for (k, c) in config.lowdim_xattn_entries.items():
            if (not(ignore_masks) or not(k.endswith('_mask'))):
                entries[k] = streams['xattn'][:, c[0]:c[1], c[2]:c[3]]
                if detach:
                    entries[k] = entries[k].detach()

    if 'adaln' in streams:
        assert streams['adaln'].ndim == 3, \
            f'adaln stream must be (B, T, C), but got {streams["adaln"].shape} instead.'
        for (k, c) in config.lowdim_adaln_entries.items():
            if (not(ignore_masks) or not(k.endswith('_mask'))):
                entries[k] = streams['adaln'][:, c[0]:c[1], c[2]:c[3]]
                if detach:
                    entries[k] = entries[k].detach()

    return entries


@torch.no_grad()
def verify_entry_shapes(
    a4d_latent: dict[str, torch.Tensor],
    config,
) -> None:
    '''
    Verify all encoded entry shapes (consistency with config)
    :return (a4d_raw, a4d_latent): Updated entries with imputed masks.
    '''
    highdim_entries = {k: c for v in config.video_entries for k, c in v.items()}  # flattened
    lowdim_entries = {**config.lowdim_sattn_entries,
                      **config.lowdim_xattn_entries,
                      **config.lowdim_adaln_entries}  # combined
    # Tv0 = a4d_latent['video_dims'][0][0]
    
    for (k, e) in a4d_latent.items():
        if k in highdim_entries:
            c = highdim_entries[k]
            assert e.ndim == 5, \
                f'video entry {k} must be (B, Cl, Tl, Hl, Wl), but got {e.shape} instead.'
            assert e.shape[1] == c[1] - c[0], \
                f'video entry {k} must have {c[1] - c[0]} channels, but got {e.shape[1]} instead.'
        
        elif k in lowdim_entries:
            c = lowdim_entries[k]
            if k in config.lowdim_sattn_entries:
                assert e.ndim == 3, \
                    f'sattn entry {k} must be (B, T, C), but got {e.shape} instead.'
                assert e.shape[1] == c[1] - c[0] and e.shape[2] == c[3] - c[2], \
                    f'sattn entry {k} must have {c[1] - c[0]} tokens and {c[3] - c[2]} channels, ' \
                    f'but got {e.shape[1]} tokens and {e.shape[2]} channels instead.'
            
            elif k in config.lowdim_xattn_entries:
                assert e.ndim == 3, \
                    f'xattn entry {k} must be (B, T, C), but got {e.shape} instead.'
                assert e.shape[1] == c[1] - c[0] and e.shape[2] == c[3] - c[2], \
                    f'xattn entry {k} must have {c[1] - c[0]} tokens and {c[3] - c[2]} channels, ' \
                    f'but got {e.shape[1]} tokens and {e.shape[2]} channels instead.'
            
            elif k in config.lowdim_adaln_entries:
                assert e.ndim == 3, \
                    f'adaln entry {k} must be (B, T, C), but got {e.shape} instead.'
                # assert e.shape[1] == Tv0 and e.shape[2] == c[1] - c[0], \
                #     f'adaln entry {k} must have {Tv0} timesteps and {c[1] - c[0]} channels, ' \
                #     f'but got {e.shape[1]} timesteps and {e.shape[2]} channels instead.'
                assert e.shape[1] == c[1] - c[0] and e.shape[2] == c[3] - c[2], \
                    f'adaln entry {k} must have {c[1] - c[0]} tokens and {c[3] - c[2]} channels, ' \
                    f'but got {e.shape[1]} tokens and {e.shape[2]} channels instead.'


@torch.no_grad()
def validate_masks(
    masks: dict[str, dict[str, torch.Tensor]],
) -> bool:
    '''
    Perform sanity checks on masks to ensure they meet expected constraints:
    1. For x_masks, input_mask and output_mask should be mutually exclusive.
    2. For y_masks, supervise_mask should be a subset of output_mask.
    3. Input and output masks should be binary.
    4. Supervise masks should be non-negative.
    :param masks (dict).
    :return success (bool).
    '''
    for k in set(masks['input'].keys()) & set(masks['output'].keys()):
        assert torch.all(masks['input'][k] * masks['output'][k] == 0.0), \
            f'Input and output masks may not overlap / must be mutually exclusive, stream {k}'
    
    for k in set(masks['output'].keys()) & set(masks['supervise'].keys()):
        assert torch.all(masks['output'][k] * masks['supervise'][k] == masks['supervise'][k]), \
            f'Supervise mask must be subset of output mask, stream {k}'

    for k in set(masks['input'].keys()):
        assert torch.all((masks['input'][k] == 0.0) | (masks['input'][k] == 1.0)), \
            f'Input mask must be binary, but got {masks["input"][k]} instead.'

    for k in set(masks['output'].keys()):
        assert torch.all((masks['output'][k] == 0.0) | (masks['output'][k] == 1.0)), \
            f'Output mask must be binary, but got {masks["output"][k]} instead.'

    for k in set(masks['supervise'].keys()):
        assert torch.all(masks['supervise'][k] >= 0.0), \
            f'Supervise mask must be non-negative (since they act as fine-grained loss weights), ' \
            f'but got {masks["supervise"][k]} instead.'
    
    return True


@torch.no_grad()
def assemble_network_inputs(
    x0_streams: dict[str, torch.Tensor],
    yt_streams: dict[str, torch.Tensor],
    masks: dict[str, dict[str, torch.Tensor]],
) -> dict[str, torch.Tensor]:
    '''
    :return xt_streams (dict).
    Assemble the network input tensors (xt_streams) by combining the conditioning inputs
    (x0_streams) and noisy target values (yt_streams), typically already scaled as needed.
    :param x0_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float.
    :param yt_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float.
    :param masks (dict).
    :return xt_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float.
    '''
    x0_masked = multiply_dicts(x0_streams, masks['input'], key_mode='equal')
    yt_masked = multiply_dicts(yt_streams, masks['output'], key_mode='equal')
    
    xt_streams = add_dicts(x0_masked, yt_masked, key_mode='union')
    
    return xt_streams


def assemble_clean_predictions(
    x0_streams: dict[str, torch.Tensor],
    y0_raw_pred_streams: dict[str, torch.Tensor],
    masks: dict[str, dict[str, torch.Tensor]],
    config=None,
) -> dict[str, torch.Tensor]:
    '''
    Assemble the virtual / final prediction tensors (y0_pred_streams) that does not reveal
    unsupervised network outputs by combining the non-augmented conditioning inputs (x0_streams)
    with the predicted outputs (y0_raw_pred_streams). Used for pretty much everything
    (loss, visualizations, metrics).
    :return y0_pred_streams (dict).
    '''
    if config is not None and config.legacy_logistics_behavior == 2:
        from custom.legacy.a4d_logistics_vpad import assemble_clean_predictions as _legacy
        return _legacy(x0_streams, y0_raw_pred_streams, masks)
    
    # Sanitize predictions: replace NaN/Inf with zeros before masking, as 0 * NaN = NaN in PyTorch.
    # (corrupted data safety/fallback to prevent NaNs in gradients).
    y0_sanitized = y0_raw_pred_streams
    sanitize_nan_inf_streams([y0_sanitized], masks=masks)

    x0_masked = multiply_dicts(x0_streams, masks['input'], key_mode='equal')
    y0_masked = multiply_dicts(y0_sanitized, masks['output'], key_mode='equal')

    y0_pred_streams = add_dicts(x0_masked, y0_masked, key_mode='intersect')
    
    return y0_pred_streams


def add_dicts(
    a: dict[str, torch.Tensor],
    b: dict[str, torch.Tensor],
    key_mode: str = 'equal',
) -> dict[str, torch.Tensor]:
    '''
    :return c (dict): a + b.
    '''
    return math_dicts(a, b, operation='add', key_mode=key_mode)


def multiply_dicts(
    a: dict[str, torch.Tensor],
    b: dict[str, torch.Tensor],
    key_mode: str = 'equal',
) -> dict[str, torch.Tensor]:
    '''
    :return c (dict): a * b.
    '''
    return math_dicts(a, b, operation='multiply', key_mode=key_mode)


def math_dicts(
    a: dict[str, torch.Tensor],
    b: dict[str, torch.Tensor],
    operation: str = 'add',
    key_mode: str = 'equal',
) -> dict[str, torch.Tensor]:
    '''
    :return c (dict): a operand b.
    '''
    
    if key_mode == 'intersect':
        keys = set(a.keys()) & set(b.keys())
    elif key_mode == 'a':
        keys = set(a.keys())
    elif key_mode == 'b':
        keys = set(b.keys())
    elif key_mode == 'union':
        keys = set(a.keys()) | set(b.keys())
    elif key_mode == 'equal':
        keys = set(a.keys())
        assert keys == set(b.keys()), \
            f'Keys must be equal, but {keys} != {set(b.keys())}'
    else:
        raise ValueError(f'Invalid key_mode: {key_mode}')
    
    result = dict()
    for k in sorted(keys):
        if k in a and k in b:
            if operation == 'add':
                result[k] = a[k] + b[k]
            elif operation == 'subtract':
                result[k] = a[k] - b[k]
            elif operation == 'multiply':
                result[k] = a[k] * b[k]
            elif operation == 'divide':
                result[k] = a[k] / b[k]
            else:
                raise ValueError(f'Invalid operation: {operation}')
        
        elif k in a:
            result[k] = a[k]
        
        elif k in b:
            result[k] = b[k]
    
    return result