import torch


def preprocess_video_panoptic_gt(
        gt_labels,
        gt_masks,
        gt_semantic_seg,
        gt_instance_ids,
        num_things,
        num_stuff,
):
    num_classes = num_things + num_stuff
    num_frames = len(gt_masks)
    mask_size = gt_masks[0].masks.shape[-2:]

    thing_masks_list = []
    for frame_id in range(num_frames):
        thing_masks_list.append(gt_masks[frame_id].pad(
            mask_size, pad_val=0).to_tensor(
            dtype=torch.bool, device=gt_labels.device)
        )
    instances = torch.unique(gt_instance_ids[:, 1])
    things_masks = []
    labels = []
    for instance in instances:
        pos_ins = torch.nonzero(torch.eq(gt_instance_ids[:, 1], instance), as_tuple=True)[0]  # 0 is for redundant tuple
        labels_instance = gt_labels[:, 1][pos_ins]
        assert torch.allclose(labels_instance, labels_instance[0])
        labels.append(labels_instance[0])
        instance_frame_ids = gt_instance_ids[:, 0][pos_ins].to(dtype=torch.int32).tolist()
        instance_masks = []
        for frame_id in range(num_frames):
            frame_instance_ids = gt_instance_ids[gt_instance_ids[:, 0] == frame_id, 1]
            if frame_id not in instance_frame_ids:
                empty_mask = torch.zeros(
                    mask_size,
                    dtype=thing_masks_list[frame_id].dtype, device=thing_masks_list[frame_id].device
                )
                instance_masks.append(empty_mask)
            else:
                pos_inner_frame = torch.nonzero(torch.eq(frame_instance_ids, instance), as_tuple=True)[0].item()
                frame_mask = thing_masks_list[frame_id][pos_inner_frame]
                instance_masks.append(frame_mask)
        things_masks.append(torch.stack(instance_masks))

    if len(instances) == 0:
        things_masks = torch.stack(thing_masks_list, dim=1)
        labels = torch.empty_like(instances)
    else:
        things_masks = torch.stack(things_masks)
        labels = torch.stack(labels)
    assert torch.all(torch.less(labels, num_things))

    if gt_semantic_seg is not None:
        things_labels = labels
        gt_semantic_seg = gt_semantic_seg.squeeze(1)

        semantic_labels = torch.unique(
            gt_semantic_seg,
            sorted=False,
            return_inverse=False,
            return_counts=False)
        stuff_masks_list = []
        stuff_labels_list = []
        for label in semantic_labels:
            if label < num_things or label >= num_classes:
                continue
            stuff_mask = gt_semantic_seg == label
            stuff_masks_list.append(stuff_mask)
            stuff_labels_list.append(label)

        if len(stuff_masks_list) > 0:
            stuff_masks = torch.stack(stuff_masks_list, dim=0)
            stuff_labels = torch.stack(stuff_labels_list, dim=0)
            assert torch.all(torch.ge(stuff_labels, num_things)) and torch.all(torch.less(stuff_labels, num_classes))
            labels = torch.cat([things_labels, stuff_labels], dim=0)
            masks = torch.cat([things_masks, stuff_masks], dim=0)
        else:
            labels = things_labels
            masks = things_masks
        assert len(labels) == len(masks)
    else:
        masks = things_masks

    labels = labels.to(dtype=torch.long)
    masks = masks.to(dtype=torch.long)
    return labels, masks