# Program 14: TFIDF - Term Frequency and Inverse Document Frequency
#   TF(W)    = number of times word W occurs in a document
#   DF(W)    = number of documents in which W occurs
#   IDF(W)   = N / DF(W), where N is the number of documents
#   TFIDF(W) = TF(W) x log(IDF(W)), with log to the base 10

import numpy as np
import pandas as pd

documents = ["Aman and Avni are stressed",
             "Aman went to a therapist",
             "Avni went to download a health chatbot"]

tokens = []
for d in documents:
    tokens.append(d.lower().split())
dictionary = []
for t in tokens:
    for word in t:
        if word not in dictionary:
            dictionary.append(word)

pd.set_option("display.width", 160)
pd.set_option("display.max_columns", None)

# Term frequency: the document vector table from Bag of Words.
tf = []
for t in tokens:
    row = []
    for word in dictionary:
        row.append(t.count(word))
    tf.append(row)
tf = pd.DataFrame(tf, columns=dictionary, index=["Doc 1", "Doc 2", "Doc 3"])
print("Term frequency (TF):")
print(tf)
print()

N = len(documents)
df = (tf > 0).sum()
print("Document frequency (DF):")
print(df.to_frame().T.to_string(index=False))
print()

idf = N / df
print("Inverse document frequency (IDF = N / DF), N =", N)
print(idf.round(2).to_frame().T.to_string(index=False))
print()

tfidf = tf * np.log10(idf)
print("TFIDF = TF x log10(IDF):")
print(tfidf.round(3))
print()

print("Highest TFIDF in each document:")
for i in range(N):
    row = tfidf.iloc[i]
    best = row.max()
    words = list(row[row == best].index)
    print("  Doc", i + 1, ":", ", ".join(words), "(" + str(round(best, 3)) + ")")
print()

# The handbook's two worked examples, for a corpus of 10 documents.
print("'and' in 10 of 10 documents      : log10(10/10) =", round(np.log10(10 / 10), 3))
print("'pollution' in 3 of 10 documents : log10(10/3)  =", round(np.log10(10 / 3), 3))