i need you to fix the following code.
I expect that, given the 5 categories, if i select category_limit to 200, the onehot encoded file will ahve rows composed by image name and (200*5) columns for the one hot. if i select category limit to 50, i will have (50*5) columns.
it seems not the case. there are many more columns than there should be.

i want that the onehot encoded file has exactly the requested number of categories, while it seems that it calculates not the reduced, but all the keywords.

Please provide the full entire corrected script.

import os
import re
import csv
import argparse
import logging
import nltk
import pandas as pd
from collections import defaultdict, Counter
from nltk.stem import WordNetLemmatizer
from typing import List, Dict, Tuple, Set
import numpy as np
from sklearn.preprocessing import OneHotEncoder

# Import the suggested libraries
from sentence_transformers import SentenceTransformer
from sklearn.cluster import AgglomerativeClustering
from rich.logging import RichHandler
from rich.progress import track
import torch


def download_nltk_data():
    """
    Ensures that required NLTK data packages are downloaded.
    """
    nltk_data_packages = ['wordnet']
    for package in nltk_data_packages:
        try:
            nltk.data.find(f'corpora/{package}')
            logging.info(f"NLTK package '{package}' already exists.")
        except LookupError:
            nltk.download(package)
            logging.info(f"Downloaded NLTK package '{package}'.")


def parse_arguments():
    """
    Parses command-line arguments.
    """
    parser = argparse.ArgumentParser(description='Process and cluster keywords from a CSV file.')
    parser.add_argument('--input_file', type=str, default="/home/danilo/Projects/IR_MD/utils/_testing_publication_cleaned_code/clothing_qwen_original.csv", required=False,
                        help='Path to the input CSV file.')
    parser.add_argument('--save_lemmatized', action='store_true', help='Save lemmatized keywords and counts.')
    parser.add_argument('--model_name', type=str, default='all-MiniLM-L6-v2',
                        help='Name of the SentenceTransformer model to use for clustering. Default is "all-MiniLM-L6-v2".')
    parser.add_argument('--device', type=str, default='cuda',
                        help='Device to run the model on ("cuda" or "cpu"). Default is "cuda".')
    parser.add_argument('--frequency_threshold', type=int, default=0,
                        help='Frequency threshold for keywords. Set to 0 to cluster all keywords. Default is 0.')
    parser.add_argument('--num_clusters', type=int, default=None,
                        help='Number of clusters for keywords. Default is None.')
    parser.add_argument('--distance_threshold', type=float, default=0.6,
                        help='Distance threshold for clustering. Default is 0.2.')
    parser.add_argument('--category_limit', type=int, default=50,
                        help='Limit of top N keywords per category. Default is 200.')
    parser.add_argument('--remove_duplicates', action='store_true',
                        help='Remove duplicates in clustered keywords per row.')
    return parser.parse_args()


