doczyai-pipelines/src/pipelines/shared/extraction/tin_npi_funcs.py

import json

import logging
import re
from typing import Optional

import src.constants.regex_patterns as regex_patterns
import src.config as config
import src.prompts.prompt_templates as prompt_templates
from src.utils import llm_utils, string_utils
from src.pipelines.saas.prompts import prompt_calls
from src.constants.delimiters import Delimiter
from src.prompts.fieldset import Field, FieldSet
from difflib import SequenceMatcher


def get_all_matches(text: str, pattern: str) -> list[str]:
    """
    Extracts all unique matches of a given pattern from the provided text.

    Args:
        text (str): The input text to search for matches.
        pattern (str): The regular expression pattern to match against the text.

    Returns:
        list[str]: A list of unique matches found in the text. If the pattern is None or empty, returns an empty list.
    """
    if string_utils.is_empty(pattern):  # check if pattern is None or empty
        return []
    matches = re.findall(pattern, text)
    return list(set(matches))  # Remove duplicates


def get_all_matches_with_ocr(
    text: str, exact_pattern: str, ocr_pattern: str, identifier_type: str = "TIN"
) -> list[tuple]:
    """Extracts all matches for a specific identifier type from the text, using both exact and OCR patterns.

    Args:
        text (str): The input text to search for matches.
        exact_pattern (str): The regular expression pattern for exact matches.
        ocr_pattern (str): The regular expression pattern for OCR-correctable matches.
        identifier_type (str, optional): The type of identifier being searched for (e.g., "TIN" or "NPI"). Defaults to "TIN".

    Returns:
        list[tuple]: A list of (value, match_quality) tuples where match_quality is "EXACT" or "OCR_CORRECTED".
    """
    results = []
    seen_values = set()

    # Get exact matches and normalize them
    exact_matches = get_all_matches(text, exact_pattern)
    for match in exact_matches:
        # Normalize: remove hyphens/dots, keep only digits
        normalized = "".join(c for c in match if c.isdigit())
        # Validate length based on identifier type
        if identifier_type == "TIN":
            expected_length = 9
        elif identifier_type == "NPI":
            expected_length = 10

        if len(normalized) == expected_length and normalized not in seen_values:
            results.append((normalized, "EXACT"))
            seen_values.add(normalized)

    # Get OCR-correctable candidates (even if we found exact matches)
    ocr_candidates = get_all_matches(text, ocr_pattern)
    for candidate in ocr_candidates:
        logging.debug(f"OCR candidate for {identifier_type}: {candidate}")
        corrected, is_valid = attempt_ocr_correction(candidate, identifier_type)
        if is_valid and corrected not in seen_values:
            results.append((corrected, "OCR_CORRECTED"))
            seen_values.add(corrected)

    return results


def attempt_ocr_correction(candidate: str, identifier_type: str) -> tuple[str, bool]:
    """Attempt OCR correction on a candidate string

    Args:
        candidate (str): The candidate string that may contain OCR errors.
        identifier_type (str): The type of identifier (e.g., "TIN" or "NPI").

    Returns:
        tuple[str, bool]:   A tuple containing the corrected string and a boolean
                            indicating if the correction was successful:
                            (corrected_value, is_valid_after_correction)
    """
    if identifier_type not in ["TIN", "NPI"]:
        return (
            candidate,
            False,
        )  # Currently only TIN and NPI OCR correction is supported

    corrected = candidate
    corrections_made = []

    # Apply OCR substitutions
    for bad_char, good_char in regex_patterns.OCR_SUBSTITUTIONS.items():
        if bad_char in corrected:
            corrected = corrected.replace(bad_char, good_char)
            corrections_made.append(f"{bad_char}→{good_char}")

    # Remove formatting (hyphens, dots) and keep only digits
    digits_only = "".join(c for c in corrected if c.isdigit())

    # Validate based on identifier type
    if identifier_type == "TIN":
        # TIN must be exactly 9 digits
        expected_length = 9
    elif identifier_type == "NPI":
        # NPI must be exactly 10 digits
        expected_length = 10
    else:
        return candidate, False  # Unsupported identifier type

    if len(digits_only) == expected_length:
        if corrections_made:
            logging.debug(
                f"OCR correction applied: {candidate} -> {corrected} ({', '.join(corrections_made)})"
            )
        return digits_only, True

    # If not valid, return original candidate
    return candidate, False


