augment_sample_resources.py

import json
import os
import sys
import yaml
from dataclasses import asdict, dataclass, field
from decimal import Decimal
from pathlib import Path
from typing import Optional
from datetime import datetime, timezone

import pandas as pd

prvdr_info_file = "sample-data/PRVDR_HSTRY_POC.csv"
df = pd.read_csv(prvdr_info_file, dtype={"PRVDR_SK": str})

cond_sk_info_file = "sample-data/CLM_RLT_COND_SGNTR_MBR_POC.csv"
cond_sk_df = pd.read_csv(cond_sk_info_file, dtype={"CLM_RLT_COND_SGNTR_SK": str})

synth_df = None
synth_prvdr_file = "out/SYNTHETIC_PRVDR_HSTRY.csv"
if os.path.exists(synth_prvdr_file):
    try:
        synth_df = pd.read_csv(synth_prvdr_file, dtype={"PRVDR_SK": str})
    except Exception:
        pass


def lookup_provider_history(npi_num: str | int | None) -> dict | None:
    if not npi_num:
        return None
    npi_str = str(npi_num).strip()
    if not npi_str or npi_str == "None":
        return None
    matching_rows = df[df["PRVDR_SK"] == npi_str]
    if not matching_rows.empty:
        return json.loads(matching_rows.iloc[0].to_json())
    if synth_df is not None:
        matching_rows = synth_df[synth_df["PRVDR_SK"] == npi_str]
        if not matching_rows.empty:
            return json.loads(matching_rows.iloc[0].to_json())
    return None


cur_sample = sys.argv[1]
cur_sample_data = {}
with Path(cur_sample).open("r") as file:
    cur_sample_data = json.load(file)

if len(sys.argv) > 2:
    cur_sample_data["profileType"] = sys.argv[2]

is_pharmacy = cur_sample_data.get("resourceType") == "ExplanationOfBenefit-Pharmacy"

OTHR_PRVDR_MEANING = (
    "supervisor" if cur_sample_data.get("CLM_TYPE_CD") in ("1700", "2700") else "otheroperating"
)

careteam_header_columns = {
    "PRVDR_RFRG_PRVDR_NPI_NUM": "referring",
    "PRVDR_OTHR_PRVDR_NPI_NUM": OTHR_PRVDR_MEANING,
    "PRVDR_ATNDG_PRVDR_NPI_NUM": "attending",
    "PRVDR_OPRTG_PRVDR_NPI_NUM": "operating",
    "PRVDR_RNDRNG_PRVDR_NPI_NUM": "rendering",
    "CLM_PRSBNG_PRVDR_GNRC_ID_NUM": "prescribing",
    # This is purposely commented out to note we do not pull it in on the careteam (for now).
    "PRVDR_SRVC_PRVDR_NPI_NUM": "",
}

def load_profile_map():
    profile_map = {}
    paths = [
        Path("dictionary-support-files/ExplanationOfBenefit.yaml"),
        Path("dictionary-support-files/ExplanationOfBenefit-Pharmacy.yaml"),
        Path("dictionary-support-files/ExplanationOfBenefit-PriorAuth.yaml"),
    ]

    for p in filter(Path.exists, paths):
        with p.open("r") as f:
            for item in yaml.safe_load(f) or []:
                if (profiles := item.get("profiles")) is None:
                    continue
                keys = {item.get("sourceColumn"), item.get("inputPath", "").split(".")[-1]} - {"", None}
                profile_map.update(dict.fromkeys(keys, profiles))
    return profile_map


def filter_by_profile(data, profile_type, profile_map):
    if isinstance(data, dict):
        for k, v in list(data.items()):
            if k == "providerList":
                continue
            if k in profile_map and profile_type not in profile_map[k]:
                del data[k]
            else:
                filter_by_profile(v, profile_type, profile_map)
    elif isinstance(data, list):
        for item in data:
            filter_by_profile(item, profile_type, profile_map)


