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

import ast
import hashlib
import os
import re
from datetime import datetime

import pandas as pd
from src import postprocessing_funcs as generic_postprocessing_funcs
from src.constants import investment_values
from src.prompts.investment_prompts import FieldSet, invoke_derived_term_date
from src.utils import string_utils
import src.prompts.investment_prompts as investment_prompts
import src.utils.llm_utils as llm_utils
from src.enums.delimiters import Delimiter
import src.config as config

# Determine the base directory for the project
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Path to the mappings folder
MAPPINGS_DIR = os.path.join(BASE_DIR, 'crosswalk', 'mappings')



def normalize_indicator_field(value: str) -> str:
    """
    Normalize the indicator field value to 'Y' or 'N'.
    
    Args:
        value (str): The input value from an _IND field.
    Returns:
        str: 'Y' if the value is 'Y', otherwise 'N'.
    """
    if isinstance(value, str) and value.strip().upper() == "Y":
        return "Y"
    return 'N'


def format_rate_fields_with_commas(value: str) -> str:
    """
    Format the rate with commas and two decimal places.
    
    Args:
        value (str): A string representing a numeric value.
    Returns:
        str: The formatted rate string.
    """
    if string_utils.is_empty(value):
        return ""
    
    try:
        numeric_value = float(value)
        return f"{numeric_value:,.2f}"
    except ValueError:
        return ""

    
def remove_hyphens(value: str) -> str:
    """
    Remove hyphens from the input string.

    Args:
        value (str): The input string.
    Returns:
        str: The cleaned string with hyphens removed.
    """
    if not isinstance(value, str):
        return ""
    return value.replace("-", "")


def flatten_singleton_string_list(list_str: str) -> str:
    """Take in a list within a string, and if it has a single element, return that element.
    If it has multiple elements, return the list as a string.

    Args:
        list_str (str): A string representation of a list.

    Returns:
        str: The single element if the list has one element, otherwise the list as a string.
    """
    if string_utils.is_empty(list_str):
        return ""
    try:
        # Convert the string to a list using ast.literal_eval
        list_obj = ast.literal_eval(list_str)
        # Check if list_obj is actually a list
        if not isinstance(list_obj, list):
            return list_str
        # If the list has one element, return that element
        if len(list_obj) == 1:
            return list_obj[0]
        # Otherwise, return the list as a string
        return ", ".join(list_obj)
    except (ValueError, SyntaxError):
        return list_str

def validate_and_reformat_date(date_str: str) -> str:
    """
    Validate if a string is in YYYY/MM/DD format or reformat it to YYYY/MM/DD if possible.
    
    Args:
        date_str (str): The input date string.
        
    Returns:
        str: The reformatted date string if valid or reformatted, otherwise the original string.
    """
    if not isinstance(date_str, str):
        return date_str

    try:
        # Check if the date is already in YYYY/MM/DD format
        datetime.strptime(date_str, "%Y/%m/%d")
        return date_str
    except ValueError:
        # Attempt to reformat the date from known formats
        known_formats = [
            "%Y-%m-%d", "%m/%d/%Y", "%d-%b-%Y", "%d/%m/%Y"
        ]
        for fmt in known_formats:
            try:
                return datetime.strptime(date_str, fmt).strftime("%Y/%m/%d")
            except ValueError:
                continue
    return date_str

def date_postprocess(df, field_json_path):
    """
    Postprocess the date fields in the DataFrame.
    """
    # Date field postprocessing using FieldSet definition
    date_fields = FieldSet(file_path=field_json_path).filter(format="date")
    for field in date_fields.fields:
        if field.field_name in df.columns:
            df[field.field_name] = list(map(generic_postprocessing_funcs.date_postprocessing, df[field.field_name]))

    # Derived termination date
    df['AARETE_DERIVED_TERMINATION_DT'] = list(map(invoke_derived_term_date, df['AARETE_DERIVED_EFFECTIVE_DT'], df['TERMINATION_DT']))
    
    return df

