modify the following code in order to perform the threshold only one time.
In particular. i perform a     frequency_threshold = 5  # Keywords with count >5 are kept as-is and a     Category_Limit = 50    # Limit to top 100 keywords per category.
 I do this because i want to obtain AT MAX 50 unique keywords for each category.

is there a more optimal and smart way to do it?
take into account that i want, if possible, to not just cut the top 50 and then put the other keywords to "other" i would like, if possible to force the clustering of each category into EXACTLY 50 unique keywords. can you do it?



import csv
import re
import os
import nltk
import pandas as pd
from collections import defaultdict, Counter
from nltk.stem import WordNetLemmatizer
from gensim.models import KeyedVectors, FastText
from sklearn.cluster import KMeans, AgglomerativeClustering
import numpy as np
from sklearn.preprocessing import OneHotEncoder

# ------------------------------#
#           Setup               #
# ------------------------------#

# Initialize the lemmatizer
lemmatizer = WordNetLemmatizer()

# Import for POS tagging
from nltk import pos_tag, word_tokenize
from nltk.corpus import wordnet as wn
from nltk.corpus import stopwords

# Set of stopwords to filter out
stop_words = set(stopwords.words('english'))

# Function to get WordNet POS tag from treebank tag
def get_wordnet_pos(treebank_tag):
    if treebank_tag.startswith('J'):
        return wn.ADJ  # Adjective
    elif treebank_tag.startswith('V'):
        return wn.VERB  # Verb
    elif treebank_tag.startswith('N'):
        return wn.NOUN  # Noun
    elif treebank_tag.startswith('R'):
        return wn.ADV  # Adverb
    else:
        return None

# Function to load the appropriate word embedding model
def load_model(model_type, model_path):
    if model_type == 'word2vec':
        print("Loading Word2Vec model. This may take a few minutes...")
        model = KeyedVectors.load_word2vec_format(model_path, binary=True)
        print("Word2Vec model loaded successfully.")
    elif model_type == 'glove':
        print("Loading GloVe model. This may take a few minutes...")
        model = {}
        with open(model_path, 'r', encoding='utf-8') as f:
            for line in f:
                values = line.split()
                word = values[0]
                vector = np.asarray(values[1:], dtype='float32')
                model[word] = vector
        print("GloVe model loaded successfully.")
    elif model_type == 'fasttext':
        print("Loading FastText model. This may take a few minutes...")
        model = FastText.load_fasttext_format(model_path)
        print("FastText model loaded successfully.")
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
    return model

# Function to get the vector for a word/phrase
def get_word_vector(word, model, model_type):
    if model_type in ['word2vec', 'fasttext']:
        if word in model:
            return model[word]
    elif model_type == 'glove':
        if word in model:
            return model[word]
    # Handle multi-word phrases by averaging their word vectors
    words = word.split()
    vectors = []
    for w in words:
        if model_type in ['word2vec', 'fasttext'] and w in model:
            vectors.append(model[w])
        elif model_type == 'glove' and w in model:
            vectors.append(model[w])
    if vectors:
        return np.mean(vectors, axis=0)
    else:
        return None

# Function to lemmatize a keyword
def lemmatize_keyword(keyword):
    return lemmatizer.lemmatize(keyword.lower())

# Function to process the CSV file and lemmatize keywords per row, and output lemmatized counts
def lemmatize_keywords(input_file, lemmatized_counts_output_file):
    # Read the CSV file using pandas
    df = pd.read_csv(input_file)
    
    # Ensure 'Image Name' and 'Keywords' columns exist
    if 'Image Name' not in df.columns or 'Keywords' not in df.columns:
        print("Error: 'Image Name' and 'Keywords' columns not found in input CSV.")
        return None, None
    
    per_row_data = []
    keyword_counts = Counter()
    
    # Iterate through each row to extract and lemmatize keywords
    for index, row in df.iterrows():
        image_name = row.get('Image Name', '')
        keywords_str = row.get('Keywords', '')
        if pd.isna(keywords_str):
            keywords_str = ''
        
        # Split keywords by comma and strip whitespace
        extracted_keywords = [kw.strip() for kw in keywords_str.split(',')]
        # Lemmatize keywords
        lemmatized_keywords = [lemmatize_keyword(kw) for kw in extracted_keywords if kw]
        # Update global counts
        keyword_counts.update(lemmatized_keywords)
        
        # Store per-row data
        per_row_data.append({
            'Image Name': image_name,
            'Keywords': ', '.join(extracted_keywords),
            'Lemmatized Keywords': ', '.join(lemmatized_keywords)
        })
    
    # Create a DataFrame from per_row_data
    per_row_df = pd.DataFrame(per_row_data)
    
    # Write the lemmatized keyword counts to CSV
    lemmatized_counts = [{'Lemmatized Keyword': kw, 'Count': count} 
                         for kw, count in keyword_counts.items()]
    lemmatized_counts_df = pd.DataFrame(lemmatized_counts)
    
    # Sort by Count in descending order
    lemmatized_counts_df = lemmatized_counts_df.sort_values(by='Count', ascending=False)
    
    # Write to CSV
    lemmatized_counts_df.to_csv(lemmatized_counts_output_file, index=False)
    print(f"Lemmatized keyword counts saved to '{lemmatized_counts_output_file}'.")
    
    return per_row_df, keyword_counts

