doczyai-pipelines/src/pipelines/shared/postprocessing/aarete_derived.py

import logging
import re
import pandas as pd
import src.config as config
import src.utils.string_utils as string_utils
from src.constants.constants import Constants
from src.constants.investment_columns import FIELD_FORMAT_MAPPING
from src.crosswalk.crosswalk_builder import CrosswalkBuilder
from src.pipelines.saas.prompts import prompt_calls
from src.prompts.fieldset import FieldSet
from src.utils.crosswalk_utils import apply_crosswalk
from src.utils.formatting_utils import normalize_field_value


def _canonical_lob(val: str, valid_lobs_lower: dict[str, str]) -> str:
    """Return the canonical-cased LOB from valid_lobs (case-insensitive match).
    Falls back to title case if no match found."""
    return valid_lobs_lower.get(val.strip().lower(), val.strip().title())


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(df: pd.DataFrame, constants: Constants) -> pd.DataFrame:
    """
    Derive AARETE_DERIVED_LOB from AARETE_DERIVED_PROGRAM and AARETE_DERIVED_PRODUCT.
    Uses crosswalk mapping when available; falls back to LLM when the mapping is empty.
    Merges unique LOB values from both sources.

    Args:
        df (pd.DataFrame): DataFrame with extracted fields.
        constants (Constants): Shared Constants instance (loaded once at startup).

    Returns:
        pd.DataFrame: DataFrame with filled/merged AARETE_DERIVED_LOB values.
    """
    if df.shape[0] == 0:
        return df

    filename = df["FILE_NAME"].iloc[0] if "FILE_NAME" in df.columns else "unknown"
    payer_name_raw = df["PAYER_NAME"].iloc[0] if "PAYER_NAME" in df.columns else ""
    payer_state_raw = df["PAYER_STATE"].iloc[0] if "PAYER_STATE" in df.columns else ""
    payer_name = (
        str(payer_name_raw) if not string_utils.is_empty(payer_name_raw) else ""
    )
    payer_state = (
        str(payer_state_raw) if not string_utils.is_empty(payer_state_raw) else ""
    )

    answer_dicts = df.to_dict(orient="records")
    # Crosswalk values are the AARETE_DERIVED_LOB values - pass these to the LLM
    lob_mapping_values = constants.CROSSWALK_LOB.mapping.values()
    valid_lobs = sorted(
        {
            item
            for v in lob_mapping_values
            for item in (v if isinstance(v, list) else [v])
            if item and item != "N/A"
        }
    )
    # Case-insensitive lookup → canonical form (e.g. "MEDICAID" → "Medicaid")
    valid_lobs_lower: dict[str, str] = {v.lower(): v for v in valid_lobs}

    def _to_str_list(val) -> list:
        """Normalize a field value to a deduplicated list of non-empty strings."""
        if string_utils.is_empty(val):
            return []
        if isinstance(val, list):
            return [str(v).strip() for v in val if str(v).strip()]
        if isinstance(val, str):
            return [val.strip()] if val.strip() else []
        return []

    # --- Pre-compute LOB caches keyed by unique input tuples ---
    # Collect unique keys across all rows so each distinct value is resolved once.
    unique_aarete_program_keys: set[tuple] = set()
    unique_program_keys: set[tuple] = set()
    unique_aarete_product_keys: set[tuple] = set()
    unique_product_keys: set[tuple] = set()

    for answer_dict in answer_dicts:
        adp = _to_str_list(answer_dict.get("AARETE_DERIVED_PROGRAM"))
        if adp:
            unique_aarete_program_keys.add(tuple(adp))
        p = _to_str_list(answer_dict.get("PROGRAM"))
        if p:
            unique_program_keys.add(tuple(p))
        adprod = _to_str_list(answer_dict.get("AARETE_DERIVED_PRODUCT"))
        if adprod:
            unique_aarete_product_keys.add(tuple(adprod))
        prod = _to_str_list(answer_dict.get("PRODUCT"))
        if prod:
            unique_product_keys.add(tuple(prod))

    # Crosswalk cache: AARETE_DERIVED_PROGRAM → LOB list
    program_crosswalk_cache: dict[tuple, list] = {}
    if constants.CROSSWALK_PROGRAM_LOB.mapping:
        for key in unique_aarete_program_keys:
            program_crosswalk_cache[key] = [
                constants.CROSSWALK_PROGRAM_LOB.mapping[v]
                for v in key
                if v in constants.CROSSWALK_PROGRAM_LOB.mapping
            ]

    # LLM cache: PROGRAM → LOB list (one call per unique program list)
    program_llm_cache: dict[tuple, list] = {}
    for key in unique_program_keys:
        program_llm_cache[key] = (
            prompt_calls.prompt_program_to_lob(
                list(key),
                valid_lobs,
                filename,
                payer_name=payer_name,
                payer_state=payer_state,
            )
            or []
        )

    # Crosswalk cache: AARETE_DERIVED_PRODUCT → LOB list
    product_crosswalk_cache: dict[tuple, list] = {}
    if constants.CROSSWALK_PRODUCT_LOB.mapping:
        for key in unique_aarete_product_keys:
            product_crosswalk_cache[key] = [
                constants.CROSSWALK_PRODUCT_LOB.mapping[v]
                for v in key
                if v in constants.CROSSWALK_PRODUCT_LOB.mapping
            ]

    # LLM cache: PRODUCT → LOB list (one call per unique product list)
    product_llm_cache: dict[tuple, list] = {}
    for key in unique_product_keys:
        product_llm_cache[key] = (
            prompt_calls.prompt_product_to_lob(
                list(key),
                valid_lobs,
                filename,
                payer_name=payer_name,
                payer_state=payer_state,
            )
            or []
        )

    # --- Apply cached results to each row ---
    for answer_dict in answer_dicts:
        all_lob_values = set()

        # Get existing AARETE_DERIVED_LOB values (if any)
        # Handle both list and string formats (JSON lists or pipe-delimited strings)
        existing_lob_list = answer_dict.get("AARETE_DERIVED_LOB", "")
        if not string_utils.is_empty(existing_lob_list):
            if isinstance(existing_lob_list, list):
                lob_items = existing_lob_list
            elif isinstance(existing_lob_list, str):
                lob_items = [existing_lob_list]
            else:
                lob_items = [str(existing_lob_list)]
            for lob_val in lob_items:
                if (
                    isinstance(lob_val, str)
                    and lob_val.strip()
                    and lob_val.strip() != "N/A"
                ):
                    all_lob_values.add(_canonical_lob(lob_val, valid_lobs_lower))

        # Get AARETE_DERIVED_LOB from AARETE_DERIVED_PROGRAM + PROGRAM (via caches)
        program_lob_values = set()
        adp_key = tuple(_to_str_list(answer_dict.get("AARETE_DERIVED_PROGRAM")))
        p_key = tuple(_to_str_list(answer_dict.get("PROGRAM")))
        if program_crosswalk_cache:
            logging.info(
                f"Using crosswalk cache for AARETE_DERIVED_PROGRAM key: {adp_key}"
            )
            program_lob_list = program_crosswalk_cache.get(adp_key, [])
        else:
            logging.info(f"Using LLM cache for PROGRAM key: {p_key}")
            program_lob_list = program_llm_cache.get(p_key, [])
        filtered = [
            _canonical_lob(v, valid_lobs_lower)
            for v in program_lob_list
            if isinstance(v, str) and v.strip() and v.strip() != "N/A"
        ]
        if filtered:
            program_lob_values.update(filtered)
            all_lob_values.update(filtered)

        # Get AARETE_DERIVED_LOB from AARETE_DERIVED_PRODUCT + PRODUCT (via caches)
        product_lob_values = set()
        adprod_key = tuple(_to_str_list(answer_dict.get("AARETE_DERIVED_PRODUCT")))
        prod_key = tuple(_to_str_list(answer_dict.get("PRODUCT")))
        if product_crosswalk_cache:
            logging.info(
                f"Using crosswalk cache for AARETE_DERIVED_PRODUCT key: {adprod_key}"
            )
            product_lob_list = product_crosswalk_cache.get(adprod_key, [])
        else:
            logging.info(f"Using LLM cache for PRODUCT key: {prod_key}")
            product_lob_list = product_llm_cache.get(prod_key, [])
        filtered = [
            _canonical_lob(v, valid_lobs_lower)
            for v in product_lob_list
            if isinstance(v, str) and v.strip() and v.strip() != "N/A"
        ]
        if filtered:
            product_lob_values.update(filtered)
            all_lob_values.update(filtered)

        # Set the merged, deduplicated value
        if all_lob_values:
            answer_dict["AARETE_DERIVED_LOB"] = sorted(all_lob_values)
            # Only set if field is empty or "N/A" (preserve LLM-determined Inclusive/Exclusive)
            if program_lob_values:
                existing_relationship = answer_dict.get("LOB_PROGRAM_RELATIONSHIP", "")
                if string_utils.is_empty(existing_relationship):
                    answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"
            if product_lob_values:
                existing_relationship = answer_dict.get("LOB_PRODUCT_RELATIONSHIP", "")
                if string_utils.is_empty(existing_relationship):
                    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

        # Handle cases where PROGRAM/PRODUCT is NA/blank
        # Set relationship to "Exclusive" as default when missing and relationship is not set
        if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
            if string_utils.is_empty(answer_dict.get("LOB_PROGRAM_RELATIONSHIP")):
                answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"

        if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PRODUCT")):
            if string_utils.is_empty(answer_dict.get("LOB_PRODUCT_RELATIONSHIP")):
                answer_dict["LOB_PRODUCT_RELATIONSHIP"] = "Exclusive"

    return pd.DataFrame(answer_dicts)