def cleanup_empty_items(data):
    if sic := data.get("supportingInfoComponents"):
        new_sic = []
        si_map = {}
        # Filter for non-empty items (items having more than just ROW_NUM)
        filtered_sic = [i for i in sic if any(k != "ROW_NUM" for k in i)]

        # Sort on avail columns for consistency.
        def get_sort_key(item):
            return sorted([(k, str(v)) for k, v in item.items() if k != "ROW_NUM"])

        sorted_sic = sorted(filtered_sic, key=get_sort_key)

        for next_row, item in enumerate(sorted_sic, start=1):
            if (old_row := item.get("ROW_NUM")) is not None:
                si_map.update({str(old_row): next_row, int(old_row): next_row})
            item["ROW_NUM"] = next_row
            new_sic.append(item)
        data["supportingInfoComponents"] = new_sic

        for li in data.get("lineItemComponents", []):
            if seq := li.get("SEQUENCE_INFO"):
                li["SEQUENCE_INFO"] = [si_map[int(s)] for s in seq if int(s) in si_map]

    if diag := data.get("diagnoses"):
        new_diag = []
        diag_map = {}
        internal_diag_keys = {"ROW_NUM", "clm_prod_type_cd_map"}
        # Filter for non-empty diagnoses
        for next_row, item in enumerate((i for i in diag if any(k not in internal_diag_keys for k in i)), start=1):
            old_row = item.get("ROW_NUM")
            item["ROW_NUM"] = str(next_row) if isinstance(old_row, str) else next_row
            if old_row is not None:
                diag_map.update({str(old_row): next_row, int(old_row): next_row})
            new_diag.append(item)
        data["diagnoses"] = new_diag

        # Update diagnosisSequence in lineItemComponents
        for li in data.get("lineItemComponents", []):
            if diag_seq := li.get("diagnosisSequence"):
                li["diagnosisSequence"] = [diag_map[int(s)] for s in diag_seq if int(s) in diag_map]



# we filter twice - once before augmentation and again after, to make it simpler.
profile_map = load_profile_map()
profile_type = cur_sample_data.get("profileType", "CMS")
profile_type = "CMS" if profile_type.lower() == "cms" else profile_type.capitalize()
cur_sample_data["profileType"] = profile_type #in case it's not specified, we rely upon this for mapping now.

filter_by_profile(cur_sample_data, profile_type, profile_map)


