doczyai-pipelines/fieldExtraction/src/utils/string_utils.py

import json
import re
import traceback
import warnings
from datetime import datetime

import numpy as np
import pandas as pd

from typing import overload, Literal

from constants.delimiters import Delimiter
from constants.regex_patterns import (
    BACKTICK_PATTERN,
    PIPE_PATTERN,
    TRIPLE_BACKTICK_PATTERN,
)


def extract_text_from_delimiters(
    raw_output: str, delimiter: Delimiter, match_index: int = -1
) -> str:
    """
    Extracts a match enclosed in specified delimiters from the raw output.
    This function searches for text enclosed by a specified delimiter in the given raw output string.
    It returns the match specified by the `which_match` index. If no match is found, it returns "N/A".

    - raw_output (str): The raw output string to search within.
    - which_match (int, optional): The index of the match to return. Defaults to -1, which returns the last match. If the index is out of range, the last match is returned.
    - str: The extracted text or "N/A" if no match is found.

    Raises:
        - TypeError: If `raw_output` is not a string or `delimiter` is not a Delimiter enum value.
        - ValueError: If an unsupported delimiter is provided.

    Returns:
        - str: The extracted answer or "N/A" if no match is found.
    """
    if not isinstance(raw_output, str):
        raise TypeError(
            "Expected a string for raw_output, got {0}.".format(
                type(raw_output).__name__
            )
        )

    if not isinstance(delimiter, Delimiter):
        raise TypeError(
            "Expected a Delimiter enum value, got {0}.".format(type(delimiter).__name__)
        )

    # Define a pattern based on the delimiter type
    if delimiter == Delimiter.PIPE:
        pattern = PIPE_PATTERN
    elif delimiter == Delimiter.BACKTICK:
        pattern = BACKTICK_PATTERN
    elif delimiter == Delimiter.TRIPLE_BACKTICK:
        pattern = TRIPLE_BACKTICK_PATTERN
    else:
        raise ValueError("Unsupported delimiter. Use one of the Delimiter enum values.")

    # Find all matches based on the pattern
    matches = re.findall(pattern=pattern, string=raw_output)

    if len(matches) == 0:
        return "N/A"

    # Adjust for negative indices
    if match_index < 0:
        match_index += len(matches)

    if not 0 <= match_index < len(matches):
        warnings.warn("Index out of range. Returning the last match by default.")
        match_index = -1

    return matches[match_index]


def page_key_sort(page_key: str) -> tuple[int, float | str]:
    """Sorts page keys, prioritizing numeric keys first.

    Args:
        page_key (str): The page key to sort.

    Returns:
        tuple[int, float|str]: A tuple where the first element is 0 for numeric keys and 1 for non-numeric keys, and the second element is the numeric value for numeric keys or the original string for non-numeric keys.

    Example:
        >>> page_key_sort("1")
        (0, 1.0)
        >>> page_key_sort("A")
        (1, "A")


    Example Usage with mixed types:
        >>> L = ["10", "2", "A", "1"]
        >>> sort(L, key=page_key_sort)
        >>> print(L)
        ['1', '2', '10', 'A']

    """
    try:
        return (0, float(page_key))  # Numeric pages first, in numerical order
    except ValueError:
        return (1, str(page_key))  # Non-numeric pages after, in alphabetical order


def json_parsing_search(response_text: str, field_list: list[str]) -> dict[str, str]:
    """
    Parses a JSON-like string to extract values for fields.

    This function attempts to clean and format the input string to resemble a valid JSON structure.
    It then searches for the specified fields within the string and extracts their corresponding values.

    Parameters:
        response_text (str): The JSON-like string to parse.
        field_list (list[str]): A list of field names to search for in the response_text.

    Returns:
        dict[str, str]: A dictionary where keys are field names and values are the extracted values from the response_text.
    """
    try:
        if response_text.split("{", 1)[1].strip()[0] == '"':
            response_text = "{" + response_text.split("{", 1)[1]
        else:
            response_text = "{" + response_text
    except:
        response_text = "{" + response_text

    if len(response_text.split("}", 1)) > 1:
        if response_text.rsplit("}", 1)[0].strip()[-1] == '"':
            response_text = response_text.rsplit("}", 1)[0] + "}"
        elif response_text.rsplit("}", 1)[0].strip()[-1] == "}":
            response_text = response_text.rsplit("}", 1)[0]
    else:
        response_text = response_text.rstrip(",") + "}"

    field_l = []
    answer_l = []
    position_dict = {}

    for f in field_list:
        location = response_text.find('"' + f + '"')
        if location != -1:
            position_dict[location] = f

    field_list = list(dict(sorted(position_dict.items())).values())

    for f in field_list:
        if f in response_text:
            field_l.append(f)
            value = response_text.split('"' + f + '"')[0]
            response_text = response_text.split('"' + f + '"')[1]
            if f != field_list[0] and f != field_list[-1]:
                answer_l.append(
                    value.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("\n")
                    .strip(" ")
                    .strip(",")
                    .strip('"')
                )
            elif f == field_list[-1]:
                answer_l.append(
                    value.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("\n")
                    .strip(" ")
                    .strip(",")
                    .strip('"')
                )
                answer_l.append(
                    response_text.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("}")
                    .strip("\n")
                    .strip(" ")
                    .strip('"')
                )

    return dict(zip(field_l, answer_l))