def normalize_auto_renewal_term(text: str) -> str:
    """
    Normalize the auto-renewal term to a standard format.
    
    Args:
        text (str): The input auto-renewal term.
        
    Returns:
        str: The normalized auto-renewal term.
    """
    if not text or not isinstance(text, str):
        return ""
    
    # Convert to lowercase and strip whitespace
    text = text.lower().strip()
    
    # Remove numbers enclosed in parentheses (e.g., "(1) ", "(12) ")
    text = re.sub(r"\(\d+\) ", "", text)
    
    # Define replacements
    replacements = {
        r"\bone\b": "1",
        r"\btwo\b": "2",
        r"\bthree\b": "3",
        r"\btwelve\b": "12",
        r"year to year": "1 year",
        r"12 months": "1 year",
        r"\btwelve months\b": "1 year",
        r"\bone year\b": "1 year",
        r"\bone-year\b": "1 year",
        r"\b1 year\b": "1 year",
        r"\bmonth to month\b": "1 month",
    }

    # Apply replacements
    for pattern, replacement in replacements.items():
        text = re.sub(pattern, replacement, text)
    
    # Remove any remaining parentheses or extra spaces
    text = text.strip().replace("(", "").replace(")", "")
    
    text = text.replace("-", " ") 
    if text == "year":
        text = "1 year"
    
    return text

def normalize_cpt_fields(value):
    """
    Normalize code values to a list format.
    
    Args:
        value (Any): Input value which could be a list, string, or other types.
    
    Returns:
        list: A normalized list of string codes.
    """
    if not value or (isinstance(value, float) and pd.isna(value)):  # Handle NaN values
        return "" # Empty string

    # First normalize to a list
    if isinstance(value, (list, tuple)):  # If already a list or tuple
        result = [str(v).strip() for v in value]
    elif isinstance(value, str):  # If it's a string
        value = value.strip()
        if value.startswith("[") and value.endswith("]"):  # String representation of a list
            try:
                # Parse it first
                parsed_list = ast.literal_eval(value)
                result = []
                # Process each item in the parsed list
                for item in parsed_list:
                    item_str = str(item).strip()
                    # If the item contains commas, split it into multiple items
                    if "," in item_str:
                        result.extend([str(v).strip() for v in item_str.split(",")])
                    else:
                        result.append(item_str)
            except (ValueError, SyntaxError):
                result = [value]  # If parsing fails, treat it as a single value
        elif "," in value: # If it's a comma-separated string
            # Split by comma and create a list
            result = [str(v).strip() for v in value.split(",")]
        elif "-" in value:
            # Range format
            result = [value]
        else:  # Single value
            result = [value]
    elif isinstance(value, (float, int)):  # If it's a number
        result = [str(int(value)).strip()]
    else:
        result = [str(value).strip()]  # Default case for other types
    
    return str(result)  # Return the list directly


def remove_update_reimbursement(df):
        
    # Remove rows where 'AARETE_DERIVED_REIMB_METHOD' value is "Medicare Member Cost Share", "Incentive Payment" or other invalid values
    # Null values are temporarily retained, as all carve-out cases currently have null values for the Aarete-derived reimbursement method
    if 'AARETE_DERIVED_REIMB_METHOD' in df.columns:
        df = df[(df['AARETE_DERIVED_REIMB_METHOD'].isin(investment_values.VALID_REIMB_TERM)) | (df['AARETE_DERIVED_REIMB_METHOD'].isna())]

    return df

