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

import logging
import os
from collections import defaultdict

import pandas as pd
import src.config as config
import src.utils.string_utils as string_utils
from constants.constants import Constants
from constants.delimiters import Delimiter
from src.crosswalk.crosswalk_builder import CrosswalkBuilder
from src.prompts.fieldset import FieldSet
from src.utils.crosswalk_utils import apply_crosswalk


def fill_na_from_field(results_dict, to_field, from_field, crosswalk_path):
    from_str = results_dict.get(from_field)
    if string_utils.is_empty(from_str):
        return results_dict.get(to_field)
    MAPPING_DICT = CrosswalkBuilder().from_json(crosswalk_path).mapping
    to_value = apply_crosswalk(from_str, MAPPING_DICT)
    return to_value if to_value else results_dict.get(to_field)


def fill_na_mapping(answer_dicts):
    """
    Fill in missing values in fields using mappings from other fields.
    Always checks all crosswalks (PROGRAM and PRODUCT) and merges unique values.

    Args:
        answer_dicts (list[dict]): List of answer dictionaries

    Returns:
        list[dict]: List of answer dictionaries with filled/merged values
    """
    for answer_dict in answer_dicts:
        # Collect all AARETE_DERIVED_LOB values from different sources
        all_lob_values = set()
        
        # Get existing AARETE_DERIVED_LOB values (if any)
        existing_lob = answer_dict.get("AARETE_DERIVED_LOB", "")
        if not string_utils.is_empty(existing_lob):
            if "|" in existing_lob:
                all_lob_values.update(v.strip() for v in existing_lob.split("|") if v.strip() and v.strip() != "N/A")
            elif existing_lob.strip() and existing_lob.strip() != "N/A":
                all_lob_values.add(existing_lob.strip())
        
        # Get AARETE_DERIVED_LOB from AARETE_DERIVED_PROGRAM crosswalk (always check if PROGRAM exists)
        program_lob_values = set()
        if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
            program_filled_value = fill_na_from_field(
                answer_dict,
                "AARETE_DERIVED_LOB",
                "AARETE_DERIVED_PROGRAM",
                "constants/mappings/crosswalk_program_lob.json",
            )
            if not string_utils.is_empty(program_filled_value) and program_filled_value != "N/A":
                if "|" in program_filled_value:
                    program_lob_values = set(v.strip() for v in program_filled_value.split("|") if v.strip() and v.strip() != "N/A")
                elif program_filled_value.strip() and program_filled_value.strip() != "N/A":
                    program_lob_values = {program_filled_value.strip()}
                all_lob_values.update(program_lob_values)
        
        # Get AARETE_DERIVED_LOB from PRODUCT crosswalk (always check if PRODUCT exists)
        product_lob_values = set()
        if not string_utils.is_empty(answer_dict.get("PRODUCT")):
            product_filled_value = fill_na_from_field(
                answer_dict,
                "AARETE_DERIVED_LOB",
                "PRODUCT",
                "constants/mappings/crosswalk_product_lob.json",
            )
            if not string_utils.is_empty(product_filled_value) and product_filled_value != "N/A":
                if "|" in product_filled_value:
                    product_lob_values = set(v.strip() for v in product_filled_value.split("|") if v.strip() and v.strip() != "N/A")
                elif product_filled_value.strip() and product_filled_value.strip() != "N/A":
                    product_lob_values = {product_filled_value.strip()}
                all_lob_values.update(product_lob_values)
        
        # Set the merged, deduplicated value
        if all_lob_values:
            answer_dict["AARETE_DERIVED_LOB"] = "|".join(sorted(all_lob_values))
            # Set relationship flags if values came from PROGRAM or PRODUCT
            if program_lob_values:
                answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"
            if product_lob_values:
                answer_dict["LOB_PRODUCT_RELATIONSHIP"] = "Exclusive"
        elif string_utils.is_empty(answer_dict.get("AARETE_DERIVED_LOB")):
            # If still empty after all attempts, keep it as is (will be N/A or empty)
            pass
    
    return answer_dicts


def get_crosswalk_fields(answer_dicts: list, constants: Constants):
    crosswalk_fields = FieldSet(file_path=config.FIELD_JSON_PATH, crosswalk=True)
    for to_field in crosswalk_fields.fields:
        to_field_name, from_field_name = to_field.field_name, to_field.base_field
        # Find crosswalk
        crosswalk = constants.get_constant(to_field.crosswalk)
        if crosswalk:
            for answer_dict in answer_dicts:
                to_field_value, from_field_value = answer_dict.get(
                    to_field_name
                ), answer_dict.get(from_field_name)
                # If from_field_value is populated and to_field_value is not populated, perform mapping
                if not string_utils.is_empty(
                    from_field_value
                ) and string_utils.is_empty(to_field_value):
                    # Get list of values to map
                    if "|" in from_field_value:
                        from_field_value_list = from_field_value.split("|")
                    elif "," in from_field_value:
                        from_field_value_list = from_field_value.split(",")
                    else:
                        from_field_value_list = [from_field_value]

                    # Map each value in the list
                    to_field_answer_list = []
                    for individual_from_field_value in from_field_value_list:
                        individual_from_field_value = individual_from_field_value.strip() 

                        if individual_from_field_value in crosswalk.mapping.keys():
                            # Value is a key in the mapping - map it to the target value
                            to_field_answer_list.append(crosswalk.mapping.get(
                                individual_from_field_value
                            ))
                        elif individual_from_field_value in crosswalk.mapping.values():
                            # Value is already in the target format (e.g., "CHIP" when target is "CHIP")
                            # Keep it as-is for AARETE_DERIVED fields, or use reverse mapping for others
                            if "AARETE_DERIVED" in to_field_name:
                                to_field_answer_list.append(individual_from_field_value)
                            else:
                                to_field_answer_list.append(crosswalk.create_reverse_mapping().get(individual_from_field_value))
                    answer_dict[to_field_name] = "|".join(to_field_answer_list)

    return answer_dicts