app/categorization/categorizer_list.py

import os
from datetime import datetime

from fastapi import HTTPException

import pandas as pd
import codecs

from app.categorization.config import CATEGORY_REFERENCE_OUTPUT_FILE
from app.categorization.categorizer import llm_list_categorizer, fuzzy_match_list_categorizer


async def categorize_list(df: pd.DataFrame) -> pd.DataFrame:
    """Asynchronously categorize a list of transactions.

    This function categorizes a list of transactions using a combination of fuzzy matching
    and a language model. It looks up new transaction descriptions in the reference file
    (a combination of user input and previous executions) to minimize API calls.
    Any uncategorized transactions are sent to the language model, and new description-category
    pairs are added to the reference file.

    Args:
        df (pd.DataFrame): The list of transactions to categorize.

    Returns:
        pd.DataFrame: The original DataFrame with an additional column for the category.
    """
    result = {"df": pd.DataFrame(), "description_category": "", "description_category_pairs": "", "uncategorized_descriptions": "", "categorized_descriptions": "", "error": ""}

    try:
        if not os.path.exists(CATEGORY_REFERENCE_OUTPUT_FILE):
            os.makedirs(CATEGORY_REFERENCE_OUTPUT_FILE)

        if os.path.exists(CATEGORY_REFERENCE_OUTPUT_FILE):
            # Read description-category pairs from the reference file

            # TODO FIRST ISSUE: Must save categorize_list categorized_descriptions items to output file
            description_category_pairs = pd.read_csv(
                CATEGORY_REFERENCE_OUTPUT_FILE, names=['name_description', 'category'], header=0
            )

            print(f"\ncategorize_list description_category_pairs \n{description_category_pairs}\n")
            df['name_description'] = df['name_description'].astype(str)
            df['category'] = df['category'].astype(str)
            df['type'] = df['type'].astype(str)
            # df['amount'] = df['amount'].str.strip().replace("\s", "").astype(float)
            # df['transaction_date'] = pd.to_datetime(df['transaction_date'].str.strip().replace("\s", ""), format='%d/%m/%Y')

            # Extract only descriptions for faster matching
            # description_category_pairs.columns = description_category_pairs.columns.str.strip()
            # descriptions = description_category_pairs['name_description'].values
            descriptions = df['name_description'].values
            print(f"\ncategorize_list descriptions \n{descriptions}\n")

            result['description_category'] = descriptions
            result['df'] = df
            temp_df = df

            print(f"\ncategorize_list df \n{df}\n")

            # Use fuzzy matching to find similar descriptions and assign the category
            # TODO SECOND ISSUE: Why isn't category being updated
            df['category'] = df['name_description'].apply(
                fuzzy_match_list_categorizer,
                args=(descriptions, description_category_pairs)
            )

            print(f"\ncategorize_list setting fuzzy match list: \n{df}\n")

            # Filter out uncategorized transactions, deduplicate, and sort by description
            uncategorized_descriptions = (
                df[df['category'].isnull()]
                .drop_duplicates(subset=['name_description'])
                .sort_values(by=['name_description'])
            )
            result['uncategorized_descriptions'] = uncategorized_descriptions
            print(f"\ncategorize_list uncategorized_descriptions: \n{uncategorized_descriptions}\n")

            # Ask the language model to categorize the remaining descriptions
            if not uncategorized_descriptions.empty:
                categorized_descriptions = await llm_list_categorizer(
                    uncategorized_descriptions[['name_description', 'category']]
                )

                categorized_descriptions.dropna(inplace=True)
                result['categorized_descriptions'] = categorized_descriptions
                print(f"\ncategorize_list categorized_descriptions: \n{categorized_descriptions}\n")

                # Update the category for uncategorized transactions based on the language model results
                if not categorized_descriptions.empty:
                    # temp_df['category'] = temp_df['name_description'].map(categorized_descriptions.set_index('name_description')['category'])
                    print(f"\ncategorize_list current dataframe:\n {temp_df}\n")

                    # df['category'] = df['category'].fillna(
                    #     df['name_description'].map(
                    #         categorized_descriptions.set_index('name_description')['category']
                    #     )
                    # )
                    # print(f"\ncategorize_list categorized_description df[category]: \n{df['category']}\n")
                
                # Fill remaining NaN values in 'category' with 'Other'
                with pd.option_context("future.no_silent_downcasting", True):
                    df['category'] = df['category'].fillna('Other').infer_objects(copy=False)
                    # print(f"\ncategorize_list Fill remaining NaN df[category]: \n{df['category']}\n")

        return df
    
    except Exception as e:
        # Return an error indicator and exception info
        print(f"ERROR categorizer_list: {str(e)}")
        result["error"] = str(e)
        raise HTTPException(status_code = 500, detail=f"categorize_list result: {result}")