def generate_reimb_ids(df: pd.DataFrame) -> pd.DataFrame:
    """Generates REIMB_ID and REIMB_LESSER_OF_ID for each row in the DataFrame based on 
        the FILE_NAME and other fields.
    The REIMB_ID is constructed using the following format:
    {FILE_NAME}_exh_pg_{EXHIBIT_PAGE}_{index}_{hash_value}

    the REIMB_LESSER_OF_ID is just the {hash_value} itself.

    Where:
    - FILE_NAME is the name of the file (without extension).
    - EXHIBIT_PAGE is the page number of the exhibit.
    - index is the index of the row in the group.
    - hash_value is the first 8 characters of the MD5 hash of a concatenated string of relevant fields.
    The relevant fields are SERVICE_TERM, REIMB_TERM, PROGRAM, PRODUCT, NETWORK, and LOB.

    Args:
        df (pd.DataFrame): Input DataFrame containing the columns FILE_NAME, EXHIBIT_PAGE, SERVICE_TERM, REIMB_TERM, PROGRAM, PRODUCT, NETWORK, and LOB.
        The FILE_NAME column should contain the name of the file (with or without extension).
        The EXHIBIT_PAGE column should contain the page number of the exhibit.

    Returns:
        pd.DataFrame: DataFrame with generated REIMB_IDs and REIMB_LESSER_OF_IDs for each row.
    """
    # First sort to ensure consistent ordering
    df = df.sort_values(by=['FILE_NAME', 'EXHIBIT_PAGE']) if "EXHIBIT_PAGE" in df.columns else df.sort_values(by=['FILE_NAME'])

    # Create a temporary copy to avoid SettingWithCopyWarning
    df_temp = df.copy()

    # Process each file+exhibit_page group to reset counter for each exhibit page
    group_cols = ['FILE_NAME', 'EXHIBIT_PAGE'] if "EXHIBIT_PAGE" in df.columns else ['FILE_NAME']
    for group_key, indices in df.groupby(group_cols).groups.items():
        if isinstance(group_key, tuple):
            filename, exhibit_page = group_key
            exhibit_page = str(exhibit_page).zfill(3)
        else:
            filename = group_key
            exhibit_page = '000'
        clean_filename = str(filename).split("/")[-1].replace(".", "_")

        for i, idx in enumerate(indices):
            row = df.loc[idx]

            # Get fields to include in hash
            service_term = str(row.get('SERVICE_TERM', ''))
            reimb_term = str(row.get('REIMB_TERM', ''))
            program = str(row.get('PROGRAM', ''))
            product = str(row.get('PRODUCT', ''))
            network = str(row.get('NETWORK', ''))
            lob = str(row.get('LOB', ''))

            # Create hash string
            hash_string = f"{clean_filename}_{service_term}_{reimb_term}_{program}_{product}_{network}_{lob}"
            hash_value = hashlib.md5(hash_string.encode()).hexdigest()[:8]

            # Generate REIMB_ID and REIMB_LESSER_OF_ID
            df_temp.at[idx, 'REIMB_ID'] = f"{clean_filename}_exh_pg_{exhibit_page}_{i+1:03d}_{hash_value}"
            df_temp.at[idx, 'REIMB_LESSER_OF_ID'] = hash_value

    return df_temp

def reorder_columns(df: pd.DataFrame, column_order: list[str]) -> pd.DataFrame:
    """
    Reorders the columns of the DataFrame based on the given column order.
    Steps:
    1. Adds any missing columns from `column_order` to the DataFrame, filled with empty strings.
    2. Reorders the columns of the DataFrame to match `column_order`.
    3. Appends any columns in the DataFrame that are not in `column_order` to the end.
    Args:
        df (pd.DataFrame): The input DataFrame.
        column_order (list[str]): The desired column order.
    Returns:
        pd.DataFrame: The reordered DataFrame.
    """
    # Add missing columns from column_order to the DataFrame, filled with empty strings
    for col in column_order:
        if col not in df.columns:
            df[col] = ""
    # Reorder columns to match column_order
    ordered_columns = [col for col in column_order if col in df.columns]
    # Add columns in df that are not in column_order to the end
    remaining_columns = [col for col in df.columns if col not in column_order]
    # Combine the ordered columns and remaining columns
    final_column_order = ordered_columns + remaining_columns
    # Return the DataFrame with reordered columns
    return df[final_column_order]

def auto_renewal(df: pd.DataFrame) -> pd.DataFrame:
    # Normalize the 'AUTO_RENEWAL_TERM' column
    if "AUTO_RENEWAL_TERM" in df.columns:
        df["AUTO_RENEWAL_TERM"] = df["AUTO_RENEWAL_TERM"].apply(normalize_auto_renewal_term)
    return df


