"""
Tf-IDF from scratch
"""
import math
import pandas as pd
from collections import Counter, OrderedDict
from nltk.tokenize import RegexpTokenizer

#Tf-idf definitions:
#tf(t,d) = count of t in d / number of words in d
#idf(t) = log(N/(df + 1)) - where N = number of total documents, df = number of documents term appears in
#tf-idf = tf(t,d) * idf(t)

#Documents
docs = ['The “cat pattern,” established very early in the evolution of modern mammals, was a successful one: early cats were already typical in form at a time when the ancestors of most other modern mammalian types were scarcely recognizable. ',
        "It's no secret that cats love to sleep. In fact, the average kitty sleeps for a whopping 12 to 16 hours per day! Because cats are nocturnal creatures, they snooze most of the day away—and who can blame them?",
        "Believe it or not, some breeds of cats also like to play in water. If your cat loves to splash around, try filling up the bathtub or a kiddie pool with a few inches of water during her playtime. This activity will help her get a little bit of exercise and cool off on particularly hot days."]

#Tokenizing documents to only include lower case words and ignoring any punctuations
tokenizer = RegexpTokenizer(r'\w+')

docs_tokenized = []#list of all documents tokenized
word_set = set() #set of all unique words across all documents
for doc in docs:
  tokenized = tokenizer.tokenize(doc)
  ls = [word.lower() for word in tokenized]
  word_set.update(ls)
  docs_tokenized.append(ls)


#Generating Term Frequency
tf_counts = []
for tokens in docs_tokenized:
  counts = Counter(tokens)
  tf_counts.append(counts)

#Generating Document Frequency
doc_freq = []
for doc_tokens in tf_counts:
  doc_freq_counter = Counter()
  for word, count in doc_tokens.items():
    for doc_number in range(len(tf_counts)):
      if word in tf_counts[doc_number]:
        doc_freq_counter[word] += 1
  doc_freq.append(doc_freq_counter)


#Creating an index of location of words to the column names
index = OrderedDict.fromkeys(word_set, None)
index_items = sorted(list(index.keys())) #list of column names (i.e. terms) in order, 'sorted' to get consistent order of column names

#Generating tf-idf matrix
#initializing an empty matrix
tf_idf_matrix = [[0 for _ in range(len(word_set))]  for _ in range(len(docs))]
for pos, word in enumerate(index_items):

  #Filling in TF
  for doc_pos, doc in enumerate(tf_counts):
    if word in doc:
      tf_idf_matrix[doc_pos][pos] = doc[word]
  
  #Filling in IDF
  for doc_pos, doc in enumerate(doc_freq):
    if word in doc:
      tf_idf_matrix[doc_pos][pos] *= math.log(len(docs) / (doc[word] + 1), 2)

# print(tf_idf_matrix)
df = pd.DataFrame(tf_idf_matrix, columns = index_items, index=['Doc {}'.format(num + 1) for num in range(len(docs))])
df.head()