# Function to cluster infrequent keywords (count <= threshold)
def cluster_infrequent_keywords(keywords, keyword_counts, model, model_type, num_clusters=5, distance_threshold=None):
    vectors = []
    valid_keywords = []
    for keyword in keywords:
        vector = get_word_vector(keyword, model, model_type)
        if vector is not None:
            vectors.append(vector)
            valid_keywords.append(keyword)
        else:
            print(f"Vector not found for keyword: '{keyword}'")
    
    if not vectors:
        print("No valid vectors found for clustering.")
        return {}
    
    vectors = np.array(vectors)
    
    # Perform Hierarchical Clustering with distance threshold
    if distance_threshold is not None:
        print(f"Performing Agglomerative Clustering with distance_threshold={distance_threshold}")
        clustering_model = AgglomerativeClustering(n_clusters=None, distance_threshold=distance_threshold)
    else:
        print(f"Performing Agglomerative Clustering with n_clusters={num_clusters}")
        clustering_model = AgglomerativeClustering(n_clusters=num_clusters)
    
    labels = clustering_model.fit_predict(vectors)
    
    # Organize keywords by clusters
    clusters = defaultdict(list)
    for label, keyword in zip(labels, valid_keywords):
        clusters[label].append(keyword)
    
    # Select the most frequent keyword in each cluster as the representative
    representative_keywords = {}
    for label, words in clusters.items():
        most_frequent = max(words, key=lambda w: keyword_counts[w])
        representative_keywords[label] = most_frequent
    
    # Create a mapping from keyword to representative
    keyword_to_representative = {}
    for label, words in clusters.items():
        rep = representative_keywords[label]
        for word in words:
            keyword_to_representative[word] = rep
    
    return keyword_to_representative

# Function to perform clustering with frequency threshold
def perform_clustering_with_threshold(keyword_counts, model, model_type, threshold=5, num_clusters=5, distance_threshold=None):
    # Separate frequent and infrequent keywords
    frequent_keywords = [kw for kw, count in keyword_counts.items() if count > threshold]
    infrequent_keywords = [kw for kw, count in keyword_counts.items() if count <= threshold]
    
    print(f"Number of frequent keywords (> {threshold}): {len(frequent_keywords)}")
    print(f"Number of infrequent keywords (<= {threshold}): {len(infrequent_keywords)}")
    
    # Keep frequent keywords as-is
    clustered = {kw: kw for kw in frequent_keywords}
    
    # Cluster infrequent keywords and map them to their representatives
    if infrequent_keywords:
        mapping = cluster_infrequent_keywords(
            infrequent_keywords, keyword_counts, model, model_type,
            num_clusters=num_clusters, distance_threshold=distance_threshold
        )
        clustered.update(mapping)
        print(f"Number of cluster representatives: {len(set(mapping.values()))}")
    else:
        print("No infrequent keywords to cluster.")
    
    return clustered