def get_crosswalk_fields(answer_dicts: list, constants: Constants):
    """
    Apply crosswalk mappings to derive target fields from source fields.

    Normalizes output based on FIELD_FORMAT_MAPPING to ensure correct format
    (str vs list[str]) for each target field.
    """

    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 = answer_dict.get(to_field_name)
                from_field_value = 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) or "N/A" in to_field_value
                ):
                    if isinstance(from_field_value, list):
                        from_field_value_list = from_field_value
                    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:
                                # Reverse mapping returns dict[str, list[str]], so .get() returns a list
                                # We need to extend, not append, to avoid nested lists
                                reverse_mapping = crosswalk.create_reverse_mapping()
                                reverse_result = reverse_mapping.get(
                                    individual_from_field_value
                                )
                                if reverse_result is not None:
                                    # reverse_result is a list, extend it
                                    to_field_answer_list.extend(reverse_result)
                                else:
                                    # If not found in reverse mapping, keep original value
                                    to_field_answer_list.append(
                                        individual_from_field_value
                                    )

                    # Normalize the result based on FIELD_FORMAT_MAPPING
                    format_type = FIELD_FORMAT_MAPPING.get(to_field_name, "list[str]")
                    normalized_value = normalize_field_value(
                        to_field_name, to_field_answer_list, format_type
                    )
                    answer_dict[to_field_name] = normalized_value

                    # DO NOT update from_field_name - preserve its original format
                    # The crosswalk should only create/update the target field (to_field_name),
                    # not modify the source field (from_field_name)
                    # If from_field_name needs to be a list, that should be handled elsewhere
    return answer_dicts