line_supporting_info_columns = [
    "CLM_LINE_PMD_UNIQ_TRKNG_NUM",
    "CLM_LINE_PA_UNIQ_TRKNG_NUM",
]
header_to_supp_info_cols = {
    "CLM_BNFT_ENHNCMT_1_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_BNFT_ENHNCMT_2_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_BNFT_ENHNCMT_3_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_BNFT_ENHNCMT_4_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_BNFT_ENHNCMT_5_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_NGACO_PBPMT_SW": "CLM_NGACO_PBPMT_SW",
    "CLM_NGACO_PDSCHRG_HCBS_SW": "CLM_NGACO_PDSCHRG_HCBS_SW",
    "CLM_NGACO_SNF_WVR_SW": "CLM_NGACO_SNF_WVR_SW",
    "CLM_NGACO_TLHLTH_SW": "CLM_NGACO_TLHLTH_SW",
    "CLM_NGACO_CPTATN_SW": "CLM_NGACO_CPTATN_SW",
    "CLM_ACO_CARE_MGMT_HCBS_SW": "CLM_ACO_CARE_MGMT_HCBS_SW",
    "CLM_PD_STUS_CD": "CLM_PD_STUS_CD",
    "CLM_NRLN_RIC_CD": "CLM_NRLN_RIC_CD",
    "CLM_RIC_CD": "CLM_RIC_CD",
    "GEO_BLG_SSA_STATE_CD": "GEO_BLG_SSA_STATE_CD",
    "CLM_MDCR_EXHSTD_DT": "CLM_MDCR_EXHSTD_DT",
    "BENE_PTNT_STUS_CD": "BENE_PTNT_STUS_CD",
    "CLM_ADMSN_SRC_CD": "CLM_ADMSN_SRC_CD",
    "CLM_ADMSN_TYPE_CD": "CLM_ADMSN_TYPE_CD",
    "CLM_FI_ACTN_CD": "CLM_FI_ACTN_CD",
    "CLM_HHA_LUP_IND_CD": "CLM_HHA_LUP_IND_CD",
    "CLM_HHA_RFRL_CD": "CLM_HHA_RFRL_CD",
    "CLM_MDCR_INSTNL_MCO_PD_SW": "CLM_MDCR_INSTNL_MCO_PD_SW",
    "CLM_MDCR_NCH_PTNT_STUS_IND_CD": "CLM_MDCR_NCH_PTNT_STUS_IND_CD",
    "CLM_MDCR_NPMT_RSN_CD": "CLM_MDCR_NPMT_RSN_CD",
    "CLM_OP_SRVC_TYPE_CD": "CLM_OP_SRVC_TYPE_CD",
    "CLM_PPS_IND_CD": "CLM_PPS_IND_CD",
    "DGNS_DRG_CD": "DGNS_DRG_CD",
    "DGNS_DRG_OUTLIER_CD": "DGNS_DRG_OUTLIER_CD",
    "CLM_ACTV_CARE_THRU_DT": "CLM_ACTV_CARE_THRU_DT",
    "CLM_CMS_PROC_DT": "CLM_CMS_PROC_DT",
    "CLM_NCH_WKLY_PROC_DT": "CLM_NCH_WKLY_PROC_DT",
    "CLM_NCVRD_FROM_DT": "CLM_NCVRD_FROM_DT",
    "CLM_NCVRD_THRU_DT": "CLM_NCVRD_THRU_DT",
    "CLM_QLFY_STAY_FROM_DT": "CLM_QLFY_STAY_FROM_DT",
    "CLM_QLFY_STAY_THRU_DT": "CLM_QLFY_STAY_THRU_DT",
    "CLM_SUBMSN_DT": "CLM_SUBMSN_DT",
    "CLM_ADJSTMT_TYPE_CD": "CLM_ADJSTMT_TYPE_CD",
    "CLM_BLOOD_PT_FRNSH_QTY": "CLM_BLOOD_PT_FRNSH_QTY",
    "CLM_CNTRCTR_NUM": "CLM_CNTRCTR_NUM",
    "CLM_DISP_CD": "CLM_DISP_CD",
    "CLM_IDR_LD_DT": "CLM_IDR_LD_DT",
    "CLM_NCH_PRMRY_PYR_CD": "CLM_NCH_PRMRY_PYR_CD",
    "CLM_QUERY_CD": "CLM_QUERY_CD",
    "CLM_CARR_PMT_DNL_CD": "CLM_CARR_PMT_DNL_CD",
    "CLM_CLNCL_TRIL_NUM": "CLM_CLNCL_TRIL_NUM",
    "CLM_MDCR_PRFNL_PRVDR_ASGNMT_SW": "CLM_MDCR_PRFNL_PRVDR_ASGNMT_SW",
    "CLM_SBMT_FRMT_CD": "CLM_SBMT_FRMT_CD",
    "CLM_SBMTR_CNTRCT_NUM": "CLM_SBMTR_CNTRCT_NUM",
    "CLM_SBMTR_CNTRCT_PBP_NUM": "CLM_SBMTR_CNTRCT_PBP_NUM",
}

line_to_supp_info_cols = {
    "CLM_LINE_BNFT_ENHNCMT_1_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_LINE_BNFT_ENHNCMT_2_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_LINE_BNFT_ENHNCMT_3_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_LINE_BNFT_ENHNCMT_4_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_LINE_BNFT_ENHNCMT_5_CD": "CLM_BNFT_ENHNCMT_CD",
    "CLM_LINE_NGACO_PBPMT_SW": "CLM_NGACO_PBPMT_SW",
    "CLM_LINE_NGACO_PDSCHRG_HCBS_SW": "CLM_NGACO_PDSCHRG_HCBS_SW",
    "CLM_LINE_NGACO_SNF_WVR_SW": "CLM_NGACO_SNF_WVR_SW",
    "CLM_LINE_NGACO_TLHLTH_SW": "CLM_NGACO_TLHLTH_SW",
    "CLM_LINE_NGACO_CPTATN_SW": "CLM_NGACO_CPTATN_SW",
    "CLM_LINE_ACO_CARE_MGMT_HCBS_SW": "CLM_ACO_CARE_MGMT_HCBS_SW",
    "GEO_RNDRG_SSA_STATE_CD": "GEO_RNDRG_SSA_STATE_CD",
    "CLM_BRND_GNRC_CD": "CLM_BRND_GNRC_CD",
    "CLM_CMPND_CD": "CLM_CMPND_CD",
    "CLM_CTSTRPHC_CVRG_IND_CD": "CLM_CTSTRPHC_CVRG_IND_CD",
    "CLM_DAW_PROD_SLCTN_CD": "CLM_DAW_PROD_SLCTN_CD",
    "CLM_DRUG_CVRG_STUS_CD": "CLM_DRUG_CVRG_STUS_CD",
    "CLM_DSPNSNG_STUS_CD": "CLM_DSPNSNG_STUS_CD",
    "CLM_LINE_AUTHRZD_FILL_NUM": "CLM_LINE_AUTHRZD_FILL_NUM",
    "CLM_LINE_DAYS_SUPLY_QTY": "CLM_LINE_DAYS_SUPLY_QTY",
    "CLM_LINE_RX_FILL_NUM": "CLM_LINE_RX_FILL_NUM",
    "CLM_LINE_RX_ORGN_CD": "CLM_LINE_RX_ORGN_CD",
    "CLM_LTC_DSPNSNG_MTHD_CD": "CLM_LTC_DSPNSNG_MTHD_CD",
    "CLM_PHRMCY_SRVC_TYPE_CD": "CLM_PHRMCY_SRVC_TYPE_CD",
    "CLM_PTNT_RSDNC_CD": "CLM_PTNT_RSDNC_CD",
    "CLM_LINE_RX_NUM": "CLM_LINE_RX_NUM",
}

