"""Post-processing functions for BIO label predictions.

Contains label correction (label_fixes), segment extraction, and BEGIN tag merging.

Adapted from segmenterv2/segmenter/inference/postprocess.py.
"""

from __future__ import annotations

import torch


def label_fixes(
    labels: torch.Tensor,
    min_segment_len: int = 2,
    min_start_len: int = 6,
) -> torch.Tensor:
    """Two-pass correction of BIO predictions.

    Pass 1 (reverse): Insert BEGIN at the start of IN-runs >= min_segment_len.
    Pass 2 (forward): Collapse consecutive BEGINs (keep first, rest become IN).

    Args:
        labels: 1D tensor of BIO predictions (0=OUT, 1=IN, 2=BEGIN).
        min_segment_len: Minimum IN-run length to trigger BEGIN insertion.
        min_start_len: Minimum IN-run length at sequence start to insert BEGIN.

    Returns:
        Corrected labels tensor (new tensor, input is not mutated).
    """
    labels = labels.clone()
    length = len(labels)
    if length == 0:
        return labels

    # Pass 1: Reverse pass to insert BEGIN at start of IN-runs
    one_counter = 0
    for index in reversed(range(length)):
        if labels[index] == 1:
            one_counter += 1
        else:
            if one_counter >= min_segment_len:
                if index + 1 < length and labels[index + 1] != 2:
                    labels[index + 1] = 2
            one_counter = 0

    # Edge case: IN-run starting from index 0
    if one_counter >= min_start_len and labels[0] != 2:
        labels[0] = 2

    # Pass 2: Forward pass to collapse consecutive BEGINs
    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


def extract_segments(prediction: torch.Tensor) -> list[tuple[int, int]]:
    """Extract (start_frame, end_frame) tuples from BIO predictions.

    A segment starts at BEGIN or at the first IN after OUT.
    A segment ends when OUT is encountered or a new BEGIN starts.

    Args:
        prediction: 1D tensor of BIO predictions (0=OUT, 1=IN, 2=BEGIN).

    Returns:
        List of (start_frame, end_frame) tuples (inclusive).
    """
    pred = prediction.clone()
    for i in reversed(range(1, len(pred))):
        if pred[i] == 2 and pred[i - 1] == 2:
            pred[i] = 1

    segments = []
    start_time = None
    for i, p in enumerate(pred):
        if p == 1 or p == 2:
            if start_time is not None and p == 2:
                segments.append((start_time, i - 1))
                start_time = None
            if start_time is None:
                start_time = i
        elif p == 0:
            if start_time is not None:
                segments.append((start_time, i - 1))
                start_time = None

    if start_time is not None:
        segments.append((start_time, len(pred) - 1))

    return segments


def merge_begin_tags(labels: torch.Tensor) -> torch.Tensor:
    """Merge BEGIN tags into IN tags. BEGIN(2) -> IN(1).

    Args:
        labels: 1D tensor of BIO predictions.

    Returns:
        New tensor with all BEGIN (2) values replaced by IN (1).
    """
    result = labels.clone()
    result[result == 2] = 1
    return result