# ---------------------------------------------------------------------------
# Helpers shared by PROGRAM and PRODUCT canonical derivation
# ---------------------------------------------------------------------------


def _split_multi_value_separators(text: str) -> list[str]:
    """Split a PROGRAM/PRODUCT cell on all recognised list separators.

    In practice these columns use three separator conventions:

    * comma  ``","``            — primary Dashboard separator
    * slash  ``"/"``            — e.g. ``"Healthy Advantage/Healthy Advantage Plus"``
    * ampersand  ``" & "``      — e.g. ``"HA & HA+"``
      (requires surrounding whitespace to avoid splitting abbreviated compound
      names such as ``"B&G"`` or ``"A&M"``)

    Returns a list of stripped, non-empty parts.
    """
    parts = re.split(r",|/|\s+&\s+", text)
    return [p.strip() for p in parts if p.strip()]


def _collect_unique_values_from_list_column(
    column: pd.Series,
) -> list[str]:
    """Collect all unique, non-empty, non-sentinel values from a PROGRAM/PRODUCT column.

    In the Dashboard DataFrame (the only place this is called), ``dashboard_postprocess``
    has already run and converted every list column to a comma-separated string, e.g.
    ``"Medicaid Managed Care, Commercial PPO"``.
    This helper handles both that string form AND the original list form so it is safe
    regardless of call order.
    """
    _SENTINELS = {"N/A", "UNKNOWN", "NONE", "NULL"}
    seen: set[str] = set()
    unique: list[str] = []
    for val in column:
        if string_utils.is_empty(val):
            continue
        if isinstance(val, list):
            # Pre-dashboard-postprocess form: Python list — each element may itself
            # contain secondary separators (/ or &), so split each item further.
            raw_items = string_utils.flatten_to_strings(val)
            items = []
            for raw in raw_items:
                items.extend(_split_multi_value_separators(raw))
        else:
            # Post-dashboard-postprocess form: comma-separated string (may also
            # contain / or & separators inside individual cells).
            items = _split_multi_value_separators(str(val))
        for item in items:
            text = item.strip()
            if not string_utils.is_empty(text) and text.upper() not in _SENTINELS:
                if text not in seen:
                    seen.add(text)
                    unique.append(text)
    return unique