def main():
    # Parse command-line arguments
    args = parse_arguments()
    input_file = args.input_file
    save_lemmatized = args.save_lemmatized
    model_name = args.model_name
    device = args.device
    frequency_threshold = args.frequency_threshold
    num_clusters = args.num_clusters
    distance_threshold = args.distance_threshold
    category_limit = args.category_limit
    remove_duplicates = args.remove_duplicates

    # Check if input file exists
    if not os.path.exists(input_file):
        raise FileNotFoundError(f"Input file '{input_file}' does not exist.")

    # Set up logging
    input_dir = os.path.dirname(os.path.abspath(input_file))
    input_filename = os.path.basename(input_file)
    input_name_no_ext = os.path.splitext(input_filename)[0]
    processed_dir = os.path.join(input_dir, 'processed')
    os.makedirs(processed_dir, exist_ok=True)
    log_file = os.path.join(processed_dir, f'log_{input_name_no_ext}.log')

    # Remove any existing handlers
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    logging.basicConfig(
        level=logging.INFO,
        format='%(message)s',
        datefmt='[%X]',
        handlers=[
            RichHandler(),               # For console output with Rich
            logging.FileHandler(log_file)  # For file logging
        ]
    )

    logging.info(f"Processing input file: {input_file}")

    # Download NLTK data if necessary
    download_nltk_data()

    # Initialize the lemmatizer
    lemmatizer = WordNetLemmatizer()

    # Labels (customizable)
    labels = ['Type', 'Color', 'Wear Location', 'Material', 'Style']  # Modify as needed

    # Define output file paths
    output_file = os.path.join(processed_dir, f'clustered_{input_filename}')
    one_hot_encoded_file = os.path.join(processed_dir, f'onehot_{input_name_no_ext}.tsv')
    lemmatized_counts_output_file = os.path.join(processed_dir, f'lemmatized_keyword_counts.csv')

    # ------------------------------#
    # Step 1: Cleaning #
    # ------------------------------#
    logging.info("Starting CSV cleaning...")
    df_clean = clean_csv(input_file, labels)
    logging.info("CSV cleaning completed.")

    # ------------------------------#
    # Step 2: Lemmatization #
    # ------------------------------#
    logging.info("Starting lemmatization...")
    if save_lemmatized:
        per_row_df, keyword_counts = lemmatize_keywords(df_clean, lemmatized_counts_output_file, lemmatizer)
    else:
        per_row_df, keyword_counts = lemmatize_keywords(df_clean, None, lemmatizer)
    # Check if per_row_df is None
    if per_row_df is None:
        logging.error("Error in lemmatization.")
        return

    # ------------------------------#
    # Step 3: Load Model #
    # ------------------------------#
    if device == 'cuda' and not torch.cuda.is_available():
        logging.warning("CUDA is not available, switching to CPU.")
        device = 'cpu'
    model = load_model(model_name, device=device)

    # ------------------------------#
    # Step 4: Clustering #
    # ------------------------------#
    logging.info("Starting clustering of keywords...")
    clustered_keyword_mapping = perform_clustering_with_threshold(
        keyword_counts, model,
        threshold=frequency_threshold, num_clusters=num_clusters,
        distance_threshold=distance_threshold
    )
    logging.info("Keyword clustering completed.")

    # Extract cluster representatives
    cluster_reps = set(clustered_keyword_mapping.values())
    logging.info(f"Cluster representatives extracted: {len(cluster_reps)} representatives.")

    # ------------------------------#
    # Step 5: Clustered Keywords #
    # ------------------------------#
    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 6: Categorization #
    # ------------------------------#
    logging.info("Starting categorization of clustered keywords...")
    category_keyword_counts = categorize_and_count_clustered_keywords(
        per_row_df, labels=labels, category_limit=category_limit
    )
    logging.info("Categorization completed.")

    # ------------------------------#
    # Step 7: Final Processing #
    # ------------------------------#
    logging.info("Generating final clustered keywords and categories...")
    top_keywords_per_category = {label: set(counts.keys()) for label, counts in category_keyword_counts.items()}

    per_row_df, other_count_rows = generate_clustered_keywords_and_categories(
        per_row_df, clustered_keyword_mapping, top_keywords_per_category, labels=labels,
        remove_duplicates=remove_duplicates, cluster_reps=cluster_reps,
        category_keyword_counts=category_keyword_counts, category_limit=category_limit
    )
    logging.info("Final processing completed.")

    # ------------------------------#
    # Step 8: Save Results #
    # ------------------------------#
    per_row_df.to_csv(output_file, index=False)
    logging.info(f"Process completed. Output saved to '{output_file}'.")

    # ------------------------------#
    # Step 9: One-Hot Encoding #
    # ------------------------------#
    logging.info("Starting one-hot encoding of categorical features...")
    encode_categorical_features(
        per_row_df, one_hot_encoded_file,
        id_column='Image Name', remove_extension='.jpg',
        labels=labels,
        sep='\t', header=False, index=False,
        top_keywords_per_category=top_keywords_per_category
    )
    logging.info(f"One-hot encoded features saved to '{one_hot_encoded_file}'.")

    # ------------------------------#
    # Step 10: Log Summary of 'Other' Count Rows #
    # ------------------------------#
    logging.info("\nSummary of 'Other' count rows:")
    for count in sorted(other_count_rows.keys()):
        logging.info(f"Rows with {count} 'Other' keywords: {other_count_rows[count]}")