npis_used = []
cur_sample_data["providerList"] = []
cur_careteam_sequence = 1
# we only use CLM_SRVC_PRVDR_GNRC_ID_NUM for part D events (we filter for PRVDR_SRVC_NPI_)
if cur_sample_data.get("CLM_TYPE_CD") not in (1, 2, 3, 4):
    billing_column = "PRVDR_BLG_PRVDR_NPI_NUM"
else:
    billing_column = "CLM_SRVC_PRVDR_GNRC_ID_NUM"

rx_line_financial_fields = [
    "CLM_LINE_INGRDNT_CST_AMT",
    "CLM_LINE_SRVC_CST_AMT",
    "CLM_LINE_SLS_TAX_AMT",
    "CLM_LINE_VCCN_ADMIN_FEE_AMT",
]


def convert_to_decimal(val: str | None) -> Decimal:
    try:
        return Decimal(val)
    except (TypeError, ValueError):
        return 0.0

@dataclass
class Provider:
    PRVDR_SK: Optional[str] = None
    PRVDR_ID_QLFYR_CD: Optional[str] = None
    NPI_TYPE: Optional[str] = None
    careTeamType: Optional[str] = None
    careTeamSequenceNumber: Optional[str] = None
    PRVDR_LAST_OR_LGL_NAME: Optional[str] = None
    PRVDR_1ST_NAME: Optional[str] = None
    PRVDR_CARETEAM_NAME: Optional[str] = None
    specialtyCode: Optional[str] = None
    PRVDR_OSCAR_NUM: Optional[str] = None
    CLM_BLG_PRVDR_ZIP5_CD: Optional[str] = None
    CLM_PRVDR_GNRC_ID_NUM: Optional[str] = None
    CLM_BLG_PRVDR_TAX_NUM: Optional[str] = None

populate_fields_except_na = [
    "PRVDR_LGL_NAME",
    "PRVDR_OSCAR_NUM",
    "PRVDR_LAST_NAME",
    "PRVDR_1ST_NAME",
    "PRVDR_MDL_NAME",
    "PRVDR_TYPE_CD",
]
provider_list = []

# we only use CLM_SRVC_PRVDR_GNRC_ID_NUM for part D events (we filter for PRVDR_SRVC_NPI_)
if cur_sample_data.get("CLM_TYPE_CD") not in (1, 2, 3, 4):
    billing_column = "PRVDR_BLG_PRVDR_NPI_NUM"
else:
    billing_column = "CLM_SRVC_PRVDR_GNRC_ID_NUM"


