doczyai-pipelines/src/codes/code_funcs.py

import ast
import json
import logging
import concurrent.futures
import os
import re

import pandas as pd
import src.config as config
import src.prompts.prompt_templates as prompt_templates
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
from src.constants.constants import Constants
from src.constants.delimiters import Delimiter
from src.prompts.fieldset import FieldSet


def clean_service(service, constants: Constants) -> str:
    """
    Cleans the service string by removing unnecessary terms and formatting it for further processing.

    Args:
        service (str): The service name or identifier to be cleaned.
    Returns:
        str: The cleaned service string. If the service is empty or contains only stop words, it returns an empty string.
    """

    if string_utils.is_empty(service):
        return ""

    service = service.strip().upper()

    # Map common values
    for acronym, full_form in constants.SYNONYM_MAP.items():
        if string_utils.is_empty(acronym) or string_utils.is_empty(full_form):
            continue
        service = re.sub(rf"\b{re.escape(acronym)}\b", full_form, service)

    # Remove other values
    for term in constants.REMOVAL_LIST:
        if term is None:
            continue
        service = re.sub(rf"\b{re.escape(term)}\b", "", service).strip()
        service = re.sub(r"\s+", " ", service)
    if service is None:
        return ""
    service = service.replace(" - ", "-")

    if any([v in service for v in ["UNLISTED", "UNCATEGORIZED", "NON-LISTED"]]):
        return "UNLISTED"

    # Check if all remaining words are stop words
    remaining_words = service.split()
    if all(word in constants.STOP_WORD_LIST for word in remaining_words):
        return ""

    return service


def get_regex_answers(service, code_answer_dict):
    """
    **UNUSED** Extracts CPT and Revenue codes from the service string using regex patterns.
    This function may be used in the future to save tokens for simple code extractions. It can be used if there is only one code match.
    If there are multiple matches, using this function risks losing information if we do not have a regex string for EVERY explicit code we need.

    Args:
        service (str): The service name or identifier.
        code_answer_dict (dict): Dictionary to store the extracted codes.
    Returns:
        dict: A dictionary containing the extracted codes. If no codes are found, returns an empty dictionary.
    """

    # Define the regex pattern for CPT codes and ranges
    # Only accept single code or single code ranges - multiple codes may have complexities that need to be prompted for
    cpt_pattern = (
        r"\b(?:[A-Z]\d{4}|\d{5})(?:\s*(?:-|through)\s*(?:[A-Z]\d{4}|\d{5}))?\b"
    )
    cpt_matches = re.findall(cpt_pattern, service)
    if len(cpt_matches) == 1:
        code_answer_dict["PROCEDURE_CD"] = [cpt_matches[0].replace("through", "-")]

    if "rev" in service.lower():
        rev_pattern = r"\b(?:\d{3}|\d{2}X)(?:\s*(?:-|through)\s*(?:\d{3}|\d{2}X))?\b"
        rev_matches = re.findall(rev_pattern, service)
        if len(rev_matches) > 0:
            code_answer_dict["REVENUE_CD"] = rev_matches

    return code_answer_dict


def code_explicit(service, methodology, filename):
    """
    Processes a service string and invokes a language model to generate a response
    based on the provided primary questions. The response is then parsed into a
    dictionary format.
    Args:
        service (str): The service name or identifier. If empty, an empty dictionary
                       is returned.
        code_primary_questions (str): The primary questions or prompts to be used
                                       for generating the response.
        filename (str): The name of the file associated with the operation, used
                        for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the parsed response from the language model.
              Returns an empty dictionary if the service string is empty.
    """

    # Prompt - field definitions are now included in CODE_EXPLICIT_INSTRUCTION() for caching
    prompt, _parser = prompt_templates.CODE_EXPLICIT(service, methodology)

    llm_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        filename,
        cache=True,
        instruction=prompt_templates.CODE_EXPLICIT_INSTRUCTION(),
        usage_label="CODE_EXPLICIT",
    )
    code_answer_dict = _parser(llm_answer_raw)
    return code_answer_dict


