"""Post-processing for sign segmentation labels.

Ported from /dev/4TB_SSD/segmentor_v2/train_Wilor.py::label_fixes()
"""

from __future__ import annotations

import torch


def label_fixes(labels: torch.Tensor) -> torch.Tensor:
    """Apply post-processing to fix segmentation labels.

    Two-pass algorithm:
    1. Reverse pass: Insert BEGIN tokens (2) at start of long IN_SIGN sequences (1s)
    2. Forward pass: Collapse consecutive BEGIN tokens into single BEGIN + IN_SIGN

    Class mapping:
        0 = OUT_SIGN (no sign)
        1 = IN_SIGN (signing)
        2 = BEGIN (start of sign)

    Args:
        labels: Predicted labels of shape (T,)

    Returns:
        Fixed labels of shape (T,)
    """
    labels = labels.clone()  # Avoid in-place mutation
    length = len(labels)

    # Pass 1: Reverse pass to insert BEGIN (2) at start of long IN_SIGN (1) sequences
    one_counter = 0
    for index in reversed(range(length)):
        if labels[index] == 1:
            one_counter += 1
        else:
            if one_counter >= 2:
                # Check if the next label isn't already a BEGIN
                if index + 1 < length and labels[index + 1] != 2:
                    labels[index + 1] = 2
            one_counter = 0

    # Edge case: if the sequence of 1s starts from index 0
    if one_counter >= 6 and labels[0] != 2:
        labels[0] = 2

    # Pass 2: Forward pass to collapse consecutive BEGINs
    # Keep only the first BEGIN, convert the rest to IN_SIGN
    i = 0
    while i < length:
        if labels[i] == 2:
            j = i + 1
            while j < length and labels[j] == 2:
                labels[j] = 1
                j += 1
            i = j
        else:
            i += 1

    return labels
