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

## Overview
This script is designed to perform topic modeling on textual data using Latent Dirichlet Allocation (LDA). It processes `.txt` files to extract key topics and their distributions.

## Folder Structure
- The script is located in the 'Code' folder.
- The 'Data' folder should contain a subfolder 'LDA_textfiles' with `.txt` files for analysis.
- The 'Results' folder will store the outputs: 'topics.csv' for topic keywords and 'topic_distribution.csv' for topic distributions per document.

## Prerequisites
- Ensure that all textual data files (.txt) are placed within the 'LDA_textfiles' subfolder under 'Data'.
- Required Python libraries including sklearn, nltk, and pandas must be installed. Run `pip install scikit-learn nltk pandas` to install these if not already available.

## Processing Steps
1. **Text Preprocessing**: Cleans and preprocesses the text by removing gibberish, emails, and applying tokenization and lemmatization.
2. **Document Vectorization**: Converts the preprocessed documents into a vector form that can be used by the LDA model, either using Count Vectorizer or TF-IDF Vectorizer.
3. **Topic Modeling**: Applies Latent Dirichlet Allocation to identify the main topics from the vectorized data.
4. **Result Compilation**: Saves the identified topics and their corresponding word distributions into 'topics.csv'. Also maps each document to its topic distribution and saves this mapping in 'topic_distribution.csv'.

## Output
- 'topics.csv': Contains identified topics along with their top words.
- 'topic_distribution.csv': Contains the distribution of topics for each processed document.
- Both files are saved in the 'Results' folder.

"""


import os
import re
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.decomposition import LatentDirichletAllocation
import nltk
nltk.download('stopwords')
nltk.download('wordnet')
from nltk.corpus import stopwords
from nltk.corpus import wordnet
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
import numpy as np
import pandas as pd
import string
np.random.seed(123)
#------------ Set up logging ------------ 
import log_config
log_config.setup_logging()
import logging
logging.info("Running script 'LDA Topics and Topic Distributions.py'")


#------------ Folder locations here ------------ 
# Source directory (User-specified)
folder_path = '../Data/LDA_textfiles'

# Define paths for generated results files
topics = '../Results/topics.csv'  # Save the topics and their top words
topic_distribution_csv = '../Results/topic_distribution.csv'  # Save topic distribution

#------------ Document Processing and Topic Modeling Workflow ------------

def preprocess_document(doc):
    lemmatizer = WordNetLemmatizer()
    stop_words = set(stopwords.words('english'))
    custom_stop_words = ['cookies', 'website','www', 'com']
    stop_words.update(custom_stop_words)
    table = str.maketrans('', '', string.punctuation)
    doc = re.sub('\S*@\S*\s?', '', doc)
    doc_lower = doc.lower()
    tokens = word_tokenize(doc_lower)
    tokens = [lemmatizer.lemmatize(word.translate(table)) for word in tokens if word not in stop_words and word.isalpha()]
    return ' '.join(tokens)

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

def vectorize_documents(processed_docs, vectorizer_type="count"):
    if vectorizer_type == "tfidf":
        vectorizer = TfidfVectorizer(lowercase=True)
    else:
        vectorizer = CountVectorizer(lowercase=True)
        
    matrix = vectorizer.fit_transform(processed_docs)
    return matrix, vectorizer

def perform_topic_modeling(count_matrix, vectorizer, document_files, num_topics):
    # 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)}')

    # Save the topics and their top words
    topics_df = pd.DataFrame([(idx, ' | '.join(words)) for idx, words in enumerate(top_words)], columns=['Topic', 'Top_Words'])
    topics_df.to_csv(topics, index=False)    

    # Map documents to topics
    doc_topic_mapping = {}
    for doc_idx, doc in enumerate(lda_matrix):
        # We store the full topic distribution of the document
        doc_topic_mapping[document_files[doc_idx]] = doc
    return doc_topic_mapping

# 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 each document into a string
documents = []
for document_file in document_files:
    with open(os.path.join(folder_path, document_file), 'r') as f:
        documents.append(f.read())
processed_docs = preprocess_documents(documents)
count_matrix, vectorizer = vectorize_documents(processed_docs, "count")

# Set the number of topics
num_topics = 10

# Perform topic modeling using LDA
doc_topic_mapping = perform_topic_modeling(count_matrix, vectorizer, document_files, num_topics)

# Create a DataFrame to store the topic distribution scores for each document
topic_columns = [f"Topic_{i}" for i in range(num_topics)]
topic_distribution_df = pd.DataFrame(columns=['Document'] + topic_columns)

topic_distribution_list = []
for doc, topic_distribution in doc_topic_mapping.items():
    row_data = [doc] + list(topic_distribution)
    topic_distribution_list.append(row_data)

topic_distribution_df = pd.DataFrame(topic_distribution_list, columns=['Document'] + topic_columns)


# Save the topic distribution DataFrame to a CSV file
topic_distribution_df.to_csv(topic_distribution_csv, index=False)
print(f'All results are saved in the Results folder')

#------------ Log completion of the script ------------
logging.info("Finished running 'LDA Topics and Topic Distributions.py'")