# There may be an opportunity to consolidate even the duplicate NPIs into a
# single careTeam reference, but we should wait to get feedback on this
# The reason being: it's possible to lose context on rendering vs ordering
def create_billing_and_service_provider(billing_col_name):
    provider_object = Provider(PRVDR_LAST_OR_LGL_NAME="N/A")

    qualifier = cur_sample_data.get("PRVDR_SRVC_ID_QLFYR_CD")
    if qualifier:
        provider_object.PRVDR_ID_QLFYR_CD = qualifier
        provider_object.CLM_PRVDR_GNRC_ID_NUM = cur_sample_data.get("CLM_SRVC_PRVDR_GNRC_ID_NUM")
    if qualifier in ("01", None):
        # Only pull NPI data if it's an NPI
        npi_num = cur_sample_data.get(billing_col_name)
        prvdr_hstry_for_npi = lookup_provider_history(npi_num)
        if prvdr_hstry_for_npi:
            provider_object.NPI_TYPE = "2" if prvdr_hstry_for_npi.get("PRVDR_LGL_NAME") else "1"
            provider_object.PRVDR_SK = npi_num
            provider_object.PRVDR_LAST_OR_LGL_NAME = (
                prvdr_hstry_for_npi["PRVDR_LGL_NAME"]
                if provider_object.NPI_TYPE == "2"
                else prvdr_hstry_for_npi["PRVDR_LAST_NAME"]
            )
            if prvdr_hstry_for_npi.get("PRVDR_1ST_NAME"):
                provider_object.PRVDR_1ST_NAME = prvdr_hstry_for_npi.get("PRVDR_1ST_NAME")

        if cur_sample_data.get("CLM_BLG_PRVDR_OSCAR_NUM"):
            provider_object.PRVDR_OSCAR_NUM = cur_sample_data.get("CLM_BLG_PRVDR_OSCAR_NUM")
        if cur_sample_data.get("CLM_BLG_PRVDR_TAX_NUM"):
            provider_object.CLM_BLG_PRVDR_TAX_NUM = cur_sample_data.get("CLM_BLG_PRVDR_TAX_NUM")
        if cur_sample_data.get("CLM_BLG_PRVDR_ZIP5_CD"):
            provider_object.CLM_BLG_PRVDR_ZIP5_CD = cur_sample_data.get("CLM_BLG_PRVDR_ZIP5_CD")

        if qualifier == "01":
            provider_object.PRVDR_SK = cur_sample_data.get("CLM_SRVC_PRVDR_GNRC_ID_NUM")
        # set other fields according to NPPES
    else:
        # only relevant to PDE.
        provider_object.CLM_PRVDR_GNRC_ID_NUM = cur_sample_data["CLM_SRVC_PRVDR_GNRC_ID_NUM"]
    return provider_object

provider_list.append(create_billing_and_service_provider(billing_column))

# For careteam elements, we don't need OSCAR number, TAX num (it's line-item), etc.
# We DO care about qualifiers + taxonomy, so we have a separate method.
def create_careteam_provider(careteam_column):
    provider_object = Provider(PRVDR_CARETEAM_NAME="N/A")
    qualifier = cur_sample_data.get("PRVDR_PRSBNG_ID_QLFYR_CD")
    if qualifier:
        provider_object.PRVDR_ID_QLFYR_CD = qualifier
        provider_object.CLM_PRVDR_GNRC_ID_NUM = cur_sample_data.get("CLM_PRSBNG_PRVDR_GNRC_ID_NUM")
    if qualifier in ("01", None):
        # Only pull NPI data if it's an NPI
        npi_num = cur_sample_data.get(careteam_column)
        prvdr_hstry_for_npi = lookup_provider_history(npi_num)
        if prvdr_hstry_for_npi:
            provider_object.NPI_TYPE = "2" if prvdr_hstry_for_npi.get("PRVDR_LGL_NAME") else "1"
            provider_object.PRVDR_SK = npi_num
            # set a default name using PRVDR_HSTRY if not available.
            provider_object.PRVDR_CARETEAM_NAME = (
                prvdr_hstry_for_npi["PRVDR_LGL_NAME"]
                if provider_object.NPI_TYPE == "2"
                else f"{prvdr_hstry_for_npi.get('PRVDR_LAST_NAME') or ''}, {prvdr_hstry_for_npi.get('PRVDR_1ST_NAME') or ''}".strip(", ")
            )
        else:
            provider_object.NPI_TYPE = "1"
            provider_object.PRVDR_SK = npi_num
            provider_object.PRVDR_CARETEAM_NAME = "Practitioner"

        provider_object.careTeamType = careteam_header_columns.get(careteam_column)
        if qualifier:
            provider_object.PRVDR_SK = cur_sample_data.get("CLM_PRSBNG_PRVDR_GNRC_ID_NUM")
        else:  # taxonomy codes only present in non-PDE
            cur_prov_type = careteam_column.split("_")[1]
            cur_specialty_code_col = "CLM_" + cur_prov_type + "_FED_PRVDR_SPCLTY_CD"
            if cur_sample_data.get(cur_specialty_code_col):
                provider_object.specialtyCode = cur_sample_data.get(cur_specialty_code_col)
            cur_prvdr_name_col = "CLM_" + cur_prov_type + "_PRVDR_NAME"
            if cur_sample_data.get(cur_prvdr_name_col):
                provider_object.PRVDR_CARETEAM_NAME = cur_sample_data.get(cur_prvdr_name_col)
    else:
        provider_object.CLM_PRVDR_GNRC_ID_NUM = cur_sample_data["CLM_PRSBNG_PRVDR_GNRC_ID_NUM"]
    return provider_object