# Function to generate clustered keywords per row and categorize them based on position,
# mapping non-top keywords to 'Other'
def generate_clustered_keywords_and_categories(per_row_df, clustered_keyword_mapping, category_keyword_counts, remove_duplicates=False):
    top_keywords_per_category = {category: set(counts.keys()) for category, counts in category_keyword_counts.items()}

    # Initialize lists for categories
    categories = ['Object', 'Adjective', 'Activity', 'Place', 'Feeling']
    category_keywords_list = {cat: [] for cat in categories}

    clustered_keywords_list = []

    for idx, row in per_row_df.iterrows():
        lemmatized_keywords = row['Lemmatized Keywords'].split(', ')
        clustered_keywords = [clustered_keyword_mapping.get(kw, kw) for kw in lemmatized_keywords]
        if remove_duplicates:
            # Remove duplicates while preserving order
            clustered_keywords = list(dict.fromkeys(clustered_keywords))
        clustered_keywords_list.append(', '.join(clustered_keywords))

        # Assign keywords to categories based on positions
        for i, category in enumerate(categories):
            if i < len(clustered_keywords):
                keyword = clustered_keywords[i]
                # Map to 'Other' if not in top keywords
                if keyword not in top_keywords_per_category[category]:
                    keyword = 'Other'
                category_keywords_list[category].append(keyword)
            else:
                category_keywords_list[category].append('')  # or 'Other'

    per_row_df['Clustered Keywords'] = clustered_keywords_list

    # Add the category columns to per_row_df
    for category in categories:
        per_row_df[category] = category_keywords_list[category]

    return per_row_df

# Function to categorize and count clustered keywords, and limit to top 50 per category
def categorize_and_count_clustered_keywords(per_row_df, Category_Limit=50):
    category_keyword_counts = {'Object': Counter(), 'Adjective': Counter(), 'Activity': Counter(),
                               'Place': Counter(), 'Feeling': Counter()}
    
    for idx, row in per_row_df.iterrows():
        clustered_keywords = row['Clustered Keywords'].split(', ')
        
        if len(clustered_keywords) > 0:
            category_keyword_counts['Object'][clustered_keywords[0]] +=1
        if len(clustered_keywords) > 1:
            category_keyword_counts['Adjective'][clustered_keywords[1]] +=1
        if len(clustered_keywords) > 2:
            category_keyword_counts['Activity'][clustered_keywords[2]] +=1
        if len(clustered_keywords) > 3:
            category_keyword_counts['Place'][clustered_keywords[3]] +=1
        if len(clustered_keywords) > 4:
            category_keyword_counts['Feeling'][clustered_keywords[4]] +=1
            
    # Limit to top 50 keywords per category
    for category in category_keyword_counts:
        category_keyword_counts[category] = Counter(dict(category_keyword_counts[category].most_common(Category_Limit)))
        
    return category_keyword_counts

# Function to one-hot encode the categorical features
def encode_categorical_features(input_csv_path, output_tsv_path, id_column='Image Name', remove_extension='.jpg', sep='\t', header=False, index=False):

    # Read only the necessary columns to improve efficiency
    columns_to_encode = ['Object', 'Adjective', 'Activity', 'Place', 'Feeling']
    columns_to_read = [id_column] + columns_to_encode
    df = pd.read_csv(input_csv_path, usecols=columns_to_read)

    # Initialize OneHotEncoder with dtype set to integer to output int values
    enc = OneHotEncoder(handle_unknown='ignore', dtype=np.int8)

    # Fit and transform the categorical columns
    emb = enc.fit_transform(df[columns_to_encode])

    # Retrieve the names of the one-hot encoded features
    feature_names = enc.get_feature_names_out(columns_to_encode)

    # Create a DataFrame with the encoded features and ensure they are integers
    emb_df = pd.DataFrame(emb.toarray(), columns=feature_names, dtype=np.int8)

    # Concatenate the id_column with the one-hot encoded features
    result_df = pd.concat([df[[id_column]], emb_df], axis=1)

    # Remove specified extension from id_column
    if remove_extension:
        result_df[id_column] = result_df[id_column].str.replace(remove_extension, '', regex=False)

    # Save the final DataFrame to a TSV file without headers and index
    result_df.to_csv(output_tsv_path, index=index, header=header, sep=sep)


# ------------------------------#
#             Main              #
# ------------------------------#