def universal_json_load(string_dict: str):
    """Matches json dicts and list of dicts - NOT simple lists
    Args:
        string_dict (str): The string containing the JSON-like content.
    Returns:
        list[dict] | dict: The dictionary or list of dictionaries obtained from parsing the cleaned and corrected string.
    Raises:
        ValueError: If no valid JSON object is found in the input string. In addition,
        the traceback as well as the offending string are printed.

    Note: if more than one json-like (or list-like) is found in the input string, this function will
    take the last one positionally.
    """
    # Try to match dictionary or list of dictionaries first
    dict_matches = re.findall(r"\{.*\}|\[\s*?\{.*\}\s*?\]", string_dict, re.DOTALL)
    if dict_matches:
        matched_str = dict_matches[-1]  # Take the last match in the string
        try:
            return json.loads(matched_str)
        except json.JSONDecodeError as e:
            print(f"Failed to parse JSON: {e}")
            print(f"Original string: {string_dict}")
            print(f"Traceback: {traceback.format_exc()}")
            # Use placeholders to preserve JSON structure characters for CSV compatibility
            csv_safe_string = (
                string_dict.replace(",", "[COMMA]")
                .replace("\n", "[NEWLINE]")
                .replace('"', "[QUOTE]")
            )
            raise ValueError(
                f"JSON parsing failed: {e.msg} at position {e.pos}; Original string:\n{csv_safe_string}"
            ) from e
    # If no dict pattern found, try to match list of strings
    list_matches = re.findall(r"\[(.*?)\]", string_dict, re.DOTALL)
    if list_matches:
        matched_str = f"[{list_matches[-1]}]"  # Take the last match in the string
        try:
            # Clean up the string list before parsing
            cleaned_str = re.sub(r'(?<!\\)"', '"', matched_str)  # Handle escaped quotes
            cleaned_str = re.sub(
                r"'", '"', cleaned_str
            )  # Replace single quotes with double quotes
            return json.loads(cleaned_str)
        except json.JSONDecodeError as e:
            print(f"Failed to parse list: {e}")
            print(f"Original string: {string_dict}")
            print(f"Traceback: {traceback.format_exc()}")
            # Use placeholders to preserve JSON structure characters for CSV compatibility
            csv_safe_string = (
                string_dict.replace(",", "[COMMA]")
                .replace("\n", "[NEWLINE]")
                .replace('"', "[QUOTE]")
            )
            raise ValueError(
                f"JSON parsing failed: {e.msg} at position {e.pos}; Original string:\n{csv_safe_string}"
            ) from e

    return {}


reimbursement_strings = [  # These are for `method='keyword'`
    "%",
    "$",
    "percent",
    "billed charges",
]
reimb_regex = r"(?<![$%])(?:\$\d+|\d+[$%])(?![$%])"


# TODO: check should this always be checking just page 1?  Bc in file_processing and table_utils it is but in preprocessing_funcs it's checking specific pages
def contains_reimbursement(text, page="1", method="keyword"):
    """
    Checks if the given text contains any reimbursement-related keywords or patterns.

    Args:
        text (str or dict): The text to check. If a dictionary is provided, it should have page numbers as keys.
        page (str, optional): The page number to check in the dictionary. Defaults to "1".
        method (str, optional): The method to use for checking. Options are "keyword"
         and "regex". Defaults to "keyword".

    Returns:
        bool: True if any reimbursement-related keyword or pattern is found, False otherwise.
    """
    if method == "keyword":
        if isinstance(text, dict):
            text_to_check = text.get(page, "").lower()
        elif isinstance(text, str):
            text_to_check = text.lower()
        else:
            print("contains_reimbursement - Invalid data type")
            return False
        return any(keyword in text_to_check for keyword in reimbursement_strings)
    elif method == "regex":
        return bool(re.search(reimb_regex, text))
    else:
        raise ValueError("Invalid method. Choose 'keyword' or 'regex'.")