def code_category(service, proc_category, hcpcs_level2_mapping, filename):
    """
    Determines if the service explicitly names a code category, such as "A-Codes" or "J-Codes", and returns the corresponding codes.
    Args:
        service (str): The service name or identifier.
        proc_category (list): List of procedure categories to check against.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the procedure codes and their descriptions. If no codes are found, returns an empty dictionary.
    If the service does not explicitly name a code category, it returns an empty dictionary.
    """

    code_category_list = [v.split(":")[1].strip() for v in proc_category if ":" in v]

    matching_values = []
    for category in code_category_list:
        matching_values += [
            value
            for key, value in hcpcs_level2_mapping.items()
            if key.startswith(category)
        ]

    prompt, _parser = prompt_templates.CODE_CATEGORY(service, matching_values)

    llm_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        filename,
        cache=True,
        instruction=prompt_templates.CODE_CATEGORY_INSTRUCTION(),
    )
    llm_answer_final = _parser(llm_answer_raw)  # Returns list

    code_answer_dict = {"PROCEDURE_CD": [], "PROCEDURE_CD_DESC": []}
    for answer in llm_answer_final:
        if "0000" in answer:
            code_answer_dict["PROCEDURE_CD"].append(answer)
            code_answer_dict["PROCEDURE_CD_DESC"].append("")
        elif answer in hcpcs_level2_mapping.values():
            code = list(hcpcs_level2_mapping.keys())[
                list(hcpcs_level2_mapping.values()).index(answer)
            ]
            code_answer_dict["PROCEDURE_CD"].append(str(code))
            code_answer_dict["PROCEDURE_CD_DESC"].append(str(answer))

    return code_answer_dict


def code_implicit_special(service, filename):
    """
    Allow certain Service classes to be handled as special cases, such as Drugs, Vaccines, Surgery, and PT/OT/ST.
    This mapping and prompt can be customized, as clients may have different requirements for these categories that are not captured by our Proc Code levels

    Args:
        service (str): The service name or identifier.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the implicit codes found for the service. If no codes are found, returns an empty dictionary.
    """
    prompt, _parser = prompt_templates.CODE_IMPLICIT_SPECIAL(service)

    llm_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        filename,
        cache=True,
        instruction=prompt_templates.CODE_IMPLICIT_SPECIAL_INSTRUCTION(),
    )
    llm_answer_final = _parser(llm_answer_raw)

    special_case_mapping = {
        "Drugs": ["J0000-J9999"],
        "Vaccines": [
            "J0000-J9999",
            "90471-90474",
            "90620-90621",
            "90633",
            "90647-90648",
            "90651",
            "90670",
            "90672",
            "90680-90681",
            "90686",
            "90696",
            "90698",
            "90700",
            "90707",
            "90710",
            "90713",
            "90714",
        ],
        "Surgery": ["10004-69990"],
        "PT/OT/ST": [
            "92507-92508",
            "92526",
            "97014",
            "97110",
            "97112",
            "97116",
            "97150",
            "97161-97168",
            "97530",
            "97535",
        ],
        "PT": ["PT Codes TBD"],
        "OT": ["OT Codes TBD"],
        "ST": ["ST Codes TBD"],
    }
    code_answer_dict = {}
    if llm_answer_final[0] in special_case_mapping:
        code_answer_dict["PROCEDURE_CD"] = special_case_mapping[llm_answer_final]
        code_answer_dict["PROCEDURE_CD_DESC"] = llm_answer_final

    return code_answer_dict


def get_embedding_levels(level_suffix, implicit_run_dict, constants: Constants):
    """
    Retrieves the embedding levels and their corresponding mappings for a given level suffix.
    Args:
        level_suffix (int): The level suffix to determine which mappings to use.
        implicit_run_dict (dict): Dictionary indicating which codes to run implicitly.
    Returns:
        tuple: A tuple containing:
            - embedding_levels (list): List of embedding levels to use.
            - cpt_mapping (dict): Mapping for CPT codes.
            - hcpcs_mapping (dict): Mapping for HCPCS codes.
            - rev_mapping (dict): Mapping for Revenue codes.
    """
    if level_suffix == 1:
        cpt_mapping = constants.CPT_LEVEL1_MAPPING
        hcpcs_mapping = constants.HCPCS_LEVEL1_MAPPING
        rev_mapping = constants.REV_LEVEL1_MAPPING
        embedding_levels = []
        if implicit_run_dict["PROCEDURE_CD"]:
            embedding_levels += [
                f"cpt_level{level_suffix}",
                f"hcpcs_level{level_suffix}",
            ]
        if implicit_run_dict["REVENUE_CD"]:
            embedding_levels += [f"rev_level{level_suffix}"]

    elif level_suffix == 2:
        cpt_mapping = constants.CPT_LEVEL2_MAPPING
        hcpcs_mapping = constants.HCPCS_LEVEL2_MAPPING
        rev_mapping = constants.REV_MAPPING

        embedding_levels = []
        if implicit_run_dict["PROCEDURE_CD"]:
            embedding_levels += [
                f"cpt_level{level_suffix}",
                f"hcpcs_level{level_suffix}",
            ]
        if implicit_run_dict["REVENUE_CD"]:
            embedding_levels += [f"rev"]

    return embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping


def get_match_list(service, embedding_levels, constants: Constants, top_k: int):
    """
    Retrieves the best matches for a given service from multiple embedding levels using FAISS index search.
    Args:
        service (str): The service name or identifier.
        embedding_levels (list): List of embedding levels to search in.
        top_k (int): The number of top matches to retrieve.
    Returns:
        tuple: A tuple containing:
            - match_list (list): List of matched service descriptions.
            - highest_similarity (float): The highest similarity score among the matches.
    """
    target_vec = constants.EMBEDDING_MODEL.encode(
        [service], normalize_embeddings=True
    ).astype("float32")
    match_list = []
    highest_similarity = 0
    for level in embedding_levels:
        embedding = constants.get_embedding(level)

        similarity_scores, match_indices = embedding.index.search(target_vec, top_k)
        if max(similarity_scores[0]) > highest_similarity:
            highest_similarity = max(similarity_scores[0])

        for idx, score in zip(match_indices[0], similarity_scores[0]):
            match_list.append(embedding.choices[idx])

    return match_list, highest_similarity


def code_implicit_rag(service, implicit_run_dict, filename, constants):
    """
    Processes a service string to find implicit codes using RAG (Retrieval-Augmented Generation) methodology.
    Args:
        service (str): The service name or identifier.
        implicit_run_dict (dict): Dictionary indicating which codes to run implicitly.
        filename (str): The name of the file associated with the operation, used for logging or tracking purposes.
    Returns:
        dict: A dictionary containing the implicit codes found for the service. If no codes are found, returns an empty dictionary.
    """

    # Get Embedding Embeddings and Mappings for Levels 1 and 2
    level_dicts = []
    for level_suffix in [1, 2]:
        # Get Mappings and Embeddings
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            get_embedding_levels(level_suffix, implicit_run_dict, constants)
        )
        if not embedding_levels:
            return {}

        # Get Best Matches
        match_list, highest_similarity = get_match_list(
            service, embedding_levels, constants, top_k=5
        )
        level_dicts.append(
            {
                "level_suffix": level_suffix,
                "match_list": match_list,
                "highest_similarity": highest_similarity,
            }
        )

    # Sort level_answers by highest_similarity in descending order
    level_dicts.sort(key=lambda x: x["highest_similarity"], reverse=True)

    # Run the Implicit prompt for each level, starting from the highest similarity
    for level_dict in level_dicts:
        embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = (
            get_embedding_levels(
                level_dict["level_suffix"], implicit_run_dict, constants
            )
        )

        prompt, _parser = prompt_templates.CODE_IMPLICIT(
            service, level_dict["match_list"]
        )

        # Run Prompt
        llm_answer_raw = llm_utils.invoke_claude(
            prompt,
            "sonnet_latest",
            filename,
            cache=True,
            instruction=prompt_templates.CODE_IMPLICIT_INSTRUCTION(),
        )

        try:
            llm_answer_final = _parser(llm_answer_raw)
        except Exception as e:
            return {"CODE_METHODOLOGY": e}

        if not llm_answer_final:
            continue

        # Populate answers, if any
        code_answer_dict = {}
        proc_codes, rev_codes = [], []
        proc_descs, rev_descs = [], []
        for description in llm_answer_final:
            if description == "INVALID_SERVICE":
                code_answer_dict["CODE_METHODOLOGY"] = "Generic - Prompt"
                continue

            if description in cpt_mapping.values():
                matching_codes = [
                    str(key) for key, val in cpt_mapping.items() if val == description
                ]
                if matching_codes:
                    # Split pipe-delimited codes into individual codes
                    for code in matching_codes:
                        proc_codes.extend(code)
                    proc_descs.append(description)
            if description in hcpcs_mapping.values():
                matching_codes = [
                    str(key) for key, val in hcpcs_mapping.items() if val == description
                ]
                if matching_codes:
                    # Split pipe-delimited codes into individual codes
                    for code in matching_codes:
                        proc_codes.extend(code)
                    proc_descs.append(description)
            if description in rev_mapping.values():
                matching_codes = [
                    str(key) for key, val in rev_mapping.items() if val == description
                ]
                if matching_codes:
                    # Split pipe-delimited codes into individual codes
                    for code in matching_codes:
                        rev_codes.extend(code)
                    rev_descs.append(description)

        # Combine answers
        if proc_codes:
            code_answer_dict["PROCEDURE_CD"] = proc_codes
            code_answer_dict["PROCEDURE_CD_DESC"] = proc_descs
            code_answer_dict["CODE_METHODOLOGY"] = (
                f"Implicit - Level {level_dict['level_suffix']}"
            )
        if rev_codes:
            code_answer_dict["REVENUE_CD"] = rev_codes
            code_answer_dict["REVENUE_CD_DESC"] = rev_descs
            code_answer_dict["CODE_METHODOLOGY"] = (
                f"Implicit - Level {level_dict['level_suffix']}"
            )

        if code_answer_dict:
            return code_answer_dict

    return {}