def clean_csv(input_file: str, labels: List[str]) -> pd.DataFrame:
    """
    Cleans the input CSV file and returns a DataFrame.

    Parameters:
        input_file (str): Path to the input CSV file.
        labels (List[str]): List of labels to process.

    Returns:
        pd.DataFrame: Cleaned DataFrame with 'Image Name' and 'Keywords' columns.
    """
    logging.info("Starting to clean the CSV file...")
    total_row_count = 0
    empty_row_count = 0
    na_count = 0
    no_comma_count = 0
    no_comma_rows = []
    cleaned_data = []
    # Read the entire content of the file
    try:
        with open(input_file, 'r', newline='', encoding='utf-8') as infile:
            content = infile.read()
    except Exception as e:
        logging.error(f"Error reading input file: {e}")
        raise

    # Remove newline characters inside quoted strings
    content = re.sub(r'("\[.*?")\n', r'\1 ', content)
    # Use StringIO to read the cleaned content as a CSV
    import io
    temp_file = io.StringIO(content)
    reader = csv.reader(temp_file)

    # Skip the first line (header)
    try:
        header = next(reader)
    except StopIteration:
        logging.error("Input file is empty.")
        raise ValueError("Input file is empty.")

    # Process each row
    for row in reader:
        total_row_count += 1
        if not row:
            continue
        jpg_file = row[0]
        if len(row) < 2:
            logging.warning(f"Row {total_row_count} doesn't have enough columns: {row}")
            keywords = []
        else:
            keywords = extract_keywords(row[1], labels)
        # Count occurrences of "N/A" and empty spaces
        na_count += keywords.count("N/A")
        na_count += sum(1 for k in keywords if k.strip() == "")
        # Check if the row is empty or contains only "N/A" or empty spaces
        if not keywords or all(k.strip().upper() == 'N/A' or k.strip() == '' for k in keywords):
            empty_row_count += 1
            continue  # Skip this row
        # Concatenate the keywords into a single string
        keywords_str = ', '.join(keywords)
        # Replace multiple spaces with a comma and a single space
        keywords_str = re.sub(r'\s{2,}', ', ', keywords_str)
        # Check if the row has no commas
        if ',' not in keywords_str:
            no_comma_count += 1
            no_comma_rows.append(f'{jpg_file},{keywords_str}')
        # Add to cleaned data
        cleaned_data.append({'Image Name': jpg_file, 'Keywords': keywords_str})
    # Now, create a DataFrame from the cleaned_data
    df = pd.DataFrame(cleaned_data)
    # Clean the 'Keywords' column further
    df = clean_keywords_in_df(df, labels)
    # Log statistics
    logging.info(f"Total rows processed: {total_row_count}")
    logging.info(f"Number of empty or N/A rows skipped: {empty_row_count}")
    logging.info(f"Occurrences of N/A (including empty spaces): {na_count}")
    logging.info(f"Number of rows without commas: {no_comma_count}")
    if no_comma_rows:
        logging.info("Rows without commas:")
        for row in no_comma_rows:
            logging.info(row)
    logging.info("CSV cleaning completed successfully.")
    return df


def extract_keywords(line: str, labels: List[str]) -> List[str]:
    """
    Extracts keywords from a line.

    Parameters:
        line (str): The line from which to extract keywords.
        labels (List[str]): List of labels to be removed.

    Returns:
        List[str]: List of extracted keywords.
    """
    # First, try to extract keywords between curly braces
    keyword_groups = re.findall(r'\{([^}]+)\}', line)
    extracted_keywords = []
    for group in keyword_groups:
        # Split by ',' to separate multiple keywords
        parts = group.split(',')
        if parts:
            # Take the first part
            first_part = parts[0].strip()
            # Further split by '/' if present
            if '/' in first_part:
                first_part = first_part.split('/')[0].strip()
            # Append the processed keyword
            extracted_keywords.append(first_part)
    return extracted_keywords


def clean_keywords_in_df(df: pd.DataFrame, labels: List[str]) -> pd.DataFrame:
    """
    Cleans the 'Keywords' column in the DataFrame.

    Parameters:
        df (pd.DataFrame): DataFrame containing the keywords to be cleaned.
        labels (List[str]): List of labels to be removed.

    Returns:
        pd.DataFrame: DataFrame with cleaned 'Keywords' column.
    """
    # Construct a regex pattern from the labels with word boundaries
    labels_pattern = r'\b(' + '|'.join(map(re.escape, labels)) + r')\b[: -]*'
    # Replace labels with an empty string using regex
    df['Keywords'] = df['Keywords'].str.replace(labels_pattern, '', regex=True)
    # Remove any '- ' (dash and space)
    df['Keywords'] = df['Keywords'].str.replace('- ', '', regex=False)
    # Remove single quotes around keywords
    df['Keywords'] = df['Keywords'].str.replace(r"'([^']+)'", r'\1', regex=True)
    # Replace multiple spaces with a comma and a single space
    df['Keywords'] = df['Keywords'].str.replace(r'\s{2,}', ', ', regex=True)
    # Remove leading/trailing commas and spaces
    df['Keywords'] = df['Keywords'].str.strip(' ,')
    return df


