import nltk
from nltk.corpus import wordnet

from nlpaug.augmenter.word import WordAugmenter
from nlpaug.util import Action, PartOfSpeech


class WordNetAug(WordAugmenter):
    def __init__(self, name='WordNet_Aug', aug_min=1, aug_p=0.3, lang='eng', tokenizer=None, stopwords=[], verbose=0):
        super(WordNetAug, self).__init__(
            action=Action.SUBSTITUTE, name=name, aug_p=aug_p, aug_min=aug_min, tokenizer=tokenizer, stopwords=stopwords,
            verbose=verbose)

        self.model = self.get_model()
        self.lang = lang

    def skip_aug(self, token_idxes, tokens):
        results = []
        for token_idx in token_idxes:
            # Some word does not come with synonym. It will be excluded in lucky draw.
            if tokens[token_idx][1] not in ['DT']:
                results.append(token_idx)

        return results

    def substitute(self, text):
        results = []

        tokens = self.tokenizer(text)

        pos = nltk.pos_tag(tokens)

        aug_idexes = self._get_aug_idxes(pos)

        for i, token in enumerate(tokens):
            # Skip if no augment for word
            if i not in aug_idexes:
                results.append(token)
                continue

            word_poses = PartOfSpeech.pos2wn(pos[i][1])
            synets = []
            if word_poses is None or len(word_poses) == 0:
                # Use every possible words as the mapping does not defined correctly
                synets.extend(self.model.synsets(pos[i][0], lang=self.lang))
            else:
                for word_pos in word_poses:
                    synets.extend(self.model.synsets(pos[i][0], pos=word_pos, lang=self.lang))

            augmented_data = []
            for synet in synets:
                for candidate in synet.lemma_names(lang=self.lang):
                    if candidate.lower() != token.lower():
                        augmented_data.append(candidate)

            if len(augmented_data) == 0:
                results.append(token)
            else:
                candidate = self.sample(augmented_data, 1)[0]
                candidate = candidate.replace("_", " ").replace("-", " ").lower()
                results.append(self.align_capitalization(token, candidate))

        return self.reverse_tokenizer(results)

    def get_model(self):
        try:
            # Check whether wordnet package is downloaded
            wordnet.synsets('computer')
        except Exception:
            nltk.download('wordnet')

        try:
            # Check whether POS package is downloaded
            nltk.pos_tag('computer')
        except Exception:
            nltk.download('averaged_perceptron_tagger')

        return wordnet

@@ -130,17 +130,68 @@
 OfSpeech
+, Warning, WarningName, WarningCode, WarningMessage
 %0A
-
 %0A%0Aclass 
@@ -354,24 +354,8 @@
 per(
-WordNetAug, self
 ).__
@@ -577,16 +577,110 @@
 = lang%0A%0A
+        ### TODO: antonym: https://arxiv.org/pdf/1809.02079.pdf%0A        self.synonyms = True%0A%0A
     def 
@@ -980,16 +980,700 @@
 esults%0A%0A
+    def _get_aug_idxes(self, tokens):%0A        aug_cnt = self.generate_aug_cnt(len(tokens))%0A        word_idxes = %5Bi for i, t in enumerate(tokens) if t%5B0%5D not in self.stopwords%5D%0A        word_idxes = self.skip_aug(word_idxes, tokens)%0A        if len(word_idxes) == 0:%0A            if self.verbose %3E 0:%0A                exception = Warning(name=WarningName.OUT_OF_VOCABULARY,%0A                                    code=WarningCode.WARNING_CODE_002, msg=WarningMessage.NO_WORD)%0A                exception.output()%0A            return None%0A        if len(word_idxes) %3C aug_cnt:%0A            aug_cnt = len(word_idxes)%0A        aug_idexes = self.sample(word_idxes, aug_cnt)%0A        return aug_idexes%0A%0A
     def 
@@ -1799,33 +1799,32 @@
 %0A%0A        aug_id
-e
 xes = self._get_
@@ -1954,17 +1954,16 @@
 n aug_id
-e
 xes:%0A   
@@ -2563,20 +2563,16 @@
         
-for 
 candidat
@@ -2576,45 +2576,339 @@
 date
- in synet.lemma_names(lang=self.lang)
+s = %5B%5D%0A                for lema in synet.lemmas():%0A                    if self.synonyms:%0A                        candidates.append(lema.name())%0A                    else:%0A                        if lema.antonyms():%0A                            candidates.append(lema.antonyms()%5B0%5D.name())%0A%0A                for candidate in candidates
 :%0A  

Fix #14