doczyai-pipelines/fieldExtraction/src/investment/file_processing.py

import csv
import json
import logging
import os

import pandas as pd
import src.investment.aarete_derived as aarete_derived
import src.investment.dynamic_funcs as dynamic_funcs
import src.investment.investment_postprocessing_funcs as investment_postprocessing_funcs
import src.investment.one_to_n_funcs as one_to_n_funcs
import src.investment.one_to_one_funcs as one_to_one_funcs
import src.investment.postprocess as postprocess
import src.investment.preprocess as preprocess
import src.investment.smart_chunking_funcs as smart_chunking_funcs
import src.utils.io_utils as io_utils
import src.utils.string_utils as string_utils
from src import config
from src.config import FIELD_JSON_PATH
from src.investment.one_to_n_funcs import (combine_one_to_n_answers,
                                           reimbursement_level)
from src.investment.tin_npi_funcs import (merge_provider_info,
                                          reimbursement_tin_npi)
from src.prompts.investment_prompts import Field, FieldSet
from src.utils.string_utils import datetime_str


################## MERGE ##################
def merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results):
    """
    Merges one_to_one_results into one_to_n_results.

    Args:
        one_to_n_results (pd.DataFrame): DataFrame with N rows and M columns, including FILE_NAME.
        one_to_one_results (dict): Dictionary with several keys, including FILE_NAME.

    Returns:
        pd.DataFrame: Updated one_to_n_results with merged values from one_to_one_results.
    """
    # Iterate over all key-value pairs in one_to_one_results
    for k, v in one_to_one_results.items():
        if k not in one_to_n_results.columns:
            # If k is not a column in one_to_n_results, create it and populate all rows with v
            one_to_n_results[k] = v
        else:
            # If k is a column in one_to_n_results, replace empty values with v
            one_to_n_results[k] = one_to_n_results[k].apply(
                lambda row_val: v if string_utils.is_empty(row_val) else row_val
            )
    
    # Handle global lesser of injection
    global_lesser_of = one_to_one_results.get("GLOBAL_LESSER_OF_STATEMENT")
    if global_lesser_of and global_lesser_of not in ["N/A", "UNKNOWN"]:
        logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] Injecting global lesser of logic: {global_lesser_of}")
        one_to_n_results = inject_global_lesser_of_rows(one_to_n_results, global_lesser_of, one_to_one_results.get("FILE_NAME", "UNKNOWN"))
    else:
        logging.debug(f"[{one_to_one_results.get('FILE_NAME', 'UNKNOWN')}] No global lesser of logic found or applicable.")
        # Initialize the tracking flag when no global lesser of logic is found
        one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N"
    return one_to_n_results

def inject_global_lesser_of_rows(one_to_n_results: pd.DataFrame, global_lesser_of_stmt: str, filename: str) -> pd.DataFrame:
    """Create new rows with global lesser of constraint for REIMB_IDs that lack lesser of logic.
    
    For each row where LESSER_OF_IND="N", creates a duplicate row with the global lesser
    of statement applied via methodology breakout.  The original row is preserved with
    LESSER_OF_IND updated to "Y" (since it's now subject to lesser of) and GLOBAL_LESSER_OF_APPLIED="N".
    The new row has LESSER_OF_IND="Y" and GLOBAL_LESSER_OF_APPLIED="Y".  Both rows 
    share the same REIMB_ID.

    Args:
        one_to_n_results (pd.DataFrame): DataFrame containing one-to-n results with a 'LESSER_OF_IND' column.
        global_lesser_of_stmt (str): Global lesser of statement extracted from contract (e.g., "lesser of billed or allowable")
        filename (str): Name of the file being processed, used for logging and debugging.

    Returns:
        pd.DataFrame: Updated one_to_n_results with new rows injected for global lesser of logic.
    """
    # Find rows that need global lesser of injection
    rows_needing_injection = one_to_n_results[one_to_n_results['LESSER_OF_IND'] == "N"].copy()

    if rows_needing_injection.empty:
        logging.debug(f"No rows needing global lesser of injection for {filename}.")
        # Still need to set tracking flag for all existing rows
        one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N"
        return one_to_n_results
    
    # Update original rows that lacked lesser of: they're now subject to lesser of due to
    # global constraint
    one_to_n_results.loc[one_to_n_results['LESSER_OF_IND'] == "N", 'LESSER_OF_IND'] = "Y"
    # Set tracking flag for all existing rows (both original injected rows and rows that already had lesser of logic) to "N"
    # New rows will get "Y" for this flag in create_global_lesser_of_row
    one_to_n_results['GLOBAL_LESSER_OF_APPLIED'] = "N"

    # Run methodology breakout on global statement (once)
    global_methodology_breakout = one_to_one_funcs.run_global_lesser_of_breakout(global_lesser_of_stmt, filename)

    # Create new rows for each REIMB_ID that was lacking lesser of logic
    new_rows = []
    for _, original_row in rows_needing_injection.iterrows():
        new_row = one_to_one_funcs.create_global_lesser_of_row(original_row, global_methodology_breakout)
        new_rows.append(new_row)
    
    # Append new rows to results
    if new_rows:
        new_rows_df = pd.DataFrame(new_rows)
        one_to_n_results = pd.concat([one_to_n_results, new_rows_df], ignore_index=True)
        logging.debug(f"Injected {len(new_rows)} global lesser of rows for {filename}.")
    else:
        logging.debug(f"No new rows created for global lesser of injection for {filename}.")
    
    return one_to_n_results

