from typing import Dict
from typing import List
from typing import Tuple


def allowed_transitions(constraint_type: str,
                        tokens: Dict[int, str]) -> List[Tuple[int, int]]:
    # pylint: disable=too-many-branches
    """
    Given tokens and a constraint type, returns the allowed transitions.

    It will additionally include transitions for the start and end states,
    which are used by the conditional random field.

    Parameters
    ----------
    constraint_type : str, required
        Indicates which constraint to apply. Current choices are "BIO" and "BIOUL".

    tokens : Dict[int, str], required
        A mapping {token_id -> token}. Most commonly this would be the value from
        Vocabulary.get_index_to_token_vocabulary()

    Returns
    -------
    List[Tuple[int, int]]
        The allowed transitions (from_token_id, to_token_id).

    """
    n_tags = len(tokens)
    start_tag = n_tags
    end_tag = n_tags + 1

    allowed = []
    if constraint_type == "BIOUL":
        for i, (from_bioul, *from_entity) in tokens.items():
            for j, (to_bioul, *to_entity) in tokens.items():
                is_allowed = any([
                    # O can transition to O, B-* or U-*
                    # L-x can transition to O, B-*, or U-*
                    # U-x can transition to O, B-*, or U-*
                    from_bioul in ('O', 'L', 'U') and to_bioul in ('O', 'B', 'U'),
                    # B-x can only transition to I-x or L-x
                    # I-x can only transition to I-x or L-x
                    from_bioul in ('B', 'I') and to_bioul in ('I', 'L') and from_entity == to_entity
                ])

                if is_allowed:
                    allowed.append((i, j))

        # start transitions
        for i, (to_bioul, *to_entity) in tokens.items():
            if to_bioul in ('O', 'B', 'U'):
                allowed.append((start_tag, i))

        # end transitions
        for i, (from_bioul, *from_entity) in tokens.items():
            if from_bioul in ('O', 'L', 'U'):
                allowed.append((i, end_tag))

    elif constraint_type == "BIO":
        for i, (from_bio, *from_entity) in tokens.items():
            for j, (to_bio, *to_entity) in tokens.items():

                is_allowed = any([
                    # Can always transition to O or B-x
                    to_bio in ('O', 'B'),
                    # Can only transition to I-x from B-x or I-x
                    to_bio == 'I' and from_bio in ('B', 'I') and from_entity == to_entity
                ])

                if is_allowed:
                    allowed.append((i, j))

        # start transitions
        for i, (to_bio, *to_entity) in tokens.items():
            if to_bio in ('O', 'B'):
                allowed.append((start_tag, i))

        # end transitions
        for i, (from_bio, *from_entity) in tokens.items():
            if from_bio in ('O', 'B', 'I'):
                allowed.append((i, end_tag))

    else:
        raise ValueError(f"Unknown constraint type: {constraint_type}")

    return allowed