def lemmatize_keywords(df: pd.DataFrame, lemmatized_counts_output_file: str, lemmatizer: WordNetLemmatizer) -> Tuple[pd.DataFrame, Counter]:
    """
    Lemmatizes the keywords in the DataFrame and returns a DataFrame with lemmatized keywords.
    Optionally saves the lemmatized keyword counts to a CSV file.

    Parameters:
        df (pd.DataFrame): DataFrame containing 'Image Name' and 'Keywords' columns.
        lemmatized_counts_output_file (str): Path to save the lemmatized keyword counts CSV.
        lemmatizer (WordNetLemmatizer): NLTK WordNetLemmatizer instance.

    Returns:
        Tuple[pd.DataFrame, Counter]: DataFrame with lemmatized keywords and keyword counts.
    """
    logging.info("Starting lemmatization of keywords...")
    per_row_data = []
    keyword_counts = Counter()
    # Iterate through each row to extract and lemmatize keywords
    for index, row in track(df.iterrows(), total=len(df), description="Lemmatizing keywords..."):
        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 = [lemmatizer.lemmatize(kw.lower()) 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)
    if lemmatized_counts_output_file:
        # 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
        try:
            lemmatized_counts_df.to_csv(lemmatized_counts_output_file, index=False)
            logging.info(f"Lemmatized keyword counts saved to '{lemmatized_counts_output_file}'.")
        except Exception as e:
            logging.error(f"Error saving lemmatized keyword counts: {e}")
            raise
    logging.info("Lemmatization completed.")
    return per_row_df, keyword_counts


def load_model(model_name: str, device: str = 'cuda') -> SentenceTransformer:
    """
    Loads a SentenceTransformer model.
    """
    try:
        logging.info(f"Loading SentenceTransformer model '{model_name}' on device '{device}'. This may take a few minutes...")
        model = SentenceTransformer(model_name, device=device)
        logging.info(f"Model '{model_name}' loaded successfully on device '{device}'.")
        return model
    except Exception as e:
        logging.error(f"Error loading model '{model_name}': {e}")
        raise


def perform_clustering_with_threshold(keyword_counts: Counter, model: SentenceTransformer, threshold: int = 0,
                                      num_clusters: int = None, distance_threshold: float = 0.2) -> Dict[str, str]:
    """
    Clusters keywords using hierarchical clustering.

    Parameters:
        keyword_counts (Counter): Counts of lemmatized keywords.
        model (SentenceTransformer): SentenceTransformer model.
        threshold (int): Frequency threshold to separate frequent and infrequent keywords.
                         Set to 0 to cluster all keywords.
        num_clusters (int): Number of clusters to form if distance_threshold is not specified. Default is None.
        distance_threshold (float): The linkage distance threshold above which clusters will not be merged.

    Returns:
        Dict[str, str]: A mapping from keywords to their cluster representatives.
    """
    logging.info(f"Clustering keywords with frequency threshold: {threshold}")
    # Apply frequency threshold
    if threshold > 0:
        filtered_keywords = [kw for kw, count in keyword_counts.items() if count >= threshold]
    else:
        filtered_keywords = list(keyword_counts.keys())
    logging.info(f"Total number of keywords to cluster: {len(filtered_keywords)}")

    # Cluster all keywords and map them to their representatives
    mapping = cluster_keywords(
        filtered_keywords, keyword_counts, model,
        num_clusters=num_clusters, distance_threshold=distance_threshold
    )
    logging.info(f"Number of cluster representatives: {len(set(mapping.values()))}")
    return mapping