def process_file(file_object, all_dataset, run_timestamp):
    filename, contract_text = file_object
    print(f"{datetime_str()} Processing {filename}...")

    ################## PREPROCESS ##################
    contract_text = preprocess.clean_text(contract_text)
    text_dict, top_sheet_dict = preprocess.split_text(contract_text)
    text_dict = preprocess.clean_tables(text_dict, filename)
    exhibit_dict, all_exhibit_headers = preprocess.one_to_n_exhibit_chunking(text_dict, filename)
    print(f"{datetime_str()} Preprocessing Complete - {filename}")

    ################## ONE TO N ##################
    if string_utils.contains_reimbursement(contract_text):
        one_to_n_results, dynamic_one_to_one_fields = run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, all_dataset) # Return df
        one_to_n_results['FILE_NAME'] = filename
        one_to_n_results = investment_postprocessing_funcs.generate_reimb_ids(one_to_n_results) # Add reimb_id
        print(f"{datetime_str()} One to N Complete - {filename}")
    else:
        one_to_n_results = pd.DataFrame([{'FILE_NAME' : filename}])
        dynamic_one_to_one_fields = FieldSet()
        print(f"{datetime_str()} No One to N Found, Skipping - {filename}")
    

    ################## ONE TO ONE ##################
    one_to_one_results = run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict, dynamic_one_to_one_fields) # Return dict
    one_to_one_results['FILE_NAME'] = filename
    print(f"{datetime_str()} One to One Complete - {filename}")

    ################## MERGE ##################
    final_results = merge_one_to_one_into_one_to_n(one_to_n_results, one_to_one_results)

    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results)
    print(f"{datetime_str()} Postprocessing Complete - {filename}")

    ################## WRITE INDIVIDUAL ##################
    if config.WRITE_TO_S3:
        io_utils.write_s3(final_df, filename, run_timestamp, "individual")
    else:
        io_utils.write_local(final_df, filename, "", "individual")

    print(f"{datetime_str()} Writing Complete - {filename}")
    return final_df