def update_termination_date_for_conditions(df: pd.DataFrame) -> pd.DataFrame:
    """
    Updates the 'AARETE_DERIVED_TERMINATION_DT' column based on specific conditions 
    involving 'TERMINATION_DT', 'AUTO_RENEWAL_IND', and 'AARETE_DERIVED_TERMINATION_DT'.

    Args:
        df (pd.DataFrame): The input DataFrame.

    Returns:
        pd.DataFrame: The updated DataFrame with modified 'AARETE_DERIVED_TERMINATION_DT' values.
    """

    if "TERMINATION_DT" in df.columns and "AUTO_RENEWAL_IND" in df.columns and "AARETE_DERIVED_TERMINATION_DT" in df.columns:
        # Replace Y with 9999/12/31
        df.loc[
            df["AUTO_RENEWAL_IND"] == "Y", 
            "AARETE_DERIVED_TERMINATION_DT"
        ] = "9999/12/31"

        # Replace "N" and "year to year" with 9999/12/31
        df.loc[
            (df["TERMINATION_DT"] == "year to year") & 
            (df["AUTO_RENEWAL_IND"] == "N") & 
            (df["AARETE_DERIVED_TERMINATION_DT"] == "N/A"), 
            "AARETE_DERIVED_TERMINATION_DT"
        ] = "9999/12/31"
    
    return df

def standardize_reimb_method_and_fee_schedule(df: pd.DataFrame) -> pd.DataFrame:
    """
    Standardizes the 'AARETE_DERIVED_REIMB_METHOD' column and updates the 
    'AARETE_DERIVED_FEE_SCHEDULE' column based on specific values.

    Args:
        df (pd.DataFrame): The input DataFrame.

    Returns:
        pd.DataFrame: The updated DataFrame with standardized columns.
    """
    if "AARETE_DERIVED_REIMB_METHOD" in df.columns and 'AARETE_DERIVED_FEE_SCHEDULE' in df.columns:
        # Standardize 'AARETE_DERIVED_REIMB_METHOD' values
        # Copy AWP, ASP, WAC to AARETE_DERIVED_FEE_SCHEDULE column for the corresponding row
        medrx_fee_schedule_values = ['AWP', 'ASP', 'WAC']
        df.loc[df['AARETE_DERIVED_REIMB_METHOD'].isin(medrx_fee_schedule_values), 'AARETE_DERIVED_FEE_SCHEDULE'] = df['AARETE_DERIVED_REIMB_METHOD']

        # Convert AARETE_DERIVED_REIMB_METHOD values [AWP, ASP, WAC] to 'MedRx'
        df['AARETE_DERIVED_REIMB_METHOD'] = df['AARETE_DERIVED_REIMB_METHOD'].replace({'AWP': 'MedRx', 'ASP': 'MedRx', 'WAC': 'MedRx'})

    # Remove rows where 'AARETE_DERIVED_REIMB_METHOD' value is "Medicare Member Cost Share" or "Incentive Payment"
    if 'AARETE_DERIVED_REIMB_METHOD' in df.columns:
        df = df[~df['AARETE_DERIVED_REIMB_METHOD'].isin(["Medicare Member Cost Share"] + investment_values.INVALID_REIMB_TERM)]

    # standardize the 'UNIT_OF_MEASURE' column
    if "UNIT_OF_MEASURE" in df.columns:
        df["UNIT_OF_MEASURE"] = df["UNIT_OF_MEASURE"].str.replace("Per Encounter", "Per Visit")

        # remove cases of Per Member Per Month (PMPM)
        df = df[~df['UNIT_OF_MEASURE'].isin(['Per Member Per Month (PMPM)'])]

        # Ensure UNIT_OF_MEASURE is one of the valid values
        df['UNIT_OF_MEASURE'] = df['UNIT_OF_MEASURE'].apply(lambda x: x if x in investment_values.VALID_UNIT_OF_MEASURE else "")

    return df