def _collect_unique_payer_value_combos(
    df: pd.DataFrame, value_col: str
) -> list[tuple[str, str, str]]:
    """Collect unique (payer_name, payer_state, raw_value) triples from the DataFrame.

    Iterates every row, splits the value_col cell on all recognised separators,
    and pairs each resulting item with the row's PAYER_NAME / PAYER_STATE.
    Combinations are deduplicated while preserving insertion order.

    Used by add_aarete_derived_program_canonical and add_aarete_derived_product_canonical
    to build the per-entry payer-context input for the batch LLM call.
    """
    _SENTINELS = {"N/A", "UNKNOWN", "NONE", "NULL"}
    seen: set[tuple[str, str, str]] = set()
    result: list[tuple[str, str, str]] = []
    has_payer_name = "PAYER_NAME" in df.columns
    has_payer_state = "PAYER_STATE" in df.columns
    for _, row in df.iterrows():
        payer_name = (
            str(row["PAYER_NAME"]).strip()
            if has_payer_name and not string_utils.is_empty(row["PAYER_NAME"])
            else ""
        )
        payer_state = (
            str(row["PAYER_STATE"]).strip()
            if has_payer_state and not string_utils.is_empty(row["PAYER_STATE"])
            else ""
        )
        val = row[value_col]
        if string_utils.is_empty(val):
            continue
        if isinstance(val, list):
            raw_items = string_utils.flatten_to_strings(val)
            items: list[str] = []
            for raw in raw_items:
                items.extend(_split_multi_value_separators(raw))
        else:
            items = _split_multi_value_separators(str(val))
        for item in items:
            text = item.strip()
            if not string_utils.is_empty(text) and text.upper() not in _SENTINELS:
                combo = (payer_name, payer_state, text)
                if combo not in seen:
                    seen.add(combo)
                    result.append(combo)
    return result


def _all_empty_in_list_column(column: pd.Series) -> bool:
    """Return True when every cell in a list[str] column is empty / null."""
    return all(string_utils.is_empty(val) for val in column)


def _map_row_list_to_canonicals(raw_val, canonical_map: dict[str, str]) -> str:
    """Map one PROGRAM/PRODUCT cell to a comma-separated string of canonical names.

    Input can be a Python list (pre-dashboard-postprocess) or a comma-separated
    string (post-dashboard-postprocess, which is the normal case for the Dashboard DF).
    Output is always a comma-separated string to match the Dashboard DF column format.
    """
    _SENTINELS = {"N/A", "UNKNOWN", "NONE", "NULL"}
    if string_utils.is_empty(raw_val):
        return ""
    if isinstance(raw_val, list):
        raw_items = string_utils.flatten_to_strings(raw_val)
        items = []
        for raw in raw_items:
            items.extend(_split_multi_value_separators(raw))
    else:
        # Post-dashboard-postprocess: comma-separated string (may also contain
        # / or & separators within individual cells, e.g. "HA/HA+").
        items = _split_multi_value_separators(str(raw_val))
    result: list[str] = []
    seen: set[str] = set()
    for v in items:
        text = v.strip()
        if string_utils.is_empty(text) or text.upper() in _SENTINELS:
            continue
        canonical = canonical_map.get(text, text.upper())
        if canonical and canonical not in seen:
            seen.add(canonical)
            result.append(canonical)
        elif canonical in seen:
            logging.debug(
                "_map_row_list_to_canonicals: dropped duplicate canonical '%s' "
                "for raw='%s' (already mapped from an earlier value in this cell)",
                canonical,
                text,
            )
    return ", ".join(result)


# ---------------------------------------------------------------------------
# Public API — called from runner.py (Dashboard only)
# ---------------------------------------------------------------------------