def count_reimbursements_in_exhibit(exhibit_text: str) -> int:  # JUST by regex
    """Counts reimbursements in an exhibit text.  Reimbursements are detected by a regex
    search as defined by `reimb_regex`.

    Args:
        exhibit_text (str): Input exhibit text

    Returns:
        int: Number of reimbursements detected
    """
    return len(re.findall(reimb_regex, exhibit_text))


# Overloads for type checking
@overload
def is_empty(value: str, pd_mask: bool = True) -> bool: ...


@overload
def is_empty(value: list, pd_mask: bool = True) -> bool: ...


@overload
def is_empty(value: pd.Series, pd_mask: Literal[True] = True) -> pd.Series: ...


@overload
def is_empty(value: pd.Series, pd_mask: Literal[False]) -> bool: ...


def is_empty(value: str | list | pd.Series, pd_mask: bool = True) -> bool | pd.Series:
    """
    Checks if a value is considered empty or invalid.

    Args:
        value: The value to check.
        Can be a string, list, or pandas Series.
        pd_mask (bool): Only relevant for Series inputs.
            If True, returns a mask for empty values in a pandas Series.
            Otherwise, the function returns True iff the entire Series is empty.

    Returns:
        bool | pd.Series:
            - When the input is a string, returns True if the value is empty or invalid, False otherwise.
            - When the input is a list, returns True if the list is empty or all elements are empty/invalid.
            - When the input is a pandas Series:
                - If `pd_mask` is True, returns a boolean mask indicating which elements are empty or invalid.
                - If `pd_mask` is False, returns True iff all elements are empty/invalid, or if the entire Series is empty.
    """
    empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan", "None"]

    if isinstance(value, list):  # Handle list inputs
        return not value or all(is_empty(v) for v in value)

    # if it's a pd.Series, return the mask (True for empty values)
    if isinstance(value, pd.Series):
        if pd_mask:
            return (
                value.isna()
                | (value.isin(empty_values))
                | (value.astype(str).str.strip() == "")
                | (
                    value.astype(str)
                    .str.lower()
                    .str.strip()
                    .isin(["n/a", "na", "null", "none", "nan"])
                )
            )
        else:
            return value.isna().all()

    if pd.isna(value):
        return True

    if isinstance(value, str):
        if not value or value.isspace():
            return True
        lower_stripped = value.lower().strip()
        if lower_stripped in ["n/a", "na", "null", "none", "nan"]:
            return True

    return value in empty_values


def datetime_str():
    return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")


def extract_signature_page(text_dict: dict, filename: str) -> dict:
    """Extracts the signature page(s) from a contract text dictionary and returns it as a dictionary.
    First tries to match the textract marker "this page has N signature(s)", then falls back to looking
    for the word "signature" if no matches are found.  "this page has N signature(s)" is Textract's marker for signatures

    Args:
        text_dict (dict): A dictionary containing text data, organized by pages.
        filename (str): The name of the file being processed.

    Returns:
        dict: A dictionary containing the extracted signature page(s).
    """

    signature_pages = {}
    # This regex pattern matches the Textract marker for signatures
    # It will look for both numeric and a few textual representations of numbers (one-ten)
    textract_pattern = re.compile(
        r"this page has (?:\d+|one|two|three|four|five|six|seven|eight|nine|ten) signature",
        re.IGNORECASE,
    )

    # First pass: Try to match the full textract pattern
    for page_num, page_text in text_dict.items():
        if textract_pattern.search(page_text):
            signature_pages[page_num] = page_text

    # If no pages found with textract pattern, try with  keyword
    if not signature_pages:
        page_items = list(text_dict.items())
        for i, (page_num, page_text) in enumerate(page_items):
            lower_text = page_text.lower()
            if (
                "signature page follow" in lower_text
                or "signature authorization" in lower_text
            ):
                # curr_page = int(float(page_num))
                if page_num not in signature_pages:
                    signature_pages[page_num] = page_text  # current page
                if i + 1 < len(page_items):
                    next_page_num, next_page_text = page_items[i + 1]
                    # next_page = int(float(next_page_num))
                    if next_page_num not in signature_pages:
                        signature_pages[next_page_num] = next_page_text  # next page

    return signature_pages


