# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
import pandas as pd
import regex, re, sys, nltk

from nltk import word_tokenize, sent_tokenize
from nltk.tokenize import RegexpTokenizer
from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords, wordnet

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation

from gensim.models.hdpmodel import HdpModel
from gensim.corpora.dictionary import Dictionary

from collections import Counter

import matplotlib.pyplot as plt

from pathlib import Path

src_path = str(Path.cwd().parent / "scripts")
sys.path.append(src_path)
# python file with all the functions (located in the src folder)
import topic_classification as tc
# -

filepath = Path.cwd().parent / "data" / "all_speeches_cleaned.txt"
df=pd.read_csv(filepath, usecols=['title','content'])
df.head()

df['text']=df['content'].str.lower()
df.head()

# tokenize into sentences
df['tokenized']=df['text'].apply(lambda text: nltk.sent_tokenize(text))
df[['text','tokenized']].head()

# the '\b(?!\d)' filters out expressions like '9th', since the first character cannot be a number
tokenizer = RegexpTokenizer(r'\b(?!\d)[a-zA-Z]+')
lemmatizer = WordNetLemmatizer()
df['normalized']=df['tokenized'].apply(lambda text: tc.normalize_text(text, tokenizer, lemmatizer))
df[['text','tokenized','normalized']].head()

# +
STOPWORDS = set(stopwords.words('english'))
df['fully_processed'] = df['normalized'].apply(lambda text: tc.remove_stopwords(text, STOPWORDS))
                
cnt = Counter()
for text in df['fully_processed'].values:
    # counts the number of speeches the word is in
    for word in set(text.split()):
        cnt[word] += 1
# words that are in most of the speeches
in_most_speeches = cnt.most_common(155)
in_most_speeches = [x[0] for x in in_most_speeches]

extra = ['mr', 'question', 'sure', 'obama', 'really', 'try', 'lot', 'important', 'million', 'talk', 'va', 'dr', 'romney',
        'folk', 'governor', 'republican', 'king', 'heart'] 

STOPWORDS_extra = set(in_most_speeches + extra)
# remove some words from the stopwords list that migth be important
STOPWORDS_extra = STOPWORDS_extra - set(['war', 'care', 'child', 'family', 'job', 'law', 'protect', 'security', 'power'])

df['fully_processed'] = df['fully_processed'].apply(lambda text: tc.remove_stopwords(text, STOPWORDS_extra))
df[['text','tokenized','normalized','fully_processed']].head()

# +
texts = [text.split() for text in df['fully_processed'].values]

# Create a dictionary
# In gensim a dictionary is a mapping between words and their integer id
dictionary = Dictionary(texts)

# Filter out extremes to limit the number of features
dictionary.filter_extremes(
    no_below=3,
    no_above=0.85,
    keep_n=5000
)

# Create the bag-of-words format (list of (token_id, token_count))
corpus = [dictionary.doc2bow(text) for text in texts]

Hdp_model = HdpModel(corpus=corpus, id2word=dictionary)
# -

from pprint import pprint
pprint(Hdp_model.print_topics(num_words=10))

StarcoderdataJupyterScriptsDedupFiltered