def code_last_check(service, filename):
    """
    For Services that did not match any codes, this function categorizes the reason as "Specific" or "Generic".

    Args:
        service (str): The service term.
        filename (str): The name of the file associated with the operation.
    Returns:
        str: "Specific" if the service is a specific case that requires manual intervention,
             "Generic" if the service is a generic case that does not match any codes.
    """
    prompt, _parser = prompt_templates.CODE_LAST_CHECK(service)

    llm_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        filename,
        cache=True,
        instruction=prompt_templates.CODE_LAST_CHECK_INSTRUCTION(),
    )
    try:
        llm_answer_final = _parser(llm_answer_raw)
        return llm_answer_final[0]
    except:
        return "Generic"


def fill_bill_type(
    service,
    answer_dict,
    BILL_TYPE_MAPPING,
    BILL_TYPE_REVERSE_MAPPING,
    reimb_term=None,
    exhibit_text=None,
):
    """
    Fills the BILL_TYPE_CD and BILL_TYPE_CD_DESC fields in the answer dictionary.

    Args:
        service (str): The service term.
        answer_dict (dict): The answer dictionary to update.
        BILL_TYPE_MAPPING (dict): Mapping from codes to descriptions.
        BILL_TYPE_REVERSE_MAPPING (dict): Mapping from descriptions to codes.
        reimb_term (str, optional): The reimbursement term for additional context.
        exhibit_text (str, optional): Full exhibit text for contextual analysis.

    Returns:
        dict: The updated answer dictionary.
    """

    valid_bill_type = sorted(list(set(BILL_TYPE_MAPPING.values())))

    prompt, _parser = prompt_templates.FILL_BILL_TYPE(
        service, valid_bill_type, reimb_term=reimb_term, exhibit_text=None
    )

    # First attempt: Service term + reimbursement term
    llm_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        "",
        cache=True,
        instruction=prompt_templates.FILL_BILL_TYPE_INSTRUCTION(),
    )
    llm_answer_final = _parser(llm_answer_raw)

    # Second attempt: Service term + reimbursement term + exhibit context (if first attempt failed and context available)
    if not llm_answer_final and exhibit_text:
        prompt, _parser = prompt_templates.FILL_BILL_TYPE(
            service,
            valid_bill_type,
            reimb_term=reimb_term,
            exhibit_text=exhibit_text,
        )

        llm_answer_raw = llm_utils.invoke_claude(
            prompt,
            "sonnet_latest",
            "",
            cache=True,
            instruction=prompt_templates.FILL_BILL_TYPE_INSTRUCTION(),
        )
        llm_answer_final = _parser(llm_answer_raw)

    # Keep legacy behavior: return answer_dict unchanged if no result found
    if not llm_answer_final:
        return answer_dict

    bill_codes, bill_descs = [], []
    for description in llm_answer_final:
        if description in BILL_TYPE_REVERSE_MAPPING:
            bill_codes.append(BILL_TYPE_REVERSE_MAPPING[description])
            bill_descs.append(description)

    if bill_codes:
        answer_dict["BILL_TYPE_CD"] = bill_codes
        answer_dict["BILL_TYPE_CD_DESC"] = bill_descs

    return answer_dict


