doczyai-pipelines/fieldExtraction/src/utils.py

import os
import re
import pandas as pd
import shutil
import time
import sys
import numpy as np
from collections import defaultdict
import config
from botocore.exceptions import ClientError
from io import StringIO


class InvalidDateException(Exception):
    pass


def read_local(file_path):
    # Check if the file is a text file
    if os.path.isfile(file_path) and file_path.endswith(".txt"):
        try:
            # First attempt to open the file with UTF-8 encoding
            with open(file_path, "r", encoding="utf-8") as file:
                file_contents = file.read()
                return file_contents
        except UnicodeDecodeError:
            # If UTF-8 fails, try reading the file with ANSI encoding
            try:
                with open(file_path, "r", encoding="cp1252") as file:
                    file_contents = file.read()
                    return file_contents
            except UnicodeDecodeError:
                # If ANSI also fails, log an error message or handle it accordingly
                print(f"Failed to decode {file_path} with UTF-8 and cp1252 encodings.")
    
    if os.path.isfile(file_path) and file_path.endswith(".csv"):
        try:
            df = pd.read_csv(file_path)
            return df
        except Exception:
            print('Failed to read .csv file.')

    if os.path.isfile(file_path) and file_path.endswith(".xlsx"):
        try:
            df = pd.read_excel(file_path)
            return df
        except Exception:
            print('Failed to read .xlsx file.')


def list_s3_files(prefix=config.S3_PREFIX, bucket=config.S3_BUCKET):
    """List all files in an S3 folder."""
    response = config.S3_CLIENT.list_objects_v2(Bucket=bucket, Prefix=prefix)
    return [item["Key"] for item in response.get("Contents", [])]


def read_s3():
    s3_client = config.S3_CLIENT
    bucket = config.S3_BUCKET
    prefix = config.S3_PREFIX

    file_list = []
    continuation_token = None

    while True:

        if continuation_token:
            response = s3_client.list_objects_v2(
                Bucket=bucket, Prefix=prefix, ContinuationToken=continuation_token
            )
        else:
            response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)

        file_list += [
            obj["Key"]
            for obj in response.get("Contents", [])
            if not obj["Key"].endswith("/")
        ]

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    contract_list = sorted(file_list)
    files = {}

    for contract in contract_list:
        print(f"Reading {contract}")
        try:
            data = s3_client.get_object(Bucket=bucket, Key=contract)
            if contract.endswith('.txt'):
                contents = data["Body"].read()
                context = contents.decode("utf-8")
                path, filename = os.path.split(contract)
                files[filename] = context
        except:
            pass

    return files

def read_s3_csv(csv_path):
    s3_client = config.S3_CLIENT
    s3_parts = csv_path[5:].split("/", 1)
    bucket = s3_parts[0]
    prefix = s3_parts[1] if len(s3_parts) > 1 else ""
    try:
        for obj in s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)["Contents"]:
            key = obj["Key"]
            if key.endswith(".csv"):
                response = s3_client.get_object(Bucket=bucket, Key=key)
                df = pd.read_csv(StringIO(response["Body"].read().decode("utf-8")))
                return df
    except ClientError as e:
        print(f"Error listing S3 objects: {e}")

def read_input(local_path=config.LOCAL_PATH):
    if config.READ_MODE == "local":
        files = {}
        for file in os.listdir(local_path):
            full_path = os.path.join(local_path, file)
            file_text = read_local(full_path)
            files[file] = file_text
        return files
    elif config.READ_MODE == "s3":
        return read_s3()
    
def read_input_csv():
    if config.DF_READ_MODE == "local":
        ac_df = read_local(config.AC_DF) if config.AC_DF else None
        b_df = read_local(config.B_DF) if config.B_DF else None
    elif config.DF_READ_MODE == "s3":
        ac_df =  read_s3_csv(config.AC_DF) if config.AC_DF else None
        b_df = read_s3_csv(config.B_DF) if config.B_DF else None

    return ac_df, b_df


def consolidate_individual(input_folder="results", output_folder="output"):
    dfs = []
    for filename in os.listdir(input_folder):
        if filename.endswith(".csv"):
            filepath = os.path.join(input_folder, filename)
            dfs.append(pd.read_csv(filepath))

    # Remove temp folder
    if input_folder == "temp" and os.path.exists(input_folder):
        shutil.rmtree(input_folder)

    consolidated_df = pd.concat(dfs, ignore_index=True)

    # Write output
    version = 1
    existing_files = [
        filename
        for filename in os.listdir(output_folder)
        if filename.startswith(f"consolidated_results_{config.TODAY}")
    ]
    if existing_files:
        versions = [
            int(file.split("_v")[1].split(".")[0])
            for file in existing_files
            if "_v" in file
        ]
        if versions:
            version = max(versions) + 1
    filename = f"consolidated_results_{config.TODAY}_v{version}.csv"
    consolidated_df.to_csv(os.path.join(output_folder, filename))