def cluster_keywords(keywords: List[str], keyword_counts: Counter, model: SentenceTransformer,
                    num_clusters: int = None, distance_threshold: float = 0.2) -> Dict[str, str]:
    """
    Clusters keywords and maps them to cluster representatives.

    Parameters:
        keywords (List[str]): List of keywords to cluster.
        keyword_counts (Counter): Counts of lemmatized keywords.
        model (SentenceTransformer): SentenceTransformer model.
        num_clusters (int): Number of clusters to form. Default is None.
        distance_threshold (float): The linkage distance threshold above which clusters will not be merged.

    Returns:
        Dict[str, str]: A mapping from each keyword to its cluster representative.
    """
    logging.info("Clustering keywords...")
    if not keywords:
        logging.warning("No keywords provided for clustering.")
        return {}
    
    # Compute embeddings for all keywords
    embeddings = model.encode(keywords, show_progress_bar=True)
    if len(embeddings) == 0:
        logging.warning("No valid embeddings found for clustering.")
        return {}
    # Compute cosine distance matrix
    from sklearn.metrics.pairwise import cosine_distances
    distance_matrix = cosine_distances(embeddings)
    # Perform Hierarchical Clustering using precomputed distances
    if distance_threshold is not None and distance_threshold > 0:
        logging.info(f"Using distance threshold: {distance_threshold}")
        clustering_model = AgglomerativeClustering(
            n_clusters=None,
            metric='precomputed',
            linkage='average',
            distance_threshold=distance_threshold
        )
    else:
        logging.info(f"Using number of clusters: {num_clusters}")
        clustering_model = AgglomerativeClustering(
            n_clusters=num_clusters,
            metric='precomputed',
            linkage='average'
        )
    labels = clustering_model.fit_predict(distance_matrix)
    logging.info(f"Clustering completed. Number of clusters formed: {len(set(labels))}")
    # Organize keywords by clusters
    clusters = defaultdict(list)
    for label, keyword in zip(labels, 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


def categorize_and_count_clustered_keywords(per_row_df: pd.DataFrame, labels: List[str], category_limit: int = 200) -> Dict[str, Counter]:
    """
    Categorizes clustered keywords and counts them per category.

    Parameters:
        per_row_df (pd.DataFrame): DataFrame containing 'Clustered Keywords'.
        labels (List[str]): List of category labels.
        category_limit (int): Limit of top N keywords per category.

    Returns:
        Dict[str, Counter]: Dictionary of keyword counts per category.
    """
    logging.info("Categorizing and counting clustered keywords...")
    category_keyword_counts = {label: Counter() for label in labels}
    for idx, row in per_row_df.iterrows():
        clustered_keywords = row['Clustered Keywords'].split(', ')
        for i, label in enumerate(labels):
            if i < len(clustered_keywords):
                keyword = clustered_keywords[i]
                category_keyword_counts[label][keyword] += 1
    # Limit to top N keywords per category
    for label in labels:
        original_count = len(category_keyword_counts[label])
        category_keyword_counts[label] = Counter(dict(category_keyword_counts[label].most_common(category_limit)))
        logging.info(f"Category '{label}': Reduced from {original_count} to {len(category_keyword_counts[label])} keywords (top {category_limit})")
    logging.info("Categorization and counting completed.")
    return category_keyword_counts


def generate_clustered_keywords_and_categories(per_row_df: pd.DataFrame, clustered_keyword_mapping: Dict[str, str],
                                               top_keywords_per_category: Dict[str, Set[str]], labels: List[str],
                                               remove_duplicates: bool = False, cluster_reps: Set[str] = set(),
                                               category_keyword_counts: Dict[str, Counter] = None,
                                               category_limit: int = 200) -> Tuple[pd.DataFrame, Dict[int, int]]:
    """
    Generates the final DataFrame with clustered keywords and category assignments.

    Parameters:
        per_row_df (pd.DataFrame): DataFrame containing 'Lemmatized Keywords'.
        clustered_keyword_mapping (Dict[str, str]): Mapping of keywords to their cluster representatives.
        top_keywords_per_category (Dict[str, Set[str]]): Top keywords per category.
        labels (List[str]): List of category labels.
        remove_duplicates (bool): Whether to remove duplicates in clustered keywords.
        cluster_reps (Set[str]): Set of cluster representative keywords.
        category_keyword_counts (Dict[str, Counter]): Keyword counts per category.
        category_limit (int): The maximum number of keywords allowed per category.

    Returns:
        Tuple[pd.DataFrame, Dict[int, int]]: Updated DataFrame with 'Clustered Keywords' and category columns,
                                            and a dictionary counting how many rows have a specific number of 'Other' keywords.
    """
    logging.info("Generating clustered keywords and assigning categories...")
    # Initialize lists for categories
    category_keywords_list = {label: [] for label in labels}
    clustered_keywords_list = []
    # Initialize a counter for the number of rows with specific "Other" counts
    max_other_count = len(labels)
    other_count_rows = {i: 0 for i in range(max_other_count + 1)}  # from 0 to len(labels)

    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 labels
        assigned_keywords = []
        row_other_count = 0
        for i, label in enumerate(labels):
            keyword = None
            if i < len(clustered_keywords):
                keyword = clustered_keywords[i]
            elif label == labels[-1]:  # If it's the last label
                remaining_keywords = clustered_keywords[i:]
                # Check if any remaining keywords are in top_keywords_per_category[label]
                found_keyword = False
                for kw in remaining_keywords:
                    if kw in top_keywords_per_category[label]:
                        keyword = kw
                        found_keyword = True
                        break
                if not found_keyword:
                    keyword = 'Other'
                    row_other_count += 1
            else:
                keyword = 'Other'
                row_other_count += 1

            # Check if keyword is 'Other' based on top_keywords_per_category or if it's a cluster representative
            if keyword != 'Other' and keyword not in top_keywords_per_category[label]:
                if keyword not in cluster_reps:
                    keyword = 'Other'
                    row_other_count += 1

            # Ensure that cluster representatives are not categorized as 'Other'
            if keyword in cluster_reps and keyword not in top_keywords_per_category[label]:
                # Replace the least frequent keyword if category_limit is reached
                if len(top_keywords_per_category[label]) >= category_limit:
                    # Identify the least frequent keyword in the category
                    least_freq_kw = min(top_keywords_per_category[label], key=lambda w: category_keyword_counts[label][w])
                    # Remove it
                    top_keywords_per_category[label].remove(least_freq_kw)
                # Add the cluster representative
                top_keywords_per_category[label].add(keyword)
                logging.info(f"Ensured cluster representative '{keyword}' is included in category '{label}'.")
                # No need to increment 'Other' count as we're adding a valid keyword

            assigned_keywords.append(keyword)
        # Add assigned keywords to the list
        for label, keyword in zip(labels, assigned_keywords):
            category_keywords_list[label].append(keyword)
        # Update the 'Other' count rows
        if row_other_count <= max_other_count:
            other_count_rows[row_other_count] += 1

    per_row_df['Clustered Keywords'] = clustered_keywords_list
    # Add the category columns to per_row_df
    for label in labels:
        per_row_df[label] = category_keywords_list[label]
    # Log summary of "Other" counts per category
    other_count = {label: 0 for label in labels}  # Counter for "Other" per category
    logging.info("\nSummary of 'Other' counts per category:")
    for label in labels:
        other_count[label] = sum(1 for keyword in category_keywords_list[label] if keyword == 'Other')
        logging.info(f"{label}: {other_count[label]}")
    logging.info("Clustered keywords and categories generated successfully.")
    return per_row_df, other_count_rows


def encode_categorical_features(input_df: pd.DataFrame, output_tsv_path: str, id_column: str = 'Image Name',
                                remove_extension: str = '.jpg', labels: List[str] = None, sep: str = '\t',
                                header: bool = False, index: bool = False,
                                top_keywords_per_category: Dict[str, Set[str]] = None):
    """
    Encodes categorical features using One-Hot Encoding and saves them to a TSV file.

    Parameters:
        input_df (pd.DataFrame): Input DataFrame containing the categorical features.
        output_tsv_path (str): Path to save the one-hot encoded TSV file.
        id_column (str): Name of the identifier column.
        remove_extension (str): Extension to remove from the identifier column values.
        labels (List[str]): List of columns to encode.
        sep (str): Delimiter to use in the TSV file.
        header (bool): Whether to include the header in the TSV file.
        index (bool): Whether to write row names (index).
        top_keywords_per_category (Dict[str, Set[str]]): Top keywords per category to define the encoder categories.
    """
    logging.info("Starting one-hot encoding of categorical features...")
    if labels is None:
        raise ValueError("Labels for encoding must be provided.")
    if top_keywords_per_category is None:
        raise ValueError("Top keywords per category must be provided.")

    # Build categories per label
    categories_per_label = []
    for label in labels:
        categories = list(top_keywords_per_category[label])
        if 'Other' not in categories:
            categories.append('Other')
        categories_per_label.append(categories)

    # Read only the necessary columns
    columns_to_encode = labels
    columns_to_read = [id_column] + columns_to_encode
    df = input_df[columns_to_read]

    # Initialize OneHotEncoder with specified categories
    enc = OneHotEncoder(categories=categories_per_label, 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(input_features=columns_to_encode)
    # Create a DataFrame with the encoded features
    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]].reset_index(drop=True), emb_df.reset_index(drop=True)], 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 to TSV file
    try:
        result_df.to_csv(output_tsv_path, index=index, header=header, sep=sep)
        logging.info(f"One-hot encoded features saved successfully to '{output_tsv_path}'.")
    except Exception as e:
        logging.error(f"Error saving one-hot encoded features: {e}")
        raise


if __name__ == "__main__":
    main()