for careteam_column in careteam_header_columns:
    if careteam_column not in cur_sample_data:
        continue
    provider_object = create_careteam_provider(careteam_column)
    provider_object.careTeamSequenceNumber = sum(
        1 for item in provider_list if provider_object.careTeamType
    )
    provider_list.append(provider_object)


# Distinct method since line items end up cleaner
def create_rendering_line_provider(npi_num):
    provider_object = Provider(
        PRVDR_SK=npi_num, careTeamType="rendering", PRVDR_CARETEAM_NAME="N/A"
    )
    prvdr_hstry_for_npi = lookup_provider_history(npi_num)
    if prvdr_hstry_for_npi:
        provider_object.NPI_TYPE = "2" if prvdr_hstry_for_npi.get("PRVDR_LGL_NAME") else "1"
        provider_object.PRVDR_CARETEAM_NAME = (
            prvdr_hstry_for_npi["PRVDR_LGL_NAME"]
            if provider_object.NPI_TYPE == "2"
            else f"{prvdr_hstry_for_npi.get('PRVDR_LAST_NAME') or ''}, {prvdr_hstry_for_npi.get('PRVDR_1ST_NAME') or ''}".strip(", ")
        )
    return provider_object

# now we go through the line items!
for line_item in cur_sample_data.get("lineItemComponents", []):
    # we only care about PRVDR_RNDRNG_PRVDR_NPI_NUM
    cur_rendering_providers = [
        x.PRVDR_SK for x in provider_list if getattr(x, "careTeamType", None) == "rendering"
    ]
    prvdr_npi_num = line_item.get("PRVDR_RNDRNG_PRVDR_NPI_NUM")
    if prvdr_npi_num and prvdr_npi_num not in cur_rendering_providers:
        # add rendering provider
        provider = create_rendering_line_provider(prvdr_npi_num)
        if line_item.get("CLM_RNDRG_FED_PRVDR_SPCLTY_CD"):
            provider.specialtyCode = line_item.get("CLM_RNDRG_FED_PRVDR_SPCLTY_CD")
        provider.careTeamSequenceNumber = sum(
            1 for item in provider_list if provider_object.careTeamType
        )
        line_item["careTeamSequence"] = provider.careTeamSequenceNumber
        provider_list.append(provider)

    if prvdr_npi_num:
        sequence = next(
            x.careTeamSequenceNumber
            for x in provider_list
            if getattr(x, "careTeamType", None) == "rendering"
            and getattr(x, "PRVDR_SK", None) == prvdr_npi_num
        )
        line_item["careTeamSequence"] = [sequence]

supporting_info_seq = 1
if "supportingInfoComponents" not in cur_sample_data:
    cur_sample_data["supportingInfoComponents"] = []
supporting_info_components = cur_sample_data["supportingInfoComponents"]

cond_sk = cur_sample_data.get("CLM_RLT_COND_SGNTR_SK")
if cond_sk:
    matching_rows = cond_sk_df[cond_sk_df["CLM_RLT_COND_SGNTR_SK"] == str(cond_sk)]
    for _, row in matching_rows.iterrows():
        cond_cd = row.get("CLM_RLT_COND_CD")
        if pd.notna(cond_cd) and str(cond_cd) != "~":
            supporting_info_components.append({"CLM_RLT_COND_CD": str(cond_cd)})