def run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict, dynamic_one_to_one_fields):

    ################## INITIALIZE FIELDS ##################
    one_to_one_fields = FieldSet(relationship="one_to_one", file_path=config.FIELD_JSON_PATH).combine(dynamic_one_to_one_fields)

    ################## RUN REGEX ##################
    regex_answers_dict = one_to_one_funcs.run_provider_info_fields(contract_text, text_dict, filename)
    one_to_one_results = {}
    one_to_one_results['PROV_INFO_JSON'] = json.dumps(regex_answers_dict)
    one_to_one_results = merge_provider_info(one_to_one_results, regex_answers_dict)

    # What if the TIN/NPI regexes don't find anything?  Add PROV_GROUP_NAME_FULL to full context
    if one_to_one_results.get("PROV_OTHER_NAME_FULL") == "NO_IDENTIFIERS_FOUND": # This is how we know the regex didn't find anything
        other_provider_name_field = Field.from_values(
            field_name="PROV_GROUP_NAME_FULL",
            relationship="one_to_one",
            field_type="full_context",
            prompt="What is the name of the group provider associated with this contract? Return the full name as it appears in the document."
        )
        one_to_one_fields.add_field(other_provider_name_field)
        one_to_one_results["PROV_OTHER_NAME_FULL"] = "UNKNOWN"
    
    ################## RUN SMART CHUNKED PROMPTS ##################
    smart_chunked_answers_dict = one_to_one_funcs.run_smart_chunked_fields(
        one_to_one_fields, contract_text, filename,text_dict
    )

    ################## RUN FULL CONTEXT PROMPTS ##################
    full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
        one_to_one_fields, contract_text, filename
    )

    ################## COMBINE ANSWERS ################
    one_to_one_results.update(smart_chunked_answers_dict)
    one_to_one_results.update(full_context_answers_dict)

    ################## Crosswalk Fields ################## 
    one_to_one_results = aarete_derived.get_crosswalk_fields([one_to_one_results])

    ################## Fill NA Mapping ##################
    one_to_one_results = aarete_derived.fill_na_mapping(one_to_one_results)

    return one_to_one_results[0]


def run_one_to_n_prompts(filename, exhibit_dict, all_exhibit_headers, all_dataset):
    
    ################## RUN PROMPTS ##################
    one_to_n_results = []
    for exhibit_page in exhibit_dict.keys():
        
        exhibit_text = exhibit_dict[exhibit_page]
        exhibit_header = all_exhibit_headers[exhibit_page]

        ################## INITIALIZE FIELDS ##################
        reimbursement_level_fields = FieldSet(relationship="one_to_n", field_type="reimbursement_level", file_path=FIELD_JSON_PATH)

        ################## GET REIMBURSEMENT TIN/NPI ##################
        tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(exhibit_text, reimbursement_level_fields, filename)

        ################## GET EXHIBIT-LEVEL ANSWERS ##################
        exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(exhibit_text, filename)
        exhibit_level_answers['EXHIBIT_PAGE'] = exhibit_page
        exhibit_level_answers['EXHIBIT_TITLE'] = exhibit_header

        ################## GET DYNAMIC-PRIMARY ANSWERS ##################
        dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_text, 
                                                                                                                    exhibit_header, 
                                                                                                                    filename,
                                                                                                                    FieldSet(relationship="one_to_n", field_type="dynamic", file_path=config.FIELD_JSON_PATH))
        exhibit_level_answers.update(dynamic_to_exhibit_level_answers)
        reimbursement_level_fields.combine(dynamic_to_reimbursement_level_fields, inplace=True)
        
        ################## GET DYNAMIC-CODE ANSWERS ##################
        code_to_exhibit_level_answers, code_to_reimbursement_level_fields  = dynamic_funcs.get_dynamic_answers(exhibit_text,
                                                                                                                exhibit_header, 
                                                                                                                filename,
                                                                                                                FieldSet(relationship="one_to_n", field_type="dynamic_code", file_path=config.FIELD_JSON_PATH))
        exhibit_level_answers.update(code_to_exhibit_level_answers)
        reimbursement_level_fields.combine(code_to_reimbursement_level_fields, inplace=True)

        ################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
        reimbursement_level_answers = reimbursement_level(exhibit_text, filename, reimbursement_level_fields, all_dataset, exhibit_page) # Return list of dictionaries

        ################# COMBINE ANSWERS ##################
        full_answer_dict = combine_one_to_n_answers(exhibit_level_answers, reimbursement_level_answers, tin_npi_answers)
        one_to_n_results += full_answer_dict
    
    ################## Crosswalk Fields ################## 
    one_to_n_results = aarete_derived.get_crosswalk_fields(one_to_n_results)

    ################## Fill NA Mapping ##################
    one_to_n_results = aarete_derived.fill_na_mapping(one_to_n_results)

    ################## Update LOB for Duals ##################
    one_to_n_results = investment_postprocessing_funcs.update_lob_for_duals(one_to_n_results)

    ################## Add N/A Dynamic or Exhibit to One-to-One ##################
    dynamic_one_to_one_fields = dynamic_funcs.get_dynamic_one_to_one_fields(one_to_n_results)

    ################## CONVERT TO DF ##################
    one_to_n_df = pd.DataFrame(one_to_n_results)
    
    return one_to_n_df, dynamic_one_to_one_fields