def preprocess_text_file(file_path):
    # Attempt to open the file with UTF-8 encoding first
    try:
        with open(file_path, "r", encoding="utf-8") as file:
            text = file.read()
    except UnicodeDecodeError:
        # If UTF-8 fails, try reading the file with ANSI (cp1252) encoding
        try:
            with open(file_path, "r", encoding="cp1252") as file:
                text = file.read()
        except UnicodeDecodeError:
            # If ANSI also fails, log an error message or handle it accordingly
            print(f"Failed to decode {file_path} with UTF-8 and cp1252 encodings.")
            return []

    # Split the text into pages based on a specific marker
    pages = re.split(r"Start of Page No\. = \d+", text)
    return pages


def format_td_check(td_dicts, dont_include_list):
    final_str = ""
    dict_count = 1
    for td_dict in td_dicts:
        final_str += str(dict_count) + ". "
        for k in td_dict.keys():
            if k not in dont_include_list:
                final_str += k + ": " + td_dict[k] + ", "
        final_str += "\n"
        dict_count += 1
    return final_str


def contains_reimbursement(text, page="1"):
    if isinstance(text, dict):
        return page.isdigit() and (
            "%" in text[page]
            or "$" in text[page]
            or "percent " in text[page]
            or "compensation schedule" in text[page].lower()
            or "reimbursement schedule" in text[page].lower()
        )
    elif isinstance(text, str):
        return (
            "%" in text
            or "$" in text
            or "percent " in text
            or "compensation schedule" in text.lower()
            or "reimbursement schedule" in text.lower()
        )
    else:
        print("contains_reimbursement - Invalid data type")


def split_and_write_csv(
    df, n, column_name, output_dir=config.CONSOLIDATED_OUTPUT_DIRECTORY
):
    """
    Splits the DataFrame into N sub-DataFrames based on unique values in a specified column
    and writes each sub-DataFrame to a separate CSV file. Each sub-DataFrame will contain
    roughly the same number of unique values from the specified column.

    Args:
    df (pd.DataFrame): The DataFrame to split.
    n (int): The number of sub-DataFrames to create.
    column_name (str): The column to split the DataFrame by.
    output_dir (str): The directory to save the CSV files to.
    """
    # Create the directory if it does not exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Get unique values and shuffle them to randomize distribution if needed
    unique_values = np.array(df[column_name].unique())
    np.random.shuffle(unique_values)

    # Split the unique values into N parts
    value_groups = np.array_split(unique_values, n)

    # Loop through each group and filter DataFrame by the values in each group
    for i, values in enumerate(value_groups, 1):
        # Filter df for the values in this group
        sub_df = df[df[column_name].isin(values)]
        # Create a filename for each sub-DataFrame
        filename = f"{config.BATCH_ID}-{i}.csv"
        path = os.path.join(output_dir, filename)

        # Write the sub-DataFrame to a CSV file
        sub_df.to_csv(path, index=False)
        print(f"Written {path}")

def contains_lol(text, page='1'):
    if isinstance(text, dict):
        return page.isdigit() and (
            re.search('les+e*o*r of+', text[page].lower())
        )
    elif isinstance(text, str):
        return (
            re.search('les+e*o*r of+', text.lower())
        )
    else:
        print("contains_lol - Invalid data type")
    


def filter_already_processed(input_dict, input_dict_b):
    already_processed_ac, already_processed_b = [], []
    for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
        # AC
        if config.AC_RESULTS_NAME in os.listdir(
            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
        ):
            already_processed_ac.append(folder_name + ".txt")
        # B
        if config.B_RESULTS_NAME in os.listdir(
            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
        ):
            already_processed_b.append(folder_name + ".txt")

    input_dict_ac = {
        key: input_dict[key]
        for key in input_dict.keys()
        if key not in already_processed_ac
    }
    input_dict_b = {
        key: input_dict_b[key]
        for key in input_dict_b.keys()
        if key not in already_processed_b
    }
    return input_dict_ac, input_dict_b


def is_empty(value):
    if pd.isna(value):
        return True
    else:
        empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan"]
        return value in empty_values


def find_regex_matches(
    pattern: str, text_dict: dict
) -> tuple[list, list]:  # TODO: Write few unit tests
    """Helper function to find matches and list of pages of a given regex pattern in text_dict.

    Args:
        pattern(str): A particular regex match that with a pattern.
        text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page

    Returns:
        matches(list): List of match(s)
        page_list(list): List of pages where the match was found

    """
    matches = []
    page_list = []
    for page, text in text_dict.items():
        found = re.findall(pattern, text)
        if found:
            matches.extend(found) # make it append
            page_list.append(page)

    return matches, page_list
    
    # matches = defaultdict(list)
    # for page, text in text_dict.items():
    #     matches_found = re.findall(pattern, text)
    #     if matches_found:
    #         matches[page].extend(matches_found)
    # return list(matches.values()), list(matches.keys())