fac_type = cur_sample_data.get("CLM_BILL_FAC_TYPE_CD")
clsfctn = cur_sample_data.get("CLM_BILL_CLSFCTN_CD")
freq = cur_sample_data.get("CLM_BILL_FREQ_CD")

if fac_type and clsfctn and freq:
    supporting_info_components.append(
        {"TYPE_OF_BILL_CD": "0" + str(fac_type) + str(clsfctn) + str(freq)}
    )

for si_comp in supporting_info_components:
    si_comp["ROW_NUM"] = supporting_info_seq
    supporting_info_seq += 1

for source_col, target_col in header_to_supp_info_cols.items():
    value = cur_sample_data.get(source_col)
    if value is None and "institutionalComponents" in cur_sample_data:
        value = cur_sample_data["institutionalComponents"].get(source_col)
    if value is None and "profComponents" in cur_sample_data:
        value = cur_sample_data["profComponents"].get(source_col)

    if value:
        temp_var = {"ROW_NUM": supporting_info_seq, target_col: value}
        supporting_info_components.append(temp_var)
        supporting_info_seq += 1

# special case - want these in a single supportinginfo component
actv_from = cur_sample_data.get("CLM_ACTV_CARE_FROM_DT")

dschrg_dt = cur_sample_data.get("CLM_DSCHRG_DT")

if actv_from or dschrg_dt:
    temp_var = {"ROW_NUM": supporting_info_seq}
    if actv_from:
        temp_var["CLM_ACTV_CARE_FROM_DT"] = actv_from
    if dschrg_dt:
        temp_var["CLM_DSCHRG_DT"] = dschrg_dt
    supporting_info_components.append(temp_var)
    supporting_info_seq += 1


# There can be line item NPIs that are not present at header level, but
# need to be added to the CareTeam. This populates those.
line_items = cur_sample_data.get("lineItemComponents", [])
for item in line_items:
    if "SEQUENCE_INFO" in item and not isinstance(item["SEQUENCE_INFO"], list):
        item["SEQUENCE_INFO"] = [item["SEQUENCE_INFO"]]

    for line_supporting_info_col in line_supporting_info_columns:
        if item.get(line_supporting_info_col):
            item.setdefault("SEQUENCE_INFO", []).append(supporting_info_seq)
            supporting_info_seq += 1

    for source_col, target_col in line_to_supp_info_cols.items():
        if source_col == "CLM_LINE_RX_NUM" and not is_pharmacy:
            continue
        value = item.get(source_col)
        if value:
            temp_var = {"ROW_NUM": supporting_info_seq, target_col: value}
            supporting_info_components.append(temp_var)
            item.setdefault("SEQUENCE_INFO", []).append(supporting_info_seq)
            supporting_info_seq += 1

    # special case - want these in a single supportinginfo component. See EOB-base map.
    hct_hgb_rslt = item.get("CLM_LINE_HCT_HGB_RSLT_NUM")
    hct_hgb_type = item.get("CLM_LINE_HCT_HGB_TYPE_CD")
    clncl_lab_num = item.get("CLM_LINE_CARR_CLNCL_LAB_NUM")

    if hct_hgb_rslt or hct_hgb_type or clncl_lab_num:
        temp_var = {"ROW_NUM": supporting_info_seq}
        if hct_hgb_rslt:
            temp_var["CLM_LINE_HCT_HGB_RSLT_NUM"] = hct_hgb_rslt
        if hct_hgb_type:
            temp_var["CLM_LINE_HCT_HGB_TYPE_CD"] = hct_hgb_type
        if clncl_lab_num:
            temp_var["CLM_LINE_CARR_CLNCL_LAB_NUM"] = clncl_lab_num
        supporting_info_components.append(temp_var)
        item.setdefault("SEQUENCE_INFO", []).append(supporting_info_seq)
        supporting_info_seq += 1

    # for part D claims, sum CLM_LINE_INGRDNT_CST_AMT, CLM_LINE_SRVC_CST_AMT, CLM_LINE_SLS_TAX_AMT,
    # CLM_LINE_VCCN_ADMIN_FEE_AMT to set TOT_RX_CST_AMT
    tot_rx_amt = sum(
        convert_to_decimal(item.get(financial_field))
        for financial_field in rx_line_financial_fields
    )
    if tot_rx_amt > 0.0:
        item["TOT_RX_CST_AMT"] = str(tot_rx_amt)

