# 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 imaginaire.utils import log

# Ensure these constants are always correct / up2date:
VAE_RATIO = (4, 8, 8)  # (T, H, W) downsampling factors between pixel & latent space


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.
    '''
    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[0] + 1
                Hp = Hl * VAE_RATIO[1]
                Wp = Wl * VAE_RATIO[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[0] + 1
                Hl = Hp // VAE_RATIO[1]
                Wl = Wp // VAE_RATIO[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).
    '''
    V_max = config.num_views
    # V_used = entries['used_views']
    # V_used = V_max
    
    B = entries['rgb0'].shape[0]
    Tv0 = entries['video_dims'][0][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()
    
    for v in range(V_max):
        Cv = config.num_video_in_channels[v]
        cur_dims = entries['video_dims'][v]
        if cur_dims is None:
            # Entries in this stream are supported, but currently all skipped (no data)
            continue
        
        (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).
    '''
    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]],
) -> 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).
    '''
    x0_masked = multiply_dicts(x0_streams, masks['input'], key_mode='equal')
    y0_masked = multiply_dicts(y0_raw_pred_streams, 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