def fill_grouper_cd_desc(answer_dict, constants: Constants):
    """
    Fills the GROUPER_CD_DESC field in the answer dictionary based on the GROUPER_CD.

    Args:
        answer_dict (dict): The answer dictionary to update.

    Returns:
        dict: The updated answer dictionary.
    """

    grouper_cd = answer_dict.get("GROUPER_CD")
    if string_utils.is_empty(grouper_cd):
        return answer_dict
    grouper_type = answer_dict.get("GROUPER_TYPE")
    if string_utils.is_empty(grouper_type):
        return answer_dict

    grouper_descs = []
    grouper_cd = eval(grouper_cd) if "[" in grouper_cd else grouper_cd

    for code in grouper_cd:
        try:
            code_without_severity = int(
                code.split("-")[0]
            )  # Remove severity level if present
        except Exception as e:
            logging.warning(
                f"Failed to parse grouper code '{code}' as integer: {type(e).__name__}: {str(e)}. "
                f"Using original code value."
            )
            code_without_severity = code

        if grouper_type == "APR-DRG" and (
            code_without_severity in constants.GROUPER_APR_DRG_MAPPING
        ):
            grouper_descs.append(
                constants.GROUPER_APR_DRG_MAPPING[code_without_severity]
            )
        elif grouper_type == "MS-DRG" and (
            code_without_severity in constants.GROUPER_MS_DRG_MAPPING
        ):
            grouper_descs.append(
                constants.GROUPER_MS_DRG_MAPPING[code_without_severity]
            )

    if grouper_descs:
        answer_dict["GROUPER_CD_DESC"] = list(set(grouper_descs))

    return answer_dict


def get_implicit_runs(answer_dict):
    """
    Determines which code runs should be executed implicitly based on the Claim Type and Bill Type

    Args:
        answer_dict (dict): Dictionary containing, at minimum, the AARETE_DERIVED_CLAIM_TYPE_CD and BILL_TYPE_CD_DESC fields.
    Returns:
        dict: Dictionary containing boolean values for REVENUE_CD and PROCEDURE_CD indicating whether to run those codes implicitly.
    """

    run_dict = {}

    claim_type = answer_dict.get("AARETE_DERIVED_CLAIM_TYPE_CD")
    bill_type = answer_dict.get("BILL_TYPE_CD_DESC")

    if claim_type == "H":
        run_dict["REVENUE_CD"] = True
    elif claim_type == "M":
        if not string_utils.is_empty(bill_type):
            run_dict["REVENUE_CD"] = True
        else:
            run_dict["REVENUE_CD"] = False
    else:
        run_dict["REVENUE_CD"] = False

    # Run proc code implicit for all bill types except 2
    if bill_type in ["Inpatient Hospital", "Skilled Nursing Facility"]:
        run_dict["PROCEDURE_CD"] = False
    else:
        run_dict["PROCEDURE_CD"] = True

    return run_dict


def normalize_answer_dict_codes(answer_dict):
    """
    Normalizes code fields in an answer dictionary to JSON list format.

    Args:
        answer_dict (dict): Dictionary containing code fields to normalize.
    Returns:
        dict: Dictionary with normalized code fields.
    """
    fields_to_normalize = ["REVENUE_CD", "PROCEDURE_CD", "CPT4_PROC_CD"]

    for field in fields_to_normalize:
        if field in answer_dict:
            answer_dict[field] = string_utils.normalize_to_json_list(answer_dict[field])

    return answer_dict