def chunk_on_matches(matches: list[str], text_dict: dict) -> str:
    """
    Extracts and concatenates text from a dictionary of pages where any of the specified matches are found.

    Args:
        matches (list[str]): A list of strings to search for in the text.
        text_dict (dict): A dictionary where keys are identifiers (e.g., page numbers) and values are strings of text.

    Returns:
        str: A single string containing the concatenated text from all pages where at least one match is found,
        separated by newline characters.
    """
    valid_pages = []
    for page_text in text_dict.values():
        if any(match in page_text for match in matches):
            valid_pages.append(page_text)
    return "\n".join(valid_pages)


def clean_provider_info(provider_info: list[dict]) -> list[dict]:
    """Cleans provider information by standardizing Name/TIN/NPI formats.

    Args:
        provider_info (list[dict]): A list of provider information dictionaries.  Each
            dictionary should be keyed by TIN, NPI, NAME

    Returns:
        list: Cleaned provider information
    """
    cleaned_providers = []
    for provider in provider_info:
        cleaned_provider = {}

        # Clean TIN - remove hyphens and dots, ensure it's 9 digits
        if "TIN" in provider and provider["TIN"]:
            cleaned_provider["TIN"] = provider["TIN"].replace("-", "").replace(".", "")
            # Validate that TIN only contains digits or "|"
            if not all(c.isdigit() or c == "|" for c in cleaned_provider["TIN"]):
                cleaned_provider["TIN"] = "UNKNOWN"
        else:
            cleaned_provider["TIN"] = "UNKNOWN"

        # Clean NPI - remove hyphens and dots, ensure it's 10 digits
        if "NPI" in provider and provider["NPI"]:
            cleaned_provider["NPI"] = provider["NPI"].replace("-", "").replace(".", "")
            # Validate that NPI only contains digits or "|"
            if not all(c.isdigit() or c == "|" for c in cleaned_provider["NPI"]):
                cleaned_provider["NPI"] = "UNKNOWN"
        else:
            cleaned_provider["NPI"] = "UNKNOWN"

        # Clean Name - remove extra spaces and ensure it's not empty
        if "NAME" in provider and provider["NAME"]:
            cleaned_provider["NAME"] = (
                " ".join(provider["NAME"].split()).strip().rstrip(".")
            )
            # If name is empty after cleaning, set to "UNKNOWN"
            if not cleaned_provider["NAME"]:
                cleaned_provider["NAME"] = "UNKNOWN"
        else:
            cleaned_provider["NAME"] = "UNKNOWN"

        cleaned_providers.append(cleaned_provider)
    return cleaned_providers


