from typing import Union
from typing import Counter


def ents_to_relations(tokens: list[str], tags: list[str]) -> Union[list[str], None]:
    """
    Convert a list of tokens and tags into relations using the head and tail format
    used by the NYT corpus.

    Parameters
    ----------
    tokens : list[str]
        List of tokens.
    tags : list[str]
        List of tags.

    Returns
    -------
    Union[list[str], None]
        List of strings with head and tail annotations.

    Example
    -------

    >>> tokens = ['New York', 'is', 'in', 'New York', ',', 'America']
    >>> tags = ['PLACE', 'O', 'O', 'PLACE', 'O', 'PLACE']

    >>> ents_to_relations(tokens, tags)[0]
    '<head> New York </head> is in <tail> New York </tail> , America'
    >>> ents_to_relations(tokens, tags)[1]
    '<head> New York </head> is in New York , <tail> America </tail>'
    """

    # check if at least two entities
    if Counter(tags)["O"] > len(tags) - 2:
        return

    loc_idxs = [idx for idx, tag in enumerate(tags) if tag != "O"]
    sequence_list = []
    for i in loc_idxs:
        for j in loc_idxs:
            if i != j:
                tokens_copy = tokens.copy()
                tokens_copy[i] = "<head> " + tokens_copy[i] + " </head>"
                tokens_copy[j] = "<tail> " + tokens_copy[j] + " </tail>"
                sequence_list.append(" ".join(tokens_copy))
    return sequence_list