def extract_codes_from_service(answer_dict, constants: Constants):
    """
    Extracts codes from the service term in the answer dictionary.

    Args:
        answer_dict (dict): Dictionary containing, at minimum, the SERVICE_TERM field.
    Returns:
        dict: Dictionary containing extracted codes and their descriptions.

    Codes Extracted:
        Explicit Codes:
            - PROCEDURE_CD: CPT, HCPCS
            - CPT4_PROC_MOD: CPT4 Modifiers
            - REVENUE_CD: Revenue codes
            - DIAG_CD: Diagnosis codes (ICD-10)
            - GROUPER_CD: MS and APR DRG Grouper Codes
            - NDC_CD: National Drug Codes (NDC)
            - CLAIM_ADMIT_TYPE_CD: Claim Admit Type Code
            - CLAIM_STATUS_CD: Claim Status Code
        Implicit Codes (Up to Level 2):
            - PROCEDURE_CD: CPT, HCPCS
            - REVENUE_CD: Revenue codes
    """

    service, bill_type = answer_dict.get("SERVICE_TERM"), answer_dict.get(
        "BILL_TYPE_CD_DESC"
    )
    filename = answer_dict.get("FILENAME", "")
    methodology = answer_dict.get("REIMB_TERM", "")

    if string_utils.is_empty(service):
        return {}

    # Fill Bill Type if not filled
    if string_utils.is_empty(bill_type):
        exhibit_text = answer_dict.get("EXHIBIT_TEXT", None)
        reimb_term = answer_dict.get("REIMB_TERM", None)
        answer_dict = fill_bill_type(
            service,
            answer_dict,
            constants.BILL_TYPE_MAPPING,
            constants.BILL_TYPE_REVERSE_MAPPING,
            reimb_term=reimb_term,
            exhibit_text=exhibit_text,
        )

    # Get list of codes we want implicit - based on claim type
    implicit_run_dict = get_implicit_runs(answer_dict)

    # Preprocessing
    service_clean = clean_service(service, constants)

    # Exit point for unlisted
    if service_clean == "UNLISTED":
        answer_dict["PROCEDURE_CD_DESC"] = "UNLISTED"
        answer_dict["CODE_METHODOLOGY"] = "UNLISTED"
        return normalize_answer_dict_codes(answer_dict)

    # Exit point if N/A or in DO_NOT_RUN list
    if string_utils.is_empty(service_clean) or any(
        [v in service_clean for v in constants.DO_NOT_RUN]
    ):
        answer_dict["CODE_METHODOLOGY"] = "Generic - Before Prompts"
        return normalize_answer_dict_codes(answer_dict)

    # Explicit Codes (run always)
    code_answer_dict = code_explicit(service_clean, methodology, filename)
    if any(
        not string_utils.is_empty(value) for value in code_answer_dict.values()
    ) and all(
        ["Category" not in v for v in code_answer_dict.get("PROCEDURE_CD", "")]
    ):  # if any code value is not empty, return
        if "NOT_ESTABLISHED" in code_answer_dict.get("PROCEDURE_CD", ""):
            code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Not Established"
            code_answer_dict["PROCEDURE_CD"] = "[]"
        else:
            code_answer_dict["CODE_METHODOLOGY"] = "Explicit"
        answer_dict.update(code_answer_dict)
        return normalize_answer_dict_codes(answer_dict)

    # Implicit Codes: Code Categories
    if any(["Category:" in v for v in code_answer_dict.get("PROCEDURE_CD", "")]):
        code_answer_dict = code_category(
            service_clean,
            code_answer_dict.get("PROCEDURE_CD", ""),
            constants.HCPCS_LEVEL2_MAPPING,
            "",
        )
        if any(
            not string_utils.is_empty(value) for value in code_answer_dict.values()
        ):  # if any code value is not empty, return
            code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Letter Category"
            answer_dict.update(code_answer_dict)
            return normalize_answer_dict_codes(answer_dict)

    # Implicit Codes: Special Categories
    code_answer_dict = code_implicit_special(service_clean, "")
    if code_answer_dict:
        code_answer_dict["CODE_METHODOLOGY"] = "Implicit - Special Case"
        answer_dict.update(code_answer_dict)
        return normalize_answer_dict_codes(answer_dict)

    # Implicit Codes: Levels 1 and 2
    code_answer_dict = code_implicit_rag(
        service_clean, implicit_run_dict, "", constants
    )
    if any(
        not string_utils.is_empty(value) for value in code_answer_dict.values()
    ):  # if any code value is not empty, return
        answer_dict.update(code_answer_dict)
        return normalize_answer_dict_codes(answer_dict)

    # No Match - Why? Generic or Specific
    last_check_answer = code_last_check(service_clean, "")
    if last_check_answer == "Generic":
        answer_dict["CODE_METHODOLOGY"] = "Generic - No Match"
    elif last_check_answer == "Specific":
        answer_dict["CODE_METHODOLOGY"] = "Specific - No Match"

    return normalize_answer_dict_codes(answer_dict)