def merge_provider_info_with_hybrid_smart_chunking(one_to_one_results, filename):
    """
    Merges provider_info into one_to_one_results.

    Args:
        one_to_one_results (dict): Dictionary of one-to-one field results.
        provider_info (list[dict]): List of dictionaries containing provider information.
            Each dictionary is keyed by TIN, NPI, NAME, IS_GROUP.

    Returns:
        dict: Updated one_to_one_results with merged values from provider_info.
    """

    # Extract group and non-group providers
    group_tins = set()
    group_npis = set()
    group_names = set()
    other_tins = []
    other_npis = []
    other_names = []

    provider_name = one_to_one_results.get("PROVIDER_NAME")
    provider_info = one_to_one_results["PROV_INFO_JSON"]
    provider_list = (
        json.loads(provider_info) if isinstance(provider_info, str) else provider_info
    )

    # Track if we found any name matches
    found_match = False

    for provider in provider_list:
        name_matches = prompt_calls.provider_name_match_check(
            provider_name, provider.get("NAME", "UNKNOWN"), filename
        )

        if name_matches:
            found_match = True
            provider["IS_GROUP"] = "Y"
            if provider.get("TIN"):
                group_tins.add(provider["TIN"])
            if provider.get("NPI"):
                group_npis.add(provider["NPI"])
            if provider.get("NAME"):
                group_names.add(provider["NAME"])
        else:
            provider["IS_GROUP"] = "N"
            if provider.get("TIN"):
                other_tins.append(provider["TIN"])
            if provider.get("NPI"):
                other_npis.append(provider["NPI"])
            if provider.get("NAME"):
                other_names.append(provider["NAME"])

    # If provider_name is not null and no match was found, add a new record
    if provider_name and not found_match:
        new_provider = {
            "NAME": provider_name,
            "TIN": "UNKNOWN",
            "NPI": "UNKNOWN",
            "IS_GROUP": "Y",
        }
        provider_list.append(new_provider)
        group_names.add(provider_name)

    # De-Unknown GROUP fields
    group_tins = [v for v in group_tins if v != "UNKNOWN"]
    group_npis = [v for v in group_npis if v != "UNKNOWN"]
    group_names = [v for v in group_names if v != "UNKNOWN"]

    # Deduplicate while preserving order and remove any "UNKNOWN" entries and overlaps between group and other
    group_tins = list(dict.fromkeys("|".join(group_tins).split("|")))
    other_tins = list(dict.fromkeys("|".join(other_tins).split("|")))
    other_tins = [
        tin for tin in other_tins if tin not in group_tins and tin != "UNKNOWN"
    ]  # remove any group TINs from other TINs

    group_npis = list(dict.fromkeys("|".join(group_npis).split("|")))
    other_npis = list(dict.fromkeys("|".join(other_npis).split("|")))
    other_npis = [
        npi for npi in other_npis if npi not in group_npis and npi != "UNKNOWN"
    ]  # remove any group NPIs from other NPIs

    group_names = list(dict.fromkeys("|".join(group_names).split("|")))
    other_names = list(dict.fromkeys("|".join(other_names).split("|")))
    other_names = [
        name for name in other_names if name not in group_names and name != "UNKNOWN"
    ]  # remove any group Names from other Names

    # Filter out records with NO_IDENTIFIERS_FOUND
    provider_list = [
        p for p in provider_list if p.get("NAME") != "NO_IDENTIFIERS_FOUND"
    ]

    # Deduplicate provider_list by NAME, keeping records with actual TIN/NPI values
    provider_list = deduplicate_providers_by_name(provider_list)

    # Reconstruct PROV_INFO_JSON and PROV_INFO_JSON_FORMATTED with IS_GROUP flags
    one_to_one_results["PROV_INFO_JSON"] = json.dumps(provider_list)
    one_to_one_results["PROV_INFO_JSON_FORMATTED"] = "\n".join(
        [json.dumps(provider) for provider in provider_list]
    )

    # Merge group provider information into one_to_one_results
    one_to_one_results["PROV_GROUP_TIN"] = "|".join(group_tins) if group_tins else ""
    one_to_one_results["PROV_GROUP_NPI"] = "|".join(group_npis) if group_npis else ""
    one_to_one_results["PROV_GROUP_NAME_FULL"] = (
        "|".join(group_names) if group_names else ""
    )
    # Merge other provider information into one_to_one_results
    one_to_one_results["PROV_OTHER_TIN"] = "|".join(other_tins) if other_tins else ""
    one_to_one_results["PROV_OTHER_NPI"] = "|".join(other_npis) if other_npis else ""
    one_to_one_results["PROV_OTHER_NAME_FULL"] = (
        "|".join(other_names) if other_names else ""
    )
    return one_to_one_results