def add_aarete_derived_program_canonical(
    df: pd.DataFrame, valid_programs: list[str]
) -> pd.DataFrame:
    """Derive canonical program names from the raw PROGRAM column.

    Architecture (single global LLM call):
    1. Collect all unique raw PROGRAM values across the entire Dashboard DataFrame.
    2. Make ONE LLM call with all unique values + payer context + valid_programs anchor.
       The LLM groups semantically equivalent variants and assigns a single canonical
       per group, resolving brand-prefix noise vs. meaningful qualifiers correctly.
    3. Build a {original_name -> canonical_name} lookup and apply it to every row.

    Trigger conditions (either is sufficient):
    - ``valid_programs`` is empty (no client-supplied valid-values list), OR
    - the ``AARETE_DERIVED_PROGRAM`` column is entirely empty/null for all rows.

    Args:
        df: Dashboard DataFrame containing at minimum a ``PROGRAM`` column.
        valid_programs: ``Constants.VALID_AARETE_DERIVED_PROGRAMS`` list.

    Returns:
        DataFrame with ``AARETE_DERIVED_PROGRAM`` populated/overwritten with
        canonical names derived from the raw PROGRAM values.
    """
    df_copy = df.copy()

    if "AARETE_DERIVED_PROGRAM" not in df_copy.columns:
        df_copy["AARETE_DERIVED_PROGRAM"] = ""

    has_valid_programs = bool(valid_programs)
    adp_col_empty = _all_empty_in_list_column(df_copy["AARETE_DERIVED_PROGRAM"])

    if has_valid_programs and not adp_col_empty:
        logging.info(
            "Skipping canonical PROGRAM derivation: valid mapping exists and "
            "AARETE_DERIVED_PROGRAM is already populated."
        )
        return df_copy

    if "PROGRAM" not in df_copy.columns:
        logging.warning(
            "PROGRAM column not found in DataFrame; skipping canonical program derivation."
        )
        return df_copy

    # Step 1 — collect unique (payer_name, payer_state, raw_value) combinations
    unique_combos = _collect_unique_payer_value_combos(df_copy, "PROGRAM")
    if not unique_combos:
        logging.info(
            "No non-empty PROGRAM values found; skipping canonical program derivation."
        )
        return df_copy

    # Step 2 — single global LLM call returns {raw: canonical} dict
    canonical_map = prompt_calls.prompt_aarete_derived_program_canonical_batch(
        unique_combos, valid_programs
    )

    # Step 4 — apply mapping to each row
    df_copy["AARETE_DERIVED_PROGRAM"] = df_copy["PROGRAM"].apply(
        lambda val: _map_row_list_to_canonicals(val, canonical_map)
    )
    return df_copy


def add_aarete_derived_product_canonical(
    df: pd.DataFrame, valid_products: list[str]
) -> pd.DataFrame:
    """Derive canonical product names from the raw PRODUCT column.

    Architecture (single global LLM call):
    1. Collect all unique raw PRODUCT values across the entire Dashboard DataFrame.
    2. Make ONE LLM call with all unique values + payer context + valid_products anchor.
       The LLM groups semantically equivalent variants and assigns a single canonical
       per group, resolving brand-prefix noise vs. meaningful qualifiers correctly.
    3. Build a {original_name -> canonical_name} lookup and apply it to every row.

    Trigger conditions (either is sufficient):
    - ``valid_products`` is empty (no client-supplied valid-values list), OR
    - the ``AARETE_DERIVED_PRODUCT`` column is entirely empty/null for all rows.

    Args:
        df: Dashboard DataFrame containing at minimum a ``PRODUCT`` column.
        valid_products: ``Constants.VALID_AARETE_DERIVED_PRODUCTS`` list.

    Returns:
        DataFrame with ``AARETE_DERIVED_PRODUCT`` populated/overwritten with
        canonical names derived from the raw PRODUCT values.
    """
    df_copy = df.copy()

    if "AARETE_DERIVED_PRODUCT" not in df_copy.columns:
        df_copy["AARETE_DERIVED_PRODUCT"] = ""

    has_valid_products = bool(valid_products)
    adp_col_empty = _all_empty_in_list_column(df_copy["AARETE_DERIVED_PRODUCT"])

    if has_valid_products and not adp_col_empty:
        logging.info(
            "Skipping canonical PRODUCT derivation: valid mapping exists and "
            "AARETE_DERIVED_PRODUCT is already populated."
        )
        return df_copy

    if "PRODUCT" not in df_copy.columns:
        logging.warning(
            "PRODUCT column not found in DataFrame; skipping canonical product derivation."
        )
        return df_copy

    # Step 1 — collect unique (payer_name, payer_state, raw_value) combinations
    unique_combos = _collect_unique_payer_value_combos(df_copy, "PRODUCT")
    if not unique_combos:
        logging.info(
            "No non-empty PRODUCT values found; skipping canonical product derivation."
        )
        return df_copy

    # Step 2 — single global LLM call returns {raw: canonical} dict
    canonical_map = prompt_calls.prompt_aarete_derived_product_canonical_batch(
        unique_combos, valid_products
    )

    # Step 4 — apply mapping to each row
    df_copy["AARETE_DERIVED_PRODUCT"] = df_copy["PRODUCT"].apply(
        lambda val: _map_row_list_to_canonicals(val, canonical_map)
    )
    return df_copy