import numpy as np
import pandas as pd

import nltk
nltk.download('brown')
from nltk.corpus import brown
from nltk.corpus import stopwords

import collections
from string import digits, punctuation

print(type(brown.words()))
words = list(brown.words())
print(f"Length of raw Brown corpus: {len(words)}")

# clean text data
def preprocessing(corpus):
    for i in range(len(corpus)):
        corpus[i] = corpus[i].translate(str.maketrans('', '', punctuation))
        #corpus[i] = corpus[i].translate(str.maketrans('', '', digits))
        corpus[i] = corpus[i].lower()
    
    stop_words = set(stopwords.words())
    corpus = [w for w in corpus if w!='' if not w in stop_words] #remove empty string and stopwords
    return corpus

words = preprocessing(words)

def extract_frequent_words(word_list, freq=1000):
    count = collections.Counter(word_list)
    word_count = {'Word':list(count.keys()), 'Count':list(count.values())}
    df = pd.DataFrame(data=word_count)
    df = df.sort_values(by='Count', ascending=False)
    selected = df[0:freq]
    selected = selected.sort_index()
    return selected

C_df = extract_frequent_words(words, freq=1000)
C_list = C_df.Word.values

V_df = extract_frequent_words(words, freq=5000)
V_df.columns = ['V_words', 'window_Count']
V_list = V_df.V_words.values

# track all the positions that word w showed up in the text stream
V_pos = [[x for x, n in enumerate(words) if n == w] for w in V_list]

Windows = []
context_word_count = []
for pos in V_pos:
    window = [] # context words surrounding w for w in V
    for i in pos:
        if i==0:
            cur_window = words[1:3]
        elif i==1:
            cur_window = [words[0]] + words[2:4]
        else:
            cur_window = words[i-2:i] + words[i+1:i+3]
        
        # exclude duplicate context words in a single window
        cur_unique = []
        for j in cur_window:
            if j not in cur_unique:
                cur_unique.append(j)
        
        window += cur_unique
    
    # count occurrence of each context word for words in window w
    context_word_count += list(collections.Counter(window).values())
    
    # remove duplicated context words
    context_word_unique = []
    for k in window:
        if k not in context_word_unique:
            context_word_unique.append(k)
            
    Windows.append(context_word_unique)

V_df['context_words'] = Windows
V_df = V_df.explode('context_words')
V_df['countext_word_count'] = context_word_count
V_df['Pr(c|w)'] = V_df['countext_word_count'].astype('float')/V_df['window_Count']
V_df['in_C'] = [1 if i in C_list else 0 for i in V_df.context_words.values]

# drop all the rows with context words that are not in our shorter list C
V_df = V_df[V_df.in_C==1]

# drop boolean column 'in_C' after we collect all context words belonged to C
V_df.drop(columns='in_C')

# calculate distribution of context words in C
C_list_uniq = list(V_df.context_words.unique())
Pr_C = {}
for c in C_list_uniq:
    Pr_C[c] = words.count(c) / len(words)

V_df['Pr(c)'] = V_df.loc[:, 'context_words'].apply(lambda x: Pr_C[x])
V_df

import math 

def pointwise_mutual_info(row):
    a = row['Pr(c|w)']
    b = row['Pr(c)']
    res = math.log(a/b)
    return max(0, res)

V_df['Phi'] = V_df.apply(pointwise_mutual_info, axis =1)
V_df

sparse_table = pd.pivot_table(V_df, index = 'V_words', columns = 'context_words', values = 'Phi')
sparse_table

# fill NaN with 0s
sparse_table = sparse_table.fillna(0)

from sklearn.decomposition import TruncatedSVD
#from sklearn.random_projection import sparse_random_matrix

X = np.asarray(sparse_table)

svd = TruncatedSVD(n_components = 100, n_iter = 5, random_state = 42)
svd.fit(X)

X_new = svd.fit_transform(X)
df_new = pd.DataFrame(X_new, index = sparse_table.index)
df_new

from sklearn.metrics.pairwise import cosine_similarity

S = 1 - cosine_similarity(X_new, X_new)
S_df = pd.DataFrame(S, index=sparse_table.index, columns=sparse_table.index)

np.fill_diagonal(S_df.values, 1)
S_df

to_check = ['communism', 'autumn', 'cigarette', 'pulmonary', 'mankind', 
            'africa', 'chicago', 'revolution', 'september', 'chemical', 
            'detergent', 'dictionary', 'storm', 'worship']

S_dict = {}
for t in to_check:
    S_dict[t] = S_df[t].idxmin()
    
for d in S_dict:
    print (f'{d} ------> {S_dict[d]}')