def main():
    # File paths (replace these with your actual file paths)
    original_input_file = '/home/danilo/Projects/IR_MD/cleaned_/home/danilo/Projects/IR_MD/cleaned_/home/danilo/Projects/IR_MD/TESTING/cleaned_graffe_better_prompt_KeyExt_Qwen2VL7B_20240912-131955_.csv'

    lemmatized_counts_output_file = '/home/danilo/Projects/IR_MD/TESTING/keywords_lemmatized_counts.csv'
    
    # Clustering parameters
    model_type = 'glove'  # Choose from 'word2vec', 'glove', 'fasttext'
    model_path = '/home/danilo/Projects/IR_MD/models/glove.6B.300d.txt'  # Path to the model file
    frequency_threshold = 5  # Keywords with count >5 are kept as-is
    num_clusters = 50       # Number of clusters for infrequent keywords (used if distance_threshold is None)
    Category_Limit = 50    # Limit to top 100 keywords per category
    distance_threshold = None  # Set this to a float value to use hierarchical clustering with a distance threshold
    
    remove_duplicates = True  # Set to True to remove duplicates in clustered keywords per row
    
    # Step 1: Lemmatization and saving lemmatized counts
    per_row_df, keyword_counts = lemmatize_keywords(original_input_file, lemmatized_counts_output_file)
    if per_row_df is None:
        return  # Exit if there was an error
    
    # Step 2: Load the appropriate model
    model = load_model(model_type, model_path)
    
    # Step 3: Perform clustering on infrequent keywords
    clustered_keyword_mapping = perform_clustering_with_threshold(
        keyword_counts, model, model_type, 
        threshold=frequency_threshold, num_clusters=num_clusters, 
        distance_threshold=distance_threshold
    )
    
    # Step 4: Generate 'Clustered Keywords' column
    per_row_df['Clustered Keywords'] = per_row_df['Lemmatized Keywords'].apply(
        lambda kws: ', '.join([clustered_keyword_mapping.get(kw, kw) for kw in kws.split(', ')])
    )
    
    # Step 5: Categorize and count clustered keywords
    category_keyword_counts = categorize_and_count_clustered_keywords(per_row_df, Category_Limit=Category_Limit)
    
    # Step 6: Generate clustered keywords per row and categorize with top keywords only
    per_row_df = generate_clustered_keywords_and_categories(
        per_row_df, clustered_keyword_mapping, category_keyword_counts, remove_duplicates=remove_duplicates
    )
    
    # Print counts for each category
    for category, counts in category_keyword_counts.items():
        print(f"\nCategory: {category}")
        for keyword, count in counts.items():
            print(f"{keyword}: {count}")
    
    # Optionally, write category counts to CSV files
    for category, counts in category_keyword_counts.items():
        category_counts_df = pd.DataFrame({'Keyword': list(counts.keys()), 'Count': list(counts.values())})
        category_counts_df = category_counts_df.sort_values(by='Count', ascending=False)
        category_counts_output_file = f'/home/danilo/Projects/IR_MD/TESTING/{category}_keywords_counts.csv'
        category_counts_df.to_csv(category_counts_output_file, index=False)
        print(f"\n{category} keyword counts saved to '{category_counts_output_file}'.")
    
    # Step 7: Construct the output file name dynamically based on parameters
    # Replace dots in distance_threshold with underscores for file naming
    if distance_threshold is not None:
        distance_threshold_str = str(distance_threshold).replace('.', '')
        clustering_info = f'DistThresh{distance_threshold_str}'
    else:
        clustering_info = f'{num_clusters}Clusters'
    
    dup_info = 'NoDup' if remove_duplicates else 'WithDup'
    
    output_file = f'/home/danilo/Projects/IR_MD/TESTING/{frequency_threshold}Freq{clustering_info}{dup_info}_keywords_clustered.csv'
    
    # Step 8: Write the output CSV
    per_row_df.to_csv(output_file, index=False)
    print(f"\nProcess completed. Output saved to '{output_file}'.")
    print(f"Lemmatized counts saved to '{lemmatized_counts_output_file}'.")
    
    # Step 9: One-hot encode the categorical features
    one_hot_encoded_file = f'/home/danilo/Projects/IR_MD/TESTING/{frequency_threshold}Freq{clustering_info}{dup_info}_keywords_clustered_onehot.csv'
    encode_categorical_features(output_file, one_hot_encoded_file, id_column='Image Name', remove_extension='.jpg', sep=',', header=False, index=False)
    print(f"\nOne-hot encoded features saved to '{one_hot_encoded_file}'.")

if __name__ == "__main__":
    main()