from scipy.stats import mannwhitneyu

df1 = pd.read_csv(WORD_FREQUENCY_CSV_FILENAME1, index_col=0) # read in the CSV
df1.rename(columns={'Unnamed: 0': 'Text'}, inplace=True) # add a label to the first column
df1 = df1.fillna(0) # replace NaNs with zeroes.

df2 = pd.read_csv(WORD_FREQUENCY_CSV_FILENAME2, index_col=0) # read in the CSV
df2.rename(columns={'Unnamed: 0': 'Text'}, inplace=True) # add a label to the first column
df2 = df2.fillna(0) # replace NaNs with zeroes.

# Make "dummy" rows of all zeroes for any words that only appear in one corpus and not the other
missingInCorpus1 = []
missingInCorpus2 = []
for i in range(0, df1.shape[1]):
    missingInCorpus1.append(0)    

for i in range(0, df2.shape[1]):
    missingInCorpus2.append(0)

# Iterate over the wordlist and the two corpora, and output to csv
with open(CORPUS_COMPARISON_FILENAME, 'w', newline='') as csvfile:
    writer = csv.writer(csvfile, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
    writer.writerow(['word', 'Mann Whitney U Value', 'Mann Whitney p-value'])
    with open(WORD_LIST_FILENAME, 'r', encoding="utf-8") as f:
        for word in f:
            word = word.strip()
            if (word in df1.index):
                countsInCorpus1 = df1.loc[word].values
            else:
                countsInCorpus1 = missingInCorpus1
            if (word in df2.index):
                countsInCorpus2 = df2.loc[word].values
            else:
                countsInCorpus2 = missingInCorpus2
            try:
                mw = mannwhitneyu(countsInCorpus1, countsInCorpus2)
                mwStat = mw.statistic
                mwP = mw.pvalue
            except ValueError:
                mwStat = -1 # in case of ties, Mann-Whitney cannot rank, and so cannot calculate U
                mwP = -1
            writer.writerow([word, mwStat, mwP])
print("Done.")