def process_patient_age_range(df):
    """
    Process the 'PATIENT_AGE_RANGE' column into 'PATIENT_AGE_MIN' and 'PATIENT_AGE_MAX'. 
    If a single number is provided, both min and max are set to that number.
    Remove PATIENT_AGE_RANGE column.
    
    Args:
        df (pd.DataFrame): The input DataFrame.
        
    Returns:
        pd.DataFrame: The updated DataFrame with 'PATIENT_AGE_MIN' and 'PATIENT_AGE_MAX' columns.
    """
    if "PATIENT_AGE_RANGE" in df.columns:
        try:
            # Create temporary columns for processing
            df['PATIENT_AGE_MIN'] = df['PATIENT_AGE_RANGE']
            df['PATIENT_AGE_MAX'] = df['PATIENT_AGE_RANGE']
            
            # Handle hyphenated ranges
            mask = df['PATIENT_AGE_RANGE'].str.contains('-', na=False)
            if mask.any():
                temp_split = df.loc[mask, 'PATIENT_AGE_RANGE'].str.split("-", expand=True)
                df.loc[mask, 'PATIENT_AGE_MIN'] = temp_split[0]
                df.loc[mask, 'PATIENT_AGE_MAX'] = temp_split[1]
            
        except Exception as e:
            # If any error occurs, keep the original values
            pass
        
        # Remove the original PATIENT_AGE_RANGE column
        df.drop(columns=['PATIENT_AGE_RANGE'], inplace=True, errors='ignore')
    
    return df

