"""
-----------------------------------
# Instructions for Code Execution
-----------------------------------

## Overview
This script analyzes textual data through Latent Dirichlet Allocation (LDA) to model topics and evaluates the model using perplexity and coherence metrics.

## Folder Structure
- The script is located in the 'Code' folder.
- The `Data` folder must contain a subfolder named 'LDA_textfiles' with `.txt` files to be processed.

## Prerequisites
- Necessary `.txt` files should be placed within the 'LDA_textfiles' folder.
- Required libraries (`scikit-learn`, `gensim`, `nltk`, `numpy`) should be installed.

## Processing Steps
1. **Text Cleaning**: Cleans text by removing emails, punctuation, and stopwords.
2. **Text Vectorization**: Transforms text into vector form using Count Vectorization or TF-IDF.
3. **Topic Modeling**: Executes LDA to identify topics and prints top words for each topic.
4. **Model Evaluation**: Calculates and displays the perplexity of the LDA model and coherence score to assess topic relevance and clarity.

## Output
- The script directly prints the top words of each topic, along with the perplexity and coherence scores.

"""

import os
import re
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.decomposition import LatentDirichletAllocation
from gensim.models import CoherenceModel
from gensim.models import LdaModel
from gensim import corpora
import nltk
from nltk.corpus import wordnet
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
import numpy as np
import nltk
from nltk.corpus import stopwords
import string
np.random.seed(123)
#------------ Set up logging ------------ 
import log_config
log_config.setup_logging()
import logging
logging.info("Running script 'Perplexity & Coherence LDA.py'")


#------------ Folder location here ------------ 
folder_path = '../Data/LDA_textfiles'

#------------ Document Preprocessing, Topic Modeling Workflow, and Calculation of Perplexity and Coherence of the LDA Model ------------

def preprocess_document(doc):
    lemmatizer = WordNetLemmatizer()  # Initialize the WordNetLemmatizer
    stop_words = set(stopwords.words('english'))  # Set stop words
    custom_stop_words = ['cookies', 'website','www', 'com']   # Add custom stop words
    stop_words.update(custom_stop_words)
    table = str.maketrans('', '', string.punctuation)   # Prepare translation table to remove punctuation
    doc = re.sub('\S*@\S*\s?', '', doc)  # Removing emails
    doc_lower = doc.lower()   # Convert document to lowercase and tokenize
    tokens = word_tokenize(doc_lower)
    tokens = [lemmatizer.lemmatize(word.translate(table)) for word in tokens if word not in stop_words and word.isalpha()] # Remove punctuation, stopwords, and lemmatize tokens
    return tokens  

def preprocess_documents(documents):
    return [preprocess_document(doc) for doc in documents]

def vectorize_documents(processed_docs, vectorizer_type="count"):
    documents = [' '.join(doc) for doc in processed_docs]  # Convert token lists back to strings
    if vectorizer_type == "tfidf":
        vectorizer = TfidfVectorizer(lowercase=True)
    else:
        vectorizer = CountVectorizer(lowercase=True)
    matrix = vectorizer.fit_transform(documents)
    return matrix, vectorizer


def perform_topic_modeling(count_matrix, vectorizer, processed_docs, num_topics=10):
    # Apply LatentDirichletAllocation for LDA
    lda_model = LatentDirichletAllocation(n_components=num_topics,learning_decay=0.5, random_state=0)
    lda_matrix = lda_model.fit_transform(count_matrix)

    # Get the most important features (words) for each topic
    feature_names = vectorizer.get_feature_names_out()
    top_words = []
    for topic_idx, topic in enumerate(lda_model.components_):
        top_words.append([feature_names[i] for i in topic.argsort()[:-21:-1]])

    # Print the topics and their top words
    for topic_idx, words in enumerate(top_words):
        print(f'Topic {topic_idx}: {" | ".join(words)}')
        
    # Print perplexity score
    print("Perplexity: ", lda_model.perplexity(count_matrix))

    # Compute Coherence Score using c_v
    dictionary = corpora.Dictionary(processed_docs)
    corpus = [dictionary.doc2bow(text) for text in processed_docs]
    lda_model_gensim = LdaModel(corpus=corpus, id2word=dictionary, num_topics=num_topics, random_state=0)
    coherence_model_lda = CoherenceModel(model=lda_model_gensim, texts=processed_docs, dictionary=dictionary, coherence='c_v')
    coherence_lda = coherence_model_lda.get_coherence()
    print('\nCoherence Score: ', coherence_lda)
    lda_model_gensim = LdaModel(corpus=corpus, id2word=dictionary, num_topics=num_topics, random_state=0)
    
if __name__ == '__main__':
# Get the list of document file names in the folder
 document_files = [file for file in os.listdir(folder_path) if file.endswith('.txt')]

# Read the documents into a list
 documents = []
 documents_tokenized = []
 for file in document_files:
    file_path = os.path.join(folder_path, file)
    with open(file_path, 'r', encoding='utf-8') as doc_file:
        document = doc_file.read()
        documents.append(document)
        documents_tokenized.append(word_tokenize(document))

# Preprocess the documents
 processed_docs = preprocess_documents(documents)
 count_matrix, vectorizer = vectorize_documents(processed_docs, "count")

# Perform topic modeling using LDA
 perform_topic_modeling(count_matrix, vectorizer, processed_docs, num_topics=10)

#------------ Log completion of the script ------------
logging.info("Finished running 'Perplexity & Coherence LDA.py'")