def reimbursement_tin_npi(
    exhibit_dict: dict[str, str],
    exhibit_page: str,
    reimbursement_level_fields: FieldSet,
    filename: str,
) -> tuple[dict, FieldSet]:
    """
    Extracts Reimbursement-Level TIN and NPI values from the given exhibit chunk using an LLM.

    Args:
        exhibit_dict (dict[str, str]): Dictionary containing the exhibit text with page numbers as keys.
        exhibit_page (str): The page number containing reimbursement details.
        reimbursement_level_fields (FieldSet): The FieldSet object to store extracted field values.
        filename (str): The name of the file being processed.

    Returns:
        tuple[dict, FieldSet]: A dictionary containing extracted TIN and NPI values, and an
        updated FieldSet with additional fields if multiple values are found.
        Fields Added:   REIMB_PROV_NAME,
                        REIMB_PROV_TIN,
                        REIMB_PROV_NPI
    """

    def get_list(field_name, filtered_providers):
        # Get all values and filter out invalid ones
        all_values = [
            provider.get(field_name, "N/A") for provider in filtered_providers
        ]
        valid_values = [
            val
            for val in all_values
            if not string_utils.is_empty(val) and val != "UNKNOWN"
        ]
        return list(set(valid_values))  # Remove duplicates

    answer_dict = {}

    # Extract all providers from the exhibit chunk
    providers = get_provider_info(
        text_dict=exhibit_dict, page_num=exhibit_page, filename=filename
    )

    # Clean, standardize, and deduplicate extracted providers
    filtered_providers = deduplicate_providers(clean_provider_info(providers))

    field_configs = [
        ("REIMB_PROV_TIN", "TIN"),
        ("REIMB_PROV_NPI", "NPI"),
        ("REIMB_PROV_NAME", "NAME"),
    ]

    # Always add providers to reimbursement_level fields for proper term attribution
    for field_name, provider_key in field_configs:
        if len(filtered_providers) > 0:
            field = Field.load_from_file(config.FIELD_JSON_PATH, field_name)
            value_list = get_list(provider_key, filtered_providers)

            if len(value_list) == 0:
                answer_dict[field_name] = "N/A"
            else:
                field.update_valid_values(value_list)
                reimbursement_level_fields.add_field(field)
        else:
            # If no providers are found, set default values
            answer_dict[field_name] = "N/A"

    return answer_dict, reimbursement_level_fields


def deduplicate_providers(
    provider_info: list[dict],
) -> list[dict]:  # TODO: add unit test
    """Deduplicates provider information based on TIN, NPI, and NAME fields.
    This function checks for duplicates in the provider information list by comparing
    the TIN, NPI, and NAME fields. If all three fields match, only one instance of the
    provider is kept. If all three fields are unknown, the provider is skipped.

    Args:
        provider_info (list[dict]): A list of provider information dictionaries to be deduplicated.

    Returns:
        list[dict]: A list of deduplicated provider information dictionaries.
    """

    seen = set()
    valid_providers = []

    # Deduplicate if all three fields match
    for provider in provider_info:
        # Skip providers where all critical identification fields are unknown
        if (
            (
                string_utils.is_empty(provider.get("TIN", ""))
                or provider.get("TIN", "") == "UNKNOWN"
            )
            and (
                string_utils.is_empty(provider.get("NPI", ""))
                or provider.get("NPI", "") == "UNKNOWN"
            )
            and (
                string_utils.is_empty(provider.get("NAME", ""))
                or provider.get("NAME", "") == "UNKNOWN"
            )
        ):
            continue
        key = (
            provider.get("TIN", ""),
            provider.get("NPI", ""),
            provider.get("NAME", ""),
        )
        if key not in seen:
            seen.add(key)
            valid_providers.append(provider)
    return valid_providers