def chunk_selected_pages(text_dict: dict[str, str], page_list: list) -> str:
    """Combines chunks for pages with regex matches for IRS, NPI, etc.

    Args:
        text_dict (dict[str, str]): A dictionary keyed by string page number and valued
            by string page contents
        page_list (str): List of pages to be chunked for context

    Returns:
        context(str): corresponding chunks of text (input strings delimited by newlines).
    """

    page_list = sorted([int(page_str) for page_str in set(page_list)])
    context = "\n".join([text_dict[str(page_num)] for page_num in page_list])
    return context


def add_hyphen_if_needed(input_str):
    """Checks for specific patterns in the input string and modifies it by adding 
    a '0' at the start and a hyphen at the 3rd index, or just adds a hyphen at 
    the 3rd index if it's a 9-digit number.

    If the input string is exactly 8 characters long, it adds a '0' at the start 
    and inserts a hyphen at the 3rd position. If the string matches the pattern 
    of a 9-digit number (exactly 9 digits), it inserts a hyphen at the 3rd index.
    If the input string matches the pattern 'x-xxxxxxx', it adds '0' at the start.

    All conditions will return a string in the format 'xx-xxxxxxx'.

    Args:
    - input_str (str): The input string (IRS).

    Returns:
    - str: The modified string in the format 'xx-xxxxxxx'.
    
    Example:
    - '1-2345678' → '01-2345678'
    - '12345678' → '01-2345678'
    - '123456789' → '12-3456789'
    """
    # Check if the string matches the pattern 'x-xxxxxxx' (1 character, hyphen, 7 characters)
    if re.fullmatch(r"\d-\d{7}", input_str):
        # Add '0' at the start
        input_str = '0' + input_str
    
    # Check if the string is exactly 'xxxxxxxx' (8 characters)
    elif bool(re.fullmatch(r"\b\d{8}\b", input_str)):
        # Add '0' at the start and insert a hyphen at the 3rd index
        input_str = '0' + input_str
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # If the string matches the pattern of a 9-digit number, add a hyphen at the 3rd index
    elif bool(re.fullmatch(r"\b\d{9}\b", input_str)):
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # Ensure the output is in the format 'xx-xxxxxxx'
    if not bool(re.fullmatch(r"\b\d{2}-\d{7}\b", input_str)):
        input_str = None

    return input_str

def convert_to_us_date_format(date_str: str) -> str:
    """
    Input date string is expected to be in the format YYYY-MM-DD
    
    Ensure the date is in MM/DD/YYYY format. If the month is greater than 12, swap the month and day.
        If the input is `N/A` or some variant (defined by is_empty), InvalidDateException will be raised.
    Args:
        date_str (str): The date string to be formatted. We expect it to come in as YYYY-MM-DD format (note the hyphens as well)

    Returns:
        str: The date string formatted as MM/DD/YYYY.

    Raises:
        InvalidDateException: If the date string is not in the expected format.
    """
    
    if not isinstance(date_str, str):
        raise TypeError(f"Date must be a string, got: {type(date_str)}")

    if is_empty(date_str) or re.fullmatch(r"\d{4}-\d{1,2}-\d{1,2}", date_str) is None:
        raise InvalidDateException(f"Invalid date format: {date_str}")

    # Split the date string into year, month, and day
    year, month, day = date_str.split("-")
    
    year_int = int(year)
    month_int = int(month)
    day_int = int(day)
    
    if year_int == 0 or month_int == 0 or day_int == 0:
        raise InvalidDateException(f"Invalid date format: {date_str}")
    
    if month_int > 12 and 0 < day_int <= 31:
        # Swap the month and day if the month is greater than 12
        month_int, day_int = day_int, month_int
    
    if 0 < month_int <= 12 and 0 < day_int <= 31:
        # Return the formatted date in MM/DD/YYYY format
        year, month, day = str(year_int).zfill(4), str(month_int).zfill(2), str(day_int).zfill(2)
        return "/".join([month, day, year])
    
    else:
        # Raise an exception if the month or day is out of range
        raise InvalidDateException(f"Invalid date format: {date_str}")


def remove_unnamed_columns(df):

    return df[[col for col in df.columns if 'Unnamed' not in col]]

def remove_txt_extension(filename):
    if isinstance(filename, str) and filename.lower().endswith(".txt"):
        return filename[:-4]
    return filename