def code_breakout(merged_results: pd.DataFrame, constants: Constants):
    """
    Processes a list of answer dictionaries to extract and fill in code-related information.
    Wrapper function to be run from main Doczy

    Args:
        merged_results (pd.DataFrame): DataFrame containing merged results from the processing pipeline
    Returns:
        pd.DataFrame: DataFrame with updated code fields.
    """
    import concurrent.futures

    answer_dicts = merged_results.to_dict(orient="records")

    # Parallelize code extraction
    def process_single_code(answer_dict):
        code_answer_dict = extract_codes_from_service(answer_dict, constants)
        answer_dict.update(code_answer_dict)
        return answer_dict

    max_workers = min(len(answer_dicts), 10)
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        answer_dicts_with_code = list(executor.map(process_single_code, answer_dicts))

    # Fill in GROUPER_CD_DESC based on GROUPER_TYPE and GROUPER_CD
    answer_dicts_with_code = [
        fill_grouper_cd_desc(answer_dict, constants)
        for answer_dict in answer_dicts_with_code
    ]

    df = pd.DataFrame(answer_dicts_with_code)

    # Normalize code columns to JSON list format
    for col in [
        "PROCEDURE_CD",
        "CPT4_PROC_MOD",
        "REVENUE_CD",
        "DIAG_CD",
        "GROUPER_CD",
        "NDC_CD",
        "CLAIM_ADMIT_TYPE_CD",
        "CLAIM_STATUS_CD",
    ]:
        if col in df.columns:
            df[col] = df[col].apply(string_utils.normalize_to_json_list)

    return df


def grouper_breakout(results_with_code: pd.DataFrame):
    """
    Processes a DataFrame to extract and fill in grouper-related information in cases where we got a grouper_cd from code_breakout
    but didn't get grouper fields from methodology + grouper breakout as in special case grouper was not identified and extracted.
    Args:
        results_with_code (pd.DataFrame): DataFrame containing the answer dictionaries with code fields
    Returns:
        pd.DataFrame: DataFrame with updated grouper fields.
    """
    # try grouper fields again if they are empty and grouper_cd is not empty
    answer_dicts_with_code = results_with_code.to_dict(orient="records")

    # Separate dicts that need grouper breakout from those that don't
    needs_breakout = []
    no_breakout_needed = []

    for answer_dict in answer_dicts_with_code:
        if string_utils.is_empty(
            answer_dict.get("GROUPER_TYPE")
        ) and not string_utils.is_empty(answer_dict.get("GROUPER_CD")):
            needs_breakout.append(answer_dict)
        else:
            no_breakout_needed.append(answer_dict)

    logging.debug(
        f"Grouper breakout: {len(needs_breakout)} rows need LLM processing, {len(no_breakout_needed)} rows skip"
    )

    # Parallelize LLM calls for rows that need breakout
    def process_grouper_breakout(answer_dict):
        grouper_breakout_prompt, grouper_parser = prompt_templates.GROUPER_BREAKOUT(
            answer_dict.get("SERVICE_TERM", ""),
            answer_dict.get("REIMB_TERM", ""),
        )
        claude_answer_raw = llm_utils.invoke_claude(
            grouper_breakout_prompt,
            "sonnet_latest",
            "",
            cache=True,
            instruction=prompt_templates.GROUPER_BREAKOUT_INSTRUCTION(),
        )
        try:
            grouper_answer = grouper_parser(claude_answer_raw)
        except Exception as e:
            grouper_answer = {"GROUPER_TYPE": e}

        answer_dict.update(grouper_answer)
        return answer_dict

    if needs_breakout:
        max_workers = min(len(needs_breakout), 10)
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            processed_breakout = list(
                executor.map(process_grouper_breakout, needs_breakout)
            )
    else:
        processed_breakout = []

    # Combine all results
    final_answer_dicts = no_breakout_needed + processed_breakout

    return pd.DataFrame(final_answer_dicts)