def extract_effective_date_pages(text_dict: dict, filename: str) -> dict:
    """
    Extracts the effective date page(s) from a contract text dictionary and returns it as a dictionary.

    This function identifies pages that either contain the Textract marker
    "this page has N signature(s)" or the keyword 'effective date:'.
    Additionally, it ensures that the first page (page 1) is always included
    in the result, as it often contains effective date.

    Args:
        text_dict (dict): A dictionary containing text data, organized by pages.
        filename (str): The name of the file being processed.

    Returns:
        dict: A dictionary containing the extracted effective date page(s),
              where keys are page numbers and values are the corresponding page text.
    """

    effective_date_pages = {}

    # Regex pattern to match the Textract marker for signature pages which often has effective dates
    # This pattern accounts for both numeric and textual representations of numbers (one-ten).
    textract_pattern = re.compile(
        r"this page has (?:\d+|one|two|three|four|five|six|seven|eight|nine|ten) signature",
        re.IGNORECASE,
    )

    # Iterate through each page in the text dictionary.
    for page_num, page_text in text_dict.items():
        # page_num = int(float(page_num))
        # Check if the page contains the Textract marker or the keyword 'effective date:'.
        if textract_pattern.search(page_text) or "effective date:" in page_text.lower():
            effective_date_pages[page_num] = page_text
        # Always include the first page (page 1) as it often contains effective date.
        elif page_num == 1 or page_num == "1" or page_num == "1.0":
            effective_date_pages[page_num] = page_text

    return effective_date_pages


def normalize_state_to_abbreviation(state_value: str) -> str:
    """Convert state names or abbreviations to standardized two-letter abbreviations.

    Handles various formats of state names and abbreviations, including full names,
    common abbreviations, and variations. Returns the standardized two-letter abbreviation
    or the original value if no match is found.

    Args:
        state_value (str): State name or abbreviation from extraction

    Returns:
        str: Two-letter state abbreviation or original value if no match found
    """
    if not state_value or state_value in ["N/A", "UNKNOWN"]:
        return state_value

    # Dictionary mapping full state names to abbreviations
    state_mapping = {
        "alabama": "AL",
        "alaska": "AK",
        "arizona": "AZ",
        "arkansas": "AR",
        "california": "CA",
        "colorado": "CO",
        "connecticut": "CT",
        "delaware": "DE",
        "florida": "FL",
        "georgia": "GA",
        "hawaii": "HI",
        "idaho": "ID",
        "illinois": "IL",
        "indiana": "IN",
        "iowa": "IA",
        "kansas": "KS",
        "kentucky": "KY",
        "louisiana": "LA",
        "maine": "ME",
        "maryland": "MD",
        "massachusetts": "MA",
        "michigan": "MI",
        "minnesota": "MN",
        "mississippi": "MS",
        "missouri": "MO",
        "montana": "MT",
        "nebraska": "NE",
        "nevada": "NV",
        "new hampshire": "NH",
        "new jersey": "NJ",
        "new mexico": "NM",
        "new york": "NY",
        "north carolina": "NC",
        "north dakota": "ND",
        "ohio": "OH",
        "oklahoma": "OK",
        "oregon": "OR",
        "pennsylvania": "PA",
        "rhode island": "RI",
        "south carolina": "SC",
        "south dakota": "SD",
        "tennessee": "TN",
        "texas": "TX",
        "utah": "UT",
        "vermont": "VT",
        "virginia": "VA",
        "washington": "WA",
        "west virginia": "WV",
        "wisconsin": "WI",
        "wyoming": "WY",
        "district of columbia": "DC",
    }

    # Clean and normalize the input
    cleaned_state = state_value.strip().lower()

    # Check if it's already a valid 2-letter abbreviation
    if len(cleaned_state) == 2 and cleaned_state.upper() in state_mapping.values():
        return cleaned_state.upper()

    # Check if it matches a full state name
    if cleaned_state in state_mapping:
        return state_mapping[cleaned_state]

    # Handle common variations
    variations = {
        "calif": "CA",
        "cal": "CA",
        "fla": "FL",
        "florida": "FL",
        "ny": "NY",
        "n.y.": "NY",
        "tx": "TX",
        "tex": "TX",
        "penn": "PA",
        "pa": "PA",
        "mass": "MA",
        "massachusetts": "MA",
        "wash": "WA",
        "washington state": "WA",
        "d.c.": "DC",
        "dc": "DC",
        "washington dc": "DC",
    }

    if cleaned_state in variations:
        return variations[cleaned_state]

    # If no match found, return original value
    return state_value


def flatten_to_strings(items) -> list[str]:
    """Helper function to flatten and stringify any nested structure

    Args:
        items (Any): The input items; can be a list, tuple, or single value

    Returns:
        list[str]: A flattened list of stringified items
    """
    result = []
    if isinstance(items, (list, tuple)):
        for item in items:
            if isinstance(item, (list, tuple)):
                result.extend(flatten_to_strings(item))
            else:
                result.append(str(item).strip())
    else:
        result.append(str(items).strip())

    return result