cur_sample_data["providerList"] = [
    {k: v for k, v in asdict(p).items() if v is not None} for p in provider_list
]

# diagnoses section
@dataclass
class Diagnosis:
    CLM_DGNS_CD: str
    CLM_PROD_TYPE_CD: str = 'D'
    CLM_POA_IND: str = '~'
    CLM_DGNS_PRCDR_ICD_IND: str = '0'
    ROW_NUM: str = '1'
    clm_prod_type_cd_map: list[str] = field(default_factory=list)

diagnosis_codes = [
    Diagnosis(
        CLM_DGNS_CD=x.get("CLM_DGNS_CD"),
        CLM_PROD_TYPE_CD=x.get("CLM_PROD_TYPE_CD"),
        CLM_POA_IND=x.get("CLM_POA_IND"),
        CLM_DGNS_PRCDR_ICD_IND=x.get("CLM_DGNS_PRCDR_ICD_IND"),
        ROW_NUM=x.get("CLM_VAL_SQNC_NUM"),
    )
    for x in cur_sample_data.get("diagnoses", [])
    if x.get("CLM_PROD_TYPE_CD") in ["D", "P", "A", "R", "E"]
]

# Sort diagnoses keys
type_priority = {"P": 1, "A": 2, "R": 3, "E": 4}

# We need to preserve the list but sort it
diagnosis_codes.sort(key=lambda d: type_priority.get(d.CLM_PROD_TYPE_CD, 99))

# Assign sequential ROW_NUM
for idx, diag in enumerate(diagnosis_codes, start=1):
    diag.ROW_NUM = str(idx)

cur_sample_data["diagnoses"] = [asdict(d) for d in diagnosis_codes]


# Resolve claim status code section
def meta_src_prefix(meta_src_sk: str | None) -> str:
    return {
        "1002": "V",
        "1001": "M",
        "1003": "F",
    }.get(str(meta_src_sk), "")


def build_claim_audit_trail_composite(sample: dict) -> str:
    meta_src_sk = str(sample.get("META_SRC_SK", ""))
    status = sample.get("CLM_AUDT_TRL_STUS_CD")

    if not status:
        return None

    prefix = meta_src_prefix(meta_src_sk)

    # VMS
    if meta_src_sk == "1002":
        location = sample.get("CLM_AUDT_TRL_LCTN_CD", "")
        return f"{prefix}{status}{location}"
    # MCS & FISS
    return f"{prefix}{status}"


def next_row_num(supporting_info):
    row_nums = [
        int(si["ROW_NUM"])
        for si in supporting_info
        if si.get("ROW_NUM") is not None and str(si["ROW_NUM"]).isdigit()
    ]
    return str(max(row_nums, default=0) + 1)


composite_status_code = build_claim_audit_trail_composite(cur_sample_data)
if composite_status_code:
    supporting_info = cur_sample_data.get("supportingInfoComponents", [])
    supporting_info.append(
        {
            "ROW_NUM": next_row_num(supporting_info),
            "CLM_AUDT_TRL_STUS_CD": build_claim_audit_trail_composite(cur_sample_data),
        }
    )

# add diagnosisSequence where necessary
for item in cur_sample_data.get("lineItemComponents", []):
    if "CLM_LINE_DGNS_CD" in item:
        line_dgns_cd = item.get("CLM_LINE_DGNS_CD")
        if match := next(
            (
                d
                for d in diagnosis_codes
                if line_dgns_cd == d.CLM_DGNS_CD and d.CLM_PROD_TYPE_CD == 'D'
            ),
            None,
        ):
            item["diagnosisSequence"] = [int(match.ROW_NUM)]

filename = "out/temporary-sample.json"

cur_sample_data["lastUpdated"] = "2026-01-01T03:02:28.000000Z"

profile_type = cur_sample_data.get("profileType", "CMS")
profile_type = "CMS" if profile_type.lower() == "cms" else profile_type.capitalize()

filter_by_profile(cur_sample_data, profile_type, profile_map)
cleanup_empty_items(cur_sample_data)

with Path(filename).open("w") as f:
    json.dump(cur_sample_data, f, indent=4)

print("Temporary augmented file created")