def add_aarete_derived_amendment_num(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add the 'AARETE_DERIVED_AMENDMENT_NUM' column to the DataFrame.

    Args:
        df (pd.DataFrame): The input DataFrame.

    Returns:
        pd.DataFrame: The updated DataFrame with the 'AARETE_DERIVED_AMENDMENT_NUM' column.
    """

    if "CONTRACT_AMENDMENT_NUM" in df.columns:
        df["AARETE_DERIVED_AMENDMENT_NUM"] = df["CONTRACT_AMENDMENT_NUM"].str.split('-').str[1].fillna(df['CONTRACT_AMENDMENT_NUM'])

    return df

def add_aarete_derived_product(df):
    """
    Add AARETE_DERIVED_PRODUCT column to df, based on PRODUCT column.
    
    The new column should be Title Case of the PRODUCT column value.
    """
    if "PRODUCT" in df.columns:
        df["AARETE_DERIVED_PRODUCT"] = df["PRODUCT"].str.title()
    return df


def update_lob_for_duals(answer_dicts):
    final_answer_dicts = []
    # Iterate through each row in the dataframe
    for answer_dict in answer_dicts:
        # Get SERVICE_TERM value
        service_term = answer_dict.get('SERVICE_TERM', "").lower()
        # Check if the SERVICE_TERM contains "Medicare" and "Medicaid" - case insensitive
        if 'medicare' in service_term and 'medicaid' in service_term:
            prompt = investment_prompts.DUAL_LOB_CHECK(service_term)
            claude_answer_raw = llm_utils.invoke_claude(prompt, model_id="sonnet_latest", filename="", max_tokens=200)
            claude_answer_extracted = string_utils.extract_text_from_delimiters(claude_answer_raw, Delimiter.PIPE)
            answer_dict["AARETE_DERIVED_LOB"] = claude_answer_extracted
        final_answer_dicts.append(answer_dict)
    return final_answer_dicts

def fill_empty_reimb_pct_rate(df):
    """
    Fill 'REIMB_PCT_RATE' based on reimbursement method:
    - Set 'REIMB_PCT_RATE' to '100' for rows where:
      - 'REIMB_PCT_RATE', 'REIMB_FEE_RATE', and 'REIMB_CONVERSION_FACTOR' are all empty.
    - Set 'REIMB_PCT_RATE' to 'N/A' for rows where:
      - 'AARETE_DERIVED_REIMB_METHOD' contains specific grouper.
    """
    required_cols = ["REIMB_PCT_RATE", "REIMB_FEE_RATE", "REIMB_CONVERSION_FACTOR", "AARETE_DERIVED_REIMB_METHOD"]
    
    if all(col in df.columns for col in required_cols):
        # Fill with '100' when all three rate fields are empty
        pct_empty = string_utils.is_empty(df["REIMB_PCT_RATE"])
        fee_empty = string_utils.is_empty(df["REIMB_FEE_RATE"])
        conversion_empty = string_utils.is_empty(df["REIMB_CONVERSION_FACTOR"])
        
        fill_mask = pct_empty & fee_empty & conversion_empty
        df.loc[fill_mask, "REIMB_PCT_RATE"] = "100"
        
        # Set to N/A for specific reimbursement methods
        na_methods = ["Grouper", "Per Diem", "Flat Rate", "Per Visit", "Case Rate"]
        
        na_mask = df["AARETE_DERIVED_REIMB_METHOD"].isin(na_methods)
        df.loc[na_mask, "REIMB_PCT_RATE"] = "N/A"
    
    return df

def remove_redundant_reimb_info(df):
    """
    Checks if "REIMB_" fields are the same as their corresponding 1:1 fields.
    If they are the same, replace the "REIMB_" values with ""

    Args:
        df (pd.DataFrame): The input DataFrame.
    
    Returns:
        pd.DataFrame: The updated DataFrame with redundant reimbursement field values removed.
    """
    COLUMN_PAIRS = [
        ("REIMB_EFFECTIVE_DT", "AARETE_DERIVED_EFFECTIVE_DT"),
        ("REIMB_TERMINATION_DT", "AARETE_DERIVED_TERMINATION_DT"),
        ("REIMB_PROV_NAME", "PROV_GROUP_NAME_FULL")
    ]
    for column_tuple in COLUMN_PAIRS:
        reimb_col = column_tuple[0]
        one_to_one_col = column_tuple[1]
        if all(col in df.columns for col in [reimb_col, one_to_one_col]):
            mask = df[reimb_col] == df[one_to_one_col]
            df.loc[mask, reimb_col] = ""
    return df



def deduplicate_provider_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Deduplicate provider fields in the DataFrame.

    This function removes PROV_GROUP_* values from the PROV_OTHER_* columns if they are already present.
    In addition, it deduplicates PROV_OTHER_* values.

    E.g. If PROV_GROUP_TIN is "123456789" and PROV_OTHER_TIN is "123456789 | 987654321 | 987654321",
    then PROV_OTHER_TIN will be updated to "987654321".


    Args:
        df (pd.DataFrame): The input DataFrame.

    Returns:
        pd.DataFrame: The updated DataFrame with deduplicated provider fields.
    """
    provider_columns = [
        "PROV_GROUP_TIN", "PROV_GROUP_NPI", "PROV_GROUP_NAME_FULL",
        "PROV_OTHER_TIN", "PROV_OTHER_NPI", "PROV_OTHER_NAME_FULL"
    ]
    other_columns = [
        "PROV_OTHER_TIN", "PROV_OTHER_NPI", "PROV_OTHER_NAME_FULL"
    ]

    # Check if all provider columns exist in the DataFrame
    if all(col in df.columns for col in provider_columns):
        # Process each row
        for index, row in df.iterrows():
            # Get GROUP values for this row
            group_tin = str(row.get("PROV_GROUP_TIN", "")).strip()
            group_npi = str(row.get("PROV_GROUP_NPI", "")).strip()
            group_name = str(row.get("PROV_GROUP_NAME_FULL", "")).strip()

            # Simple mapping of OTHER columns to their corresponding GROUP values
            column_mappings = {
                "PROV_OTHER_TIN": group_tin,
                "PROV_OTHER_NPI": group_npi,
                "PROV_OTHER_NAME_FULL": group_name
            }

            # Process each OTHER column
            for other_col, group_val in column_mappings.items():
                other_value = str(row.get(other_col, "")).strip()
                
                if other_value and not string_utils.is_empty(other_value) and other_value != "UNKNOWN":
                    # Split by pipe and strip whitespace
                    other_list = [v.strip() for v in other_value.split("|") if v.strip()]

                    # Remove corresponding GROUP value if present
                    if group_val and not string_utils.is_empty(group_val) and group_val != "UNKNOWN":
                        other_list = [v for v in other_list if v != group_val]

                    # Remove duplicates while preserving order
                    seen = set()
                    deduped_list = []
                    for item in other_list:
                        if item not in seen and not string_utils.is_empty(item) and item != "UNKNOWN":
                            seen.add(item)
                            deduped_list.append(item)
                    
                    # Update the DataFrame
                    df.at[index, other_col] = " | ".join(deduped_list) if deduped_list else ""
    return df