import spacy
from spacy.language import Language
import regex


@Language.component("entity_punctuation_removal")


def entity_punctuation_removal(doc):
    # On liste nos entités
    ents = list(doc.ents)

    i = 0
    while i < len(ents):
        current_ent = ents[i]

        # On attrape l'entité si :
        # - Elle consiste en un unique signe de ponctuation
        # - Et que son tag IOB est B (début d'entité)
        # - Et qu'elle n'est pas suivie par une entité de tag IOB I (suite d'entité) OU qu'elle n'est suivie par aucun token
        # Si l'entité est attrapée, elle est supprimée du doc
        if i + 1 < len(ents) and regex.match(r'^\p{P}$', current_ent.text) and current_ent.root.ent_iob_ == "B" :
          ents.pop(i)
        elif i == len(ents) - 1 and regex.match(r'^\p{P}$', current_ent.text) and current_ent.root.ent_iob_ == "B" :
          ents.pop(i)
        else:
            i += 1

    # On met à jour le doc avec les entités modifiées
    doc.ents = tuple(ents)

    return doc

Language.component("entity_punctuation_removal", func=entity_punctuation_removal)