import spacy
import random


class Delexicalizer:
    def __init__(self, prob_pos_tag, prob_ner_tag, spacy_model="pt_core_news_sm") -> None:

        if not spacy_model in spacy.util.get_installed_models():
            spacy.cli.download(spacy_model)

        self.nlp = spacy.load(spacy_model, enable=["parser", "tagger", "ner"])

        if prob_pos_tag < 0 or prob_pos_tag > 1:
            raise ValueError("prob_pos_tag must be between 0 and 1")

        if prob_ner_tag < 0 or prob_ner_tag > 1:
            raise ValueError("prob_ner_tag must be between 0 and 1")

        self.prob_pos_tag = prob_pos_tag
        self.prob_ner_tag = prob_ner_tag

    def delexicalize(self, text):
        doc = self.nlp(text)

        list_tokens = []

        for token in doc:

            if token.ent_type > 0 and random.uniform(0, 1) < self.prob_ner_tag:
                list_tokens.append(token.ent_type_)

            elif random.uniform(0, 1) < self.prob_pos_tag:
                list_tokens.append(token.pos_)

            else:
                list_tokens.append(token.text)

        return ' '.join(list_tokens)