# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# # Getting nouns for every image

# +
import os.path
import re
import sys
import tarfile
import time
import multiprocessing as mp
import itertools
from collections import Counter

import tensorflow.python.platform
from six.moves import urllib
import numpy as np
import tensorflow as tf
import h5py
import glob
import cPickle as pickle
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from tensorflow.python.platform import gfile

from run_inference import predict_star, predict
import pandas as pd
import seaborn as sns

# %matplotlib inline
# %load_ext autoreload
# %autoreload 2
# -

# ## Make one dictionary of noun phrases

a = {1:2,3:4}
b = {4:5,6:7}
a.update(b)
a

# +
files = glob.glob("/data/nearest_neighbor_tagging/idx_to_noun_phrases_*")

combined_dict = {}
for fn in files:
    a = pickle.load(open(fn, 'r'))
    combined_dict.update(a)
    
pickle.dump(combined_dict, open("/data/idx_to_noun_phrases.pkl", 'w'))
# -

# ## Calculate tags for all images

idx_to_noun_phrases = pickle.load(open("/data/idx_to_noun_phrases.pkl", 'r'))

idx_to_noun_phrases = combined_dict

# +
import enchant
import difflib
eng = enchant.Dict("en_US")
br = enchant.Dict("en_GB")
fr = enchant.Dict("fr_FR")
sp = enchant.Dict("es")
ger = enchant.Dict("de_DE")

from nltk.stem.snowball import SnowballStemmer
stemmer = SnowballStemmer("english")

from nltk.corpus import stopwords
stop_words = set(stopwords.words('english'))

# +
image_to_top_words = {}
nearest_image_files = glob.glob("/data/nearest_neighbor_tagging/nearest_neighbors/*.pkl")
ctr = 0

for fn in nearest_image_files:
    print fn
    idx_to_neighbors = pickle.load(open(fn, 'r'))
    for idx in idx_to_neighbors:
        ctr += 1
        print ctr
        if ctr % 10000 == 0: print ctr
            
        # get words from all neighbors
        all_words = list(idx_to_noun_phrases[idx])
        for neighbor in idx_to_neighbors[idx]:
            all_words.extend(idx_to_noun_phrases[neighbor[0]])  # neighbor = (idx, score)
        print len(all_words)

        # check for mispellings, stopwords, etc and stem the word
        all_words_pared = []
        for word in all_words:
            #print word
            if eng.check(word):  # spelled correctly
                if word in stop_words: continue
                all_words_pared.append(str(stemmer.stem(word)))
            else:  # take the top ranked suggestion
                try:
                    new_word = br.suggest(word)[0].lower()
                    new_score = enchant.utils.levenshtein(word, br_word)
                    if new_score > 3: continue  # too far away
                except:
                    continue  # if there are no suggestions from the English dictionary, the word is probably malformed
                if new_word in stop_words: continue 
                    
                # other languages
                if fr.check(word) or ger.check(word) or sp.check(word): continue

                all_words_pared.append(str(stemmer.stem(new_word)))
                
        final_tags = Counter(all_words_pared).most_common(20)
        image_to_top_words[idx] = final_tags
        print final_tags
    break


# +
def process(word):
    if eng.check(word) or br.check(word):  # spelled correctly
        if word in stop_words: 
            return
    else:  # take the top ranked suggestion
        try:
            new_word = br.suggest(word)[0].lower()
            new_score = enchant.utils.levenshtein(word, new_word)
#             eng_word = eng.suggest(word)[0].lower()
#             eng_score = enchant.utils.levenshtein(word, eng_word)
            
#             br_word = br.suggest(word)[0].lower()
#             br_score = enchant.utils.levenshtein(word, br_word)
            
#             if eng_score < br_score:
#                 new_word = eng_word
#                 new_score = eng_score
#             else:
#                 new_word = br_word
#                 new_score = br_score
                
#             print new_word, new_score
        except:
            return  # if there are no suggestions from the English dictionary, the word is probably malformed
        if new_word in stop_words: return 

        # other languages
        if fr.check(word) or ger.check(word) or sp.check(word): return
        print new_word
        
#         try:
#             fr_word = fr.suggest(word)[0].lower()
#             fr_score = enchant.utils.levenshtein(word, fr_word)
#             print fr_word, fr_score
#             if fr_score < new_score: return  # don't worry about this word
#         except: pass

#         try:
#             ger_word = ger.suggest(word)[0].lower()
#             ger_score = enchant.utils.levenshtein(word, ger_word)
#             print ger_word, ger_score
#             if ger_score < new_score: return
#         except: pass

#         try:
#             sp_word = sp.suggest(word)[0].lower()
#             sp_score = enchant.utils.levenshtein(word, sp_word)
#             print sp_word, sp_score
#             if sp_score < new_score: return  
#         except: pass


# -

# %time process("berliner")

enchant.utils.levenshtein("berliner", "Berliner")

# %time eng.suggest("berliner")

# ## Cleaning up set of noun phrases

StarcoderdataJupyterScriptsDedupFiltered