import torch
from typing import Dict


def forward_pass(model: torch.nn.Module,
                 sample: Dict[str, torch.Tensor],
                 device: torch.device) -> Dict[str, torch.Tensor]:
    """ Forward pass of the network for a given sample
    """
    img = sample['img'].to(device)
    bboxes = sample['bboxes'].to(device)
    labels = sample['labels'].to(device)
    bboxes_mask = sample['bboxes_mask'].to(device)

    # Forward pass
    ploc, plabel, *aux_out = model(img)

    # Create the loss sample
    loss_sample = {}
    loss_sample['ploc'] = ploc
    loss_sample['plabel'] = plabel
    loss_sample['bboxes'] = bboxes
    loss_sample['labels'] = labels
    loss_sample['bboxes_mask'] = bboxes_mask

    # Check if we should include auxiliary
    if sample['auxiliary'].all() and aux_out is not None:
        loss_sample['scene_gt'] = sample['scene_id'].to(device)
        loss_sample['scene_pred'] = aux_out[0]

        loss_sample['depth_gt'] = sample['depth'].to(device)
        loss_sample['depth_pred'] = aux_out[1]

        loss_sample['normals_gt'] = sample['normals'].to(device)
        loss_sample['normals_mask'] = sample['normals_mask'].to(device)
        loss_sample['normals_pred'] = aux_out[2]

    return loss_sample