from typing import Tuple, List

from pytorch_ie.annotations import BinaryRelation, Span, LabeledMultiSpan
from pytorch_ie.core import Annotation


def decode_span_without_label(ann: Annotation) -> Tuple[Tuple[int, int], ...]:
    if isinstance(ann, Span):
        return (ann.start, ann.end),
    elif isinstance(ann, LabeledMultiSpan):
        return ann.slices
    else:
        raise ValueError("Annotation must be a Span or LabeledMultiSpan")


def to_binary_relation_without_argument_labels(ann: Annotation) -> Tuple[Tuple[Tuple[int, int], ...], Tuple[Tuple[int, int], ...], str]:
    if not isinstance(ann, BinaryRelation):
        raise ValueError("Annotation must be a BinaryRelation")
    return (
        decode_span_without_label(ann.head),
        decode_span_without_label(ann.tail),
        ann.label,
    )