def get_provider_info(
    text_dict: dict, page_num: str, filename: str, payer_name: str
) -> list:
    """
    Extracts provider information from a specified page in a document using a language model.

    Args:
        text_dict (dict): Dictionary of text content with page numbers as keys.
        page_num (str): The page number to extract provider information from.
        filename (str): The name of the file being processed, used for logging or tracking purposes.
        first_page_provider_group_name (str, optional): Expected group name from page 1 for name verification.

    Returns:
        list: A list of extracted provider information dictionaries, each containing TIN, NPI,
                NAME, and flags for presence on signature or intro pages.

    Raises:
        ValueError: If the response from the language model cannot be parsed as valid JSON.
    """
    chunk = text_dict[page_num]

    provider_fields = FieldSet(
        file_path=config.FIELD_JSON_PATH, field_type="provider_info"
    )
    prompt = prompt_templates.TIN_NPI_TEMPLATE(
        chunk, provider_fields.print_prompt_dict(), payer_name
    )
    claude_answer_raw = llm_utils.invoke_claude(
        prompt,
        "sonnet_latest",
        filename,
        max_tokens=8192,
        cache=True,
        instruction=prompt_templates.TIN_NPI_TEMPLATE_INSTRUCTION(),
    )  # Sometimes rosters can have 50+ provider

    logging.debug(
        f"Raw LLM response for provider info on page {page_num}: {claude_answer_raw}"
    )

    # Defensive programming for any json parsing errors
    try:
        providers = string_utils.universal_json_load(claude_answer_raw)
        logging.debug(f"Parsed provider info on page {page_num}: {providers}")
        # Ensure we have a list of providers (not a dict or other type)
        if isinstance(providers, dict):
            # Single provider returned as dict - wrap in list
            providers = [providers]
        elif not isinstance(providers, list):
            # If not a list or dict, create a default list with error info
            logging.warning(
                f"Warning: Unexpected format in provider info response: {type(providers)}"
            )
            providers = [{"TIN": "UNKNOWN", "NPI": "UNKNOWN", "NAME": "PARSING_ERROR"}]

        # VALIDATION LAYER - Cross-check with regex findings
        page_text = text_dict[page_num]

        # Get all TINs and NPIs found via regex on this page
        regex_tins = [
            tin
            for tin, _ in get_all_matches_with_ocr(
                page_text,
                regex_patterns.TIN_PATTERN,
                regex_patterns.TIN_OCR_PATTERN,
                identifier_type="TIN",
            )
        ]
        regex_npis = [
            npi
            for npi, _ in get_all_matches_with_ocr(
                page_text,
                regex_patterns.NPI_PATTERN,
                regex_patterns.NPI_OCR_PATTERN,
                identifier_type="NPI",
            )
        ]

        # Collect all TINs/NPIs that LLM found
        llm_found_tins = []
        llm_found_npis = []
        for provider in providers:
            if provider.get("TIN") and provider["TIN"] != "UNKNOWN":
                # Normalize for comparison (remove hyphens/dots)
                raw_tins = [tin.strip() for tin in provider["TIN"].split("|")]
                normalized_tins = [
                    tin.replace("-", "").replace(".", "") for tin in raw_tins
                ]
                llm_found_tins.extend(normalized_tins)
            if provider.get("NPI") and provider["NPI"] != "UNKNOWN":
                raw_npis = [npi.strip() for npi in provider["NPI"].split("|")]
                normalized_npis = [
                    npi.replace("-", "").replace(".", "") for npi in raw_npis
                ]
                llm_found_npis.extend(normalized_npis)

        # Note: regex_tins and regex_npis are already normalized in get_all_matches_with_ocr
        missed_tins = [tin for tin in regex_tins if tin not in llm_found_tins]
        missed_npis = [npi for npi in regex_npis if npi not in llm_found_npis]
        if missed_tins or missed_npis:
            # WARN ONLY if there are missed identifiers
            logging.warning(
                f"LLM missed identifiers on page {page_num} of {filename}. Missed TINs: {missed_tins}, Missed NPIs: {missed_npis}"
            )

        return providers

    except ValueError as e:  # Handle JSON parsing error
        logging.error(f"Error parsing JSON response from LLM: {str(e)}")
        logging.error(f"Raw response: {claude_answer_raw}")
        # Return a default value
        return [{"TIN": "UNKNOWN", "NPI": "UNKNOWN", "NAME": "PARSING_ERROR"}]


def run_provider_info_fields(
    contract_text: str, text_dict: dict, filename: str, payer_name: str
):
    """
    Extracts provider information from a contract text using regex-based identifier extraction,
    processes the relevant text chunks, and cleans the results.
    Args:
        contract_text (str): The full text of the contract to process.
        text_dict (dict): A dictionary containing text data, organized by pages
        filename (str): The name of the file being processed, used for logging or reference.
        payer_name (str): The name of the payer extracted from the full context answers.
            This is not being used right now, but it might be useful for filtering providers later.
    Returns:
        list[dict]: A cleaned and standardized list containing provider information extracted
                    from the contract text.
    """
    tin_matches = get_all_matches_with_ocr(
        contract_text,
        regex_patterns.TIN_PATTERN,
        regex_patterns.TIN_OCR_PATTERN,
        identifier_type="TIN",
    )

    npi_matches = get_all_matches_with_ocr(
        contract_text,
        regex_patterns.NPI_PATTERN,
        regex_patterns.NPI_OCR_PATTERN,
        identifier_type="NPI",
    )

    # Extract just the values from the (value, match_quality) tuples for downstream processing
    all_tins = [tin for tin, quality in tin_matches]
    all_npis = [npi for npi, quality in npi_matches]

    logging.debug(f"Extracted TINs: {all_tins}")
    logging.debug(f"Extracted NPIs: {all_npis}")

    # Log OCR corrections for monitoring
    exact_tins = [tin for tin, quality in tin_matches if quality == "EXACT"]
    ocr_corrected_tins = [
        tin for tin, quality in tin_matches if quality == "OCR_CORRECTED"
    ]

    exact_npis = [npi for npi, quality in npi_matches if quality == "EXACT"]
    ocr_corrected_npis = [
        npi for npi, quality in npi_matches if quality == "OCR_CORRECTED"
    ]

    logging.debug(
        f"TIN extraction summary for {filename}: {len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
    )
    logging.debug(
        f"NPI extraction summary for {filename}: {len(exact_npis)} exact, {len(ocr_corrected_npis)} OCR-corrected"
    )

    # If there are no identifiers, return early with default values
    if not all_tins and not all_npis:
        # if there are no TINs or NPIs, we need to extract signature pages
        relevant_pages = list(
            string_utils.extract_signature_page(text_dict, filename).keys()
        )
    else:
        # Identify relevant pages containing either NPIs or TINs (with or without hyphen formatting, e.g., 68-0640053)
        relevant_pages = [
            page_num
            for page_num, page_text in text_dict.items()
            if any(
                re.sub(r"[\s-]", "", tin) in re.sub(r"[\s-]", "", page_text)
                for tin in all_tins
            )
            or any(
                re.sub(r"[\s-]", "", npi) in re.sub(r"[\s-]", "", page_text)
                for npi in all_npis
            )
        ]

        # Stopgap for when the TIN is written in the header or footer of every page (or nearly every page)
        if (
            len(relevant_pages) >= len(text_dict) - 1
            and (len(all_tins) == 1 or len(all_tins) == 0)
            and (len(all_npis) == 1 or len(all_npis) == 0)
        ):
            relevant_pages = [relevant_pages[0]]

    if not relevant_pages:
        deduplicated_provider_info = [
            {
                "TIN": "UNKNOWN",
                "NPI": "UNKNOWN",
                "NAME": "NO_IDENTIFIERS_FOUND",
            }
        ]
    else:
        # Process each page separately and combine results
        combined_provider_info = []
        for page_num in relevant_pages:
            page_provider_info = get_provider_info(
                text_dict, page_num, filename, payer_name
            )
            combined_provider_info.extend(page_provider_info)

        # Clean and standardize results
        cleaned_provider_info = clean_provider_info(combined_provider_info)
        # Deduplicate providers based on TIN and NPI
        deduplicated_provider_info = deduplicate_providers(cleaned_provider_info)

    # Create columns
    one_to_one_results = {}
    one_to_one_results["PROV_INFO_JSON"] = json.dumps(deduplicated_provider_info)
    one_to_one_results["PROV_INFO_JSON_FORMATTED"] = "\n".join(
        [json.dumps(provider) for provider in deduplicated_provider_info]
    )
    one_to_one_results["TIN_EXTRACTION_QUALITY"] = (
        f"{len(exact_tins)} exact, {len(ocr_corrected_tins)} OCR-corrected"
    )
    one_to_one_results["NPI_EXTRACTION_QUALITY"] = (
        f"{len(exact_npis)} exact, {len(ocr_corrected_npis)} OCR-corrected"
    )

    # Add FILENAME_TIN
    filename_tin = get_all_matches(filename, regex_patterns.FILE_NAME_TIN_PATTERN)
    one_to_one_results["FILENAME_TIN"] = filename_tin[0] if filename_tin else "N/A"

    return one_to_one_results


def deduplicate_providers_by_name(provider_list: list[dict]) -> list[dict]:
    """
    Deduplicates a list of provider dictionaries by NAME, preferring records with actual TIN/NPI values.

    Args:
        provider_list (list[dict]): List of provider dictionaries with NAME, TIN, NPI keys.

    Returns:
        list[dict]: Deduplicated list with one entry per unique NAME, preferring non-UNKNOWN values.
    """
    seen_names = set()
    deduplicated_list = []
    # Sort so providers with actual TIN/NPI values come first
    for provider in sorted(
        provider_list,
        key=lambda p: (p.get("TIN") == "UNKNOWN", p.get("NPI") == "UNKNOWN"),
    ):
        name = provider.get("NAME")
        if name not in seen_names:
            seen_names.add(name)
            deduplicated_list.append(provider)
    return deduplicated_list