import pandas as pd
import numpy as np
import logging
import math 
import re 
from typing import Optional
from .general import DataProcessorGeneral

logging.basicConfig(
    level = logging.INFO,                                 
    format = '%(asctime)s - %(levelname)s - %(message)s'  
)

class DataProcessorProstate(DataProcessorGeneral):

    GROUP_STAGE_MAPPING = {        
        # Stage IV
        'IV': 'IV',
        'IVA': 'IV',
        'IVB': 'IV',

        # Stage III
        'III': 'III',
        'IIIA': 'III',
        'IIIB': 'III',
        'IIIC': 'III',

        # Stage II
        'II': 'II',
        'IIA': 'II',
        'IIB': 'II',
        'IIC': 'II',

        # Stage I
        'I': 'I',
        
        # Unknown
        'Unknown / Not documented': 'unknown'
    }

    T_STAGE_MAPPING = {
        'T4': 'T4',
        'T3': 'T3',
        'T3a': 'T3',
        'T3b': 'T3',
        'T2': 'T2',
        'T2a': 'T2',
        'T2b': 'T2',
        'T2c': 'T2',
        'T1': 'T1',
        'T1a': 'T1',
        'T1b': 'T1',
        'T1c': 'T1',
        'T0': 'T1',
        'TX': 'unknown',
        'Unknown / Not documented': 'unknown'
    }

    N_STAGE_MAPPING = {
        'N1': 'N1',
        'N0': 'N0',
        'NX': 'unknown',
        'Unknown / Not documented': 'unknown'
    }

    M_STAGE_MAPPING = {
        'M1': 'M1',
        'M1a': 'M1',
        'M1b': 'M1',
        'M1c': 'M1',
        'M0': 'M0',
        'Unknown / Not documented': 'unknown'
    }

    GLEASON_MAPPING = {
        '10': 5,
        '9': 5,
        '8': 4,  
        '4 + 3 = 7': 3,
        '7 (when breakdown not available)': 3,
        '3 + 4 = 7': 2, 
        'Less than or equal to 6': 1,  
        'Unknown / Not documented': 'unknown'
    }

    LOINC_MAPPINGS = {
        'hemoglobin': ['718-7', '20509-6'],
        'wbc': ['26464-8', '6690-2'],
        'platelet': ['26515-7', '777-3', '778-1'],
        'creatinine': ['2160-0', '38483-4'],
        'bun': ['3094-0'],
        'sodium': ['2947-0', '2951-2'],
        'bicarbonate': ['1963-8', '1959-6', '14627-4', '1960-4', '2028-9'],
        'chloride': ['2075-0'],
        'potassium': ['6298-4', '2823-3'],
        'albumin': ['1751-7', '35706-1', '13980-8'],
        'calcium': ['17861-6', '49765-1'],
        'total_bilirubin': ['42719-5', '1975-2'],
        'ast': ['1920-8', '30239-8'],
        'alt': ['1742-6', '1743-4', '1744-2'],
        'alp': ['6768-6'],
        'psa': ['2857-1', '35741-8']
    }

    ICD_9_EXLIXHAUSER_MAPPING = {
        # Congestive heart failure
        r'^39891|^40201|^40211|^40291|^40401|^40403|^40411|^40413|^40491|^40493|^4254|^4255|^4256|^4257|^4258|^4259|^428': 'chf',
        
        # Cardiac arrhythmias
        r'^4260|^42613|^4267|^4269|^42610|^42612|^4270|^4271|^4272|^4273|^4274|^4276|^4277|^4278|^4279|^7850|^99601|^99604|^V450|^V533': 'cardiac_arrhythmias',
        
        # Valvular disease
        r'^0932|^394|^395|^396|^397|^424|^7463|^7464|^7465|^7466|^V422|^V433': 'valvular_disease',
        
        # Pulmonary circulation disorders
        r'^4150|^4151|^416|^4170|^4178|^4179': 'pulm_circulation',
        
        # Peripheral vascular disorders
        r'^0930|^4373|^440|^441|^4431|^4432|^4433|^4434|^4435|^4436|^4437|^4438|^4439|^4471|^5571|^5579|^V434': 'pvd',
        
        # Hypertension, uncomplicated
        r'^401': 'htn_uncomplicated',
        
        # Hypertension, complicated
        r'^402|^403|^404|^405': 'htn_complicated',
        
        # Paralysis
        r'^3341|^342|^343|^3440|^3441|^3442|^3443|^3444|^3445|^3446|^3449': 'paralysis',
        
        # Other neurological disorders
        r'^3319|^3320|^3321|^3334|^3335|^33392|^334|^335|^3362|^340|^341|^345|^3481|^3483|^7803|^7843': 'other_neuro',
        
        # Chronic pulmonary disease
        r'^4168|^4169|^490|^491|^492|^493|^494|^495|^496|^497|^498|^499|^500|^501|^502|^503|^504|^505|^5064|^5081|^5088': 'chronic_pulm_disease',
        
        # Diabetes, uncomplicated
        r'^2500|^2501|^2502|^2503': 'diabetes_uncomplicated',
        
        # Diabetes, complicated
        r'^2504|^2505|^2506|^2507|^2508|^2509': 'diabetes_complicated',
        
        # Hypothyroidism
        r'^2409|^243|^244|^2461|^2468': 'hypothyroid',
        
        # Renal failure
        r'^40301|^40311|^40391|^40402|^40403|^40412|^40413|^40492|^40493|^585|^586|^5880|^V420|^V451|^V56': 'renal_failure',
        
        # Liver disease
        r'^07022|^07023|^07032|^07033|^07044|^07054|^0706|^0709|^4560|^4561|^4562|^570|^571|^5722|^5723|^5724|^5725|^5726|^5727|^5728|^5733|^5734|^5738|^5739|^V427': 'liver_disease',
        
        # Peptic ulcer disease excluding bleeding
        r'^5317|^5319|^5327|^5329|^5337|^5339|^5347|^5349': 'pud',
        
        # AIDS/HIV
        r'^042|^043|^044': 'aids_hiv',
        
        # Lymphoma
        r'^200|^201|^202|^2030|^2386': 'lymphoma',
        
        # Rheumatoid arthritis/collagen vascular diseases
        r'^446|^7010|^7100|^7101|^7102|^7103|^7104|^7108|^7109|^7112|^714|^7193|^720|^725|^7285|^72889|^72930': 'rheumatic',
        
        # Coagulopathy
        r'^286|^2871|^2873|^2874|^2875': 'coagulopathy',
        
        # Obesity
        r'^2780': 'obesity',
        
        # Weight loss
        r'^260|^261|^262|^263|^7832|^7994': 'weight_loss',
        
        # Fluid and electrolyte disorders
        r'^2536|^276': 'fluid',
        
        # Blood loss anemia
        r'^2800': 'blood_loss_anemia',
        
        # Deficiency anemia
        r'^2801|^2802|^2803|^2804|^2805|^2806|^2807|^2808|^2809|^281': 'deficiency_anemia',
        
        # Alcohol abuse
        r'^2652|^2911|^2912|^2913|^2915|^2916|^2917|^2918|^2919|^3030|^3039|^3050|^3575|^4255|^5353|^5710|^5711|^5712|^5713|^980|^V113': 'alcohol_abuse',
        
        # Drug abuse
        r'^292|^304|^3052|^3053|^3054|^3055|^3056|^3057|^3058|^3059|^V6542': 'drug_abuse',
        
        # Psychoses
        r'^2938|^295|^29604|^29614|^29644|^29654|^297|^298': 'psychoses',
        
        # Depression
        r'^2962|^2963|^2965|^3004|^309|^311': 'depression'
    }

    ICD_10_ELIXHAUSER_MAPPING = {
        # Congestive heart failure
        r'^I099|^I110|^I130|^I132|^I255|^I420|^I425|^I426|^I427|^I428|^I429|^I43|^I50|^P290': 'chf',
        
        # Cardiac arrhythmias
        r'^I441|^I442|^I443|^I456|^I459|^I47|^I48|^I49|^R000|^R001|^R008|^T821|^Z450|^Z950': 'cardiac_arrhythmias',
        
        # Valvular disease
        r'^A520|^I05|^I06|^I07|^I08|^I091|^I098|^I34|^I35|^I36|^I37|^I38|^I39|^Q230|^Q231|^Q232|^Q233|^Z952|^Z953|^Z954': 'valvular_disease',
        
        # Pulmonary circulation disorders
        r'^I26|^I27|^I280|^I288|^I289': 'pulm_circulation',
        
        # Peripheral vascular disorders
        r'^I70|^I71|^I731|^I738|^I739|^I771|^I790|^I792|^K551|^K558|^K559|^Z958|^Z959': 'pvd',
        
        # Hypertension, uncomplicated
        r'^I10': 'htn_uncomplicated',
        
        # Hypertension, complicated
        r'^I11|^I12|^I13|^I15': 'htn_complicated',
        
        # Paralysis
        r'^G041|^G114|^G801|^G802|^G81|^G82|^G830|^G831|^G832|^G833|^G834|^G839': 'paralysis',
        
        # Other neurological disorders
        r'^G10|^G11|^G12|^G13|^G20|^G21|^G22|^G254|^G255|^G312|^G318|^G319|^G32|^G35|^G36|^G37|^G40|^G41|^G931|^G934|^R470|^R56': 'other_neuro',
        
        # Chronic pulmonary disease
        r'^I278|^I279|^J40|^J41|^J42|^J43|^J44|^J45|^J46|^J47|^J60|^J61|^J62|^J63|^J64|^J65|^J66|^J67|^J684|^J701|^J703': 'chronic_pulm_disease',
        
        # Diabetes, uncomplicated
        r'^E100|^E101|^E109|^E110|^E111|^E119|^E120|^E121|^E129|^E130|^E131|^E139|^E140|^E141|^E149': 'diabetes_uncomplicated',
        
        # Diabetes, complicated
        r'^E102|^E103|^E104|^E105|^E106|^E107|^E108|^E112|^E113|^E114|^E115|^E116|^E117|^E118|^E122|^E123|^E124|^E125|^E126|^E127|^E128|^E132|^E133|^E134|^E135|^E136|^E137|^E138|^E142|^E143|^E144|^E145|^E146|^E147|^E148': 'diabetes_complicated',
        
        # Hypothyroidism
        r'^E00|^E01|^E02|^E03|^E890': 'hypothyroid',
        
        # Renal failure
        r'^I120|^I131|^N18|^N19|^N250|^Z490|^Z491|^Z492|^Z940|^Z992': 'renal_failure',
        
        # Liver disease
        r'^B18|^I85|^I864|^I982|^K70|^K711|^K713|^K714|^K715|^K717|^K72|^K73|^K74|^K760|^K762|^K763|^K764|^K765|^K766|^K767|^K768|^K769|^Z944': 'liver_disease',
        
        # Peptic ulcer disease excluding bleeding
        r'^K257|^K259|^K267|^K269|^K277|^K279|^K287|^K289': 'pud',
        
        # AIDS/HIV
        r'^B20|^B21|^B22|^B24': 'aids_hiv',
        
        # Lymphoma
        r'^C81|^C82|^C83|^C84|^C85|^C88|^C96|^C900|^C902': 'lymphoma',
        
        # Rheumatoid arthritis/collagen vascular diseases
        r'^L940|^L941|^L943|^M05|^M06|^M08|^M120|^M123|^M30|^M310|^M311|^M312|^M313|^M32|^M33|^M34|^M35|^M45|^M461|^M468|^M469': 'rheumatic',
        
        # Coagulopathy
        r'^D65|^D66|^D67|^D68|^D691|^D693|^D694|^D695|^D696': 'coagulopathy',
        
        # Obesity
        r'^E66': 'obesity',
        
        # Weight loss
        r'^E40|^E41|^E42|^E43|^E44|^E45|^E46|^R634|^R64': 'weight_loss',
        
        # Fluid and electrolyte disorders
        r'^E222|^E86|^E87': 'fluid',
        
        # Blood loss anemia
        r'^D500': 'blood_loss_anemia',
        
        # Deficiency anemia
        r'^D508|^D509|^D51|^D52|^D53': 'deficiency_anemia',
        
        # Alcohol abuse
        r'^F10|^E52|^G621|^I426|^K292|^K700|^K703|^K709|^T51|^Z502|^Z714|^Z721': 'alcohol_abuse',
        
        # Drug abuse
        r'^F11|^F12|^F13|^F14|^F15|^F16|^F18|^F19|^Z715|^Z722': 'drug_abuse',
        
        # Psychoses
        r'^F20|^F22|^F23|^F24|^F25|^F28|^F29|^F302|^F312|^F315': 'psychoses',
        
        # Depression
        r'^F204|^F313|^F314|^F315|^F32|^F33|^F341|^F412|^F432': 'depression'
    }
        
    VAN_WALRAVEN_WEIGHTS = {
        'chf': 7,
        'cardiac_arrhythmias': 5,
        'valvular_disease': -1,
        'pulm_circulation': 4,
        'pvd': 2,
        'htn_uncomplicated': 0,
        'htn_complicated': 0,
        'paralysis': 7,
        'other_neuro': 6,
        'chronic_pulm_disease': 3,
        'diabetes_uncomplicated': 0,
        'diabetes_complicated': 0,
        'hypothyroid': 0,
        'renal_failure': 5,
        'liver_disease': 11,
        'pud': 0,
        'aids_hiv': 0,
        'lymphoma': 9,
        'rheumatic': 0,
        'coagulopathy': 3,
        'obesity': -4,
        'weight_loss': 6,
        'fluid': 5,
        'blood_loss_anemia': -2,
        'deficiency_anemia': -2,
        'alcohol_abuse': 0,
        'drug_abuse': -7,
        'psychoses': 0,
        'depression': -3
    }

    ICD_9_METS_MAPPING = {
        # Lymph nodes
        r'^196': 'lymph_met',
        
        # Thoracic
        r'^1970|^1971|^1972|^1973': 'thoracic_met',

        # Liver
        r'^1977': 'liver_met',
        
        # Bone
        r'^1985': 'bone_met',
        
        # Brain/CNS
        r'^1983|^1984': 'brain_met',
        
        # Other GI tract
        r'^1974|^1975|^1976|^1978|^1987': 'other_gi_met',
        
        # Other sites
        r'^1980|^1981|^1982|^1986|^1988|^199': 'other_met'
    }

    ICD_10_METS_MAPPING = {
        # Lymph nodes
        r'^C77': 'lymph_met',

        # Thoracic
        r'^C780|^C781|^C782|^C783': 'thoracic_met',

        # Liver
        r'^C787': 'liver_met',

        # Bone
        r'^C795': 'bone_met',

        # Brain/CNS
        r'^C793|^C794': 'brain_met',

        # Other GI tract
        r'^C784|^C785|^C786|^C788|^C797': 'other_gi_met',

        # Other sites
        r'^C790|^C791|^C792|^C796|^C798|^C799|^C80': 'other_met'
    }

    def __init__(self):
        super().__init__() 

        # prostate-speicfic attributres 
        self.enhanced_df = None
        self.biomarkers_df = None
        self.labs_df = None 
        self.diagnosis_df = None
        self.mortality_df = None
        self.adt_df = None

    def process_enhanced(self,
                         file_path: str,
                         index_date_column: str = 'MetDiagnosisDate',
                         patient_ids: list = None,
                         index_date_df: pd.DataFrame = None,
                         primary_treatment_path: str = None, 

                         drop_stages: bool = True,
                         drop_dates: bool = True) -> Optional[pd.DataFrame]: 
        """
        Processes Enhanced_MetProstate.csv to standardize categories, consolidate staging information, and calculate time-based metrics between key clinical events.
        
        The index date is used to determine castrate-resistance status by that time and to calculate time from diagnosis to castrate resistance. 
        The default index date is 'MetDiagnosisDate.' For an alternative index date, provide an index_date_df and specify the index_date_column accordingly.
        
        To process only specific patients, either:
        1. Provide patient_ids when using the default index date ('MetDiagnosisDate')
        2. Include only the desired PatientIDs in index_date_df when using a custom index date
        
        Parameters
        ----------
        file_path : str
            Path to Enhanced_MetProstate.csv file
        index_date_column : str, default = 'MetDiagnosisDate'
            name of column for index date of interest 
        patient_ids : list, optional
            List of specific PatientIDs to process. If None, processes all patients
        index_date_df : pd.DataFrame, optional 
            DataFrame containing unique PatientIDs and their corresponding index dates. Only data for PatientIDs present in this DataFrame will be processed
        primary_treatment_path : str, optional 
            Path to Enhanced_MetPC_PrimaryTreatment.csv file which is used to calculate time from diagnosis to primary treatment 
        drop_stages : bool, default=True
            If True, drops original staging columns (GroupStage, TStage, and MStage) after creating modified versions
        drop_dates : bool, default=True
            If True, drops date columns (DiagnosisDate, MetDiagnosisDate, and CRPCDate) after calculating durations

        Returns
        -------
        pd.DataFrame or None 
            - PatientID : object
                unique patient identifier
            - GroupStage_mod : category
                consolidated overall staging (I-IV and unknown) at time of first diagnosis
            - TStage_mod : category
                consolidated tumor staging (T1-T4 and unknown) at time of first diagnosis
            - NStage_mod : category
                consolidated lymph node staging (N0, N1, and unknown) at time of first diagnosis
            - MStage_mod : category
                consolidated metastasis staging (M0, M1, and unknown) at time of first diagnosis
            - GleasonScore_mod : category
                consolidated Gleason scores into Grade Groups (1-5 and unknown) at time of first diagnosis 
            - Histology : category
                histology (adenocarcinoma and NOS) at time of initial diagnosis 
            - days_diagnosis_to_met : float
                days from first diagnosis to metastatic disease 
            - met_diagnosis_year : category
                year of metastatic diagnosis 
            - TreatmentType_mod : cateogry
                consolidation primary treatments (surgery, radiation, and other), if Enhanced_MetPC_PrimaryTreatment.csv file is included as parameter
            - days_diagnosis_to_primary_treatment : float 
                days from first diagnosis to primary treatment, if Enhanced_MetPC_PrimaryTreatment.csv file is included as parameter
            - IsCRPC_index : Int64
                binary (0/1) indicator for CRPC, determined by whether CRPC date is earlier than the index date 
            - days_diagnosis_to_crpc : float
                days from diagnosis to CRPC, calculated only when CRPC date is prior to index date (i.e., IsCRPC == 1)
            - PSADiagnosis : float, ng/mL
                PSA at time of first diagnosis
            - PSAMetDiagnosis : float
                PSA at time of metastatic diagnosis 
            - psa_doubling_diag_met : float, months
                PSA doubling time from first diagnosis to metastatic disease for those with both a PSA at time of first and metastatic diagnosis, calculated only when PSA was higher at metastatic diagnosis than at initial diagnosis 
            - psa_velocity_diag_met : float, ng/mL/month
                PSA velocity from first diagnosis to metastatic disease for those with both a PSA at time of first and metastatic diagnosis

            Original date columns (DiagnosisDate, MetDiagnosisDate, CRPCDate, and TreatmentDate) retained if drop_dates = False

        Notes
        -----
        Notable T-Stage consolidation decisions:
            - T0 is included in T1 
            - TX and Unknown/not documented are categorized as 'unknown' 

        Notable Gleanson score consolidation decisions: 
            - 7 (when breakdown not available) was placed into Grade Group 3

        If days_diagnosis_to_primary_treatment is < 0, the value is set to 0; otherwise, the value is left as is. 
        Cases where days_diagnosis_to_primary_treatment is < 0 occur when the date of surgery is set to the first of the year 
        and likely reflects uknown exact date of primary treatment. By setting these cases to 0, we'll presume that the 
        primary treatment date occurred on the diagnosis date. 

        PSA doubling time calculation:
            - Only calculated for patients with PSA values at both initial diagnosis and metastatic diagnosis
            - Only valid when PSA at metastatic diagnosis is higher than at initial diagnosis
            - Formula: (ln(2)/PSA slope), measured in months
            - Where PSA slope = [ln(PSAMetDiagnosis) - ln(PSADiagnosis)]/[(MetDiagnosisDate - DiagnosisDate) in months]

        PSA velocity calculation:
            - Only calculated for patients with PSA values at both initial diagnosis and metastatic diagnosis
            - Formula: (PSAMetDiagnosis - PSADiagnosis)/[(MetDiagnosisDate - DiagnosisDate) in months], measured in ng/mL/month

        Output handling: 
        - Duplicate PatientIDs are logged as warnings if found but retained in output
        - Processed DataFrame is stored in self.enhanced_df
        """
        # Input validation
        if patient_ids is not None:
            if not isinstance(patient_ids, list):
                raise TypeError("patient_ids must be a list or None")
        
        if index_date_df is not None:
            if not isinstance(index_date_df, pd.DataFrame):
                raise ValueError("index_date_df must be a pandas DataFrame")
            if 'PatientID' not in index_date_df.columns:
                raise ValueError("index_date_df must contain a 'PatientID' column")
            if not index_date_column or index_date_column not in index_date_df.columns:
                raise ValueError('index_date_column not found in index_date_df')
            if index_date_df['PatientID'].duplicated().any():
                raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
            
            index_date_df = index_date_df.copy()
            # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
            for col in index_date_df.columns:
                if col != 'PatientID':  # Keep PatientID unchanged for merging
                    index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

            # Update index_date_column name
            index_date_column = f'imported_{index_date_column}'

        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Enhanced_MetProstate.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            # Case 1: Using default MetDiagnosisDate with specific patients
            if index_date_column == 'MetDiagnosisDate' and patient_ids is not None:
                logging.info(f"Filtering for {len(patient_ids)} specific PatientIDs")
                df = df[df['PatientID'].isin(patient_ids)]
                logging.info(f"Successfully filtered Enhanced_MetProstate.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            # Case 2: Using custom index date with index_date_df
            elif index_date_column != 'MetDiagnosisDate' and index_date_df is not None:
                index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])
                df = df[df.PatientID.isin(index_date_df.PatientID)]
                df = pd.merge(
                    df,
                    index_date_df[['PatientID', index_date_column]],
                    on = 'PatientID',
                    how = 'left'
                )
                logging.info(f"Successfully merged Enhanced_MetProstate.csv df with index_date_df resulting in shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            # Case 3: Using default MetDiagnosisDate with all patients (no filtering)
            elif index_date_column == 'MetDiagnosisDate' and patient_ids is None:
                logging.info(f"No filtering applied. Using all {df['PatientID'].nunique()} patients in the dataset")

            # Case 4: Error case - custom index date without index_date_df
            else:
                logging.error("If index_date_column is not 'MetDiagnosisDate', an index_date_df must be provided")
                return None
        
            # Convert categorical columns
            categorical_cols = ['GroupStage',
                                'TStage', 
                                'NStage',
                                'MStage', 
                                'GleasonScore', 
                                'Histology']
            
            df[categorical_cols] = df[categorical_cols].astype('category')

            # Recode stage variables using class-level mapping and create new column
            df['GroupStage_mod'] = df['GroupStage'].map(self.GROUP_STAGE_MAPPING).astype('category')
            df['TStage_mod'] = df['TStage'].map(self.T_STAGE_MAPPING).astype('category')
            df['NStage_mod'] = df['NStage'].map(self.N_STAGE_MAPPING).astype('category')
            df['MStage_mod'] = df['MStage'].map(self.M_STAGE_MAPPING).astype('category')
            df['GleasonScore_mod'] = df['GleasonScore'].map(self.GLEASON_MAPPING).astype('category')

            # Drop original stage variables if specified
            if drop_stages:
                df = df.drop(columns=['GroupStage', 'TStage', 'NStage', 'MStage', 'GleasonScore'])

            # Convert date columns to datetime
            date_cols = ['DiagnosisDate', 'MetDiagnosisDate', 'CRPCDate']
            for col in date_cols:
                df[col] = pd.to_datetime(df[col])

            # Generate new time-based variables 
            df['days_diagnosis_to_met'] = (df['MetDiagnosisDate'] - df['DiagnosisDate']).dt.days
            df['met_diagnosis_year'] = pd.Categorical(df['MetDiagnosisDate'].dt.year)

            if primary_treatment_path is not None: 
                try: 
                    primary_treatment_df = pd.read_csv(primary_treatment_path)
                    logging.info(f"Successfully read Enhanced_MetPC_PrimaryTreatment.csv file with shape: {primary_treatment_df.shape} and unique PatientIDs: {(primary_treatment_df['PatientID'].nunique())}")

                    df = pd.merge(df,primary_treatment_df, on = 'PatientID', how = 'left')
                    df['TreatmentDate'] = pd.to_datetime(df['TreatmentDate'])
                    df['days_diagnosis_to_primary_treatment'] = (df['TreatmentDate'] - df['DiagnosisDate']).dt.days
            
                    # For those with value <0, set to 0; otherwise, leave as is
                    df['days_diagnosis_to_primary_treatment'] = np.where(df['days_diagnosis_to_primary_treatment'] < 0, 0, df['days_diagnosis_to_primary_treatment'])

                    df['TreatmentType_mod'] = np.where((df['TreatmentType'] == 'Cryotherapy') | (df['TreatmentType'] == 'High Intensity Focused Ultrasound (HIFU)'), 
                                                       'other',
                                                       df['TreatmentType'])
                    
                    df = df.drop(columns = ['TreatmentType'])

                    if drop_dates:
                        df = df.drop(columns = ['TreatmentDate'])

                except Exception as e:
                    logging.error(f"Error processing Enhanced_MetPC_PrimaryTreatment.csv: {e}") 

            # Generate IsCRPC_index where 1 if CRPCDate is less than or equal to index date 
            # Calculate time from diagnosis to CRPC (presuming before metastatic diagnosis or index)
            if index_date_column == "MetDiagnosisDate":
                df['IsCRPC_index'] = np.where(df['CRPCDate'] <= df['MetDiagnosisDate'], 1, 0)
                df['days_diagnosis_to_crpc'] = np.where(df['IsCRPC_index'] == 1,
                                                        (df['CRPCDate'] - df['DiagnosisDate']).dt.days,
                                                        np.nan)
            else:
                df['IsCRPC_index'] = np.where(df['CRPCDate'] <= df[index_date_column], 1, 0)
                df['days_diagnosis_to_crpc'] = np.where(df['IsCRPC_index'] == 1,
                                                        (df['CRPCDate'] - df['DiagnosisDate']).dt.days,
                                                        np.nan)
            
            df = df.drop(columns = ['IsCRPC'])

            num_cols = ['PSADiagnosis', 'PSAMetDiagnosis']
            for col in num_cols:
                df[col] = pd.to_numeric(df[col], errors = 'coerce').astype('float')

            # Calculating PSA doubling time in months 
            df_doubling = (
                df
                .query('DiagnosisDate.notna()')
                .query('days_diagnosis_to_met > 30') # At least 30 days from first diagnosis to metastatic diagnosis
                .query('PSADiagnosis.notna()')
                .query('PSAMetDiagnosis.notna()')
                .query('PSAMetDiagnosis > PSADiagnosis') # Doubling time formula only makes sense for rising numbers 
                .query('PSADiagnosis > 0') 
                .query('PSAMetDiagnosis > 0')  
                .assign(psa_doubling_diag_met = lambda x: 
                        ((x['days_diagnosis_to_met']/30) * math.log(2))/
                        (np.log(x['PSAMetDiagnosis']) - 
                        np.log(x['PSADiagnosis']))
                        )
                [['PatientID', 'psa_doubling_diag_met']]
            )

            # Calculating PSA velocity with time in months 
            df_velocity = (
                df
                .query('DiagnosisDate.notna()')
                .query('days_diagnosis_to_met > 30') # At least 30 days from first diagnosis to metastatic diagnosis
                .query('PSADiagnosis.notna()')
                .query('PSAMetDiagnosis.notna()')
                .query('PSADiagnosis > 0') 
                .query('PSAMetDiagnosis > 0') 
                .assign(psa_velocity_diag_met = lambda x: (x['PSAMetDiagnosis'] - x['PSADiagnosis']) / (x['days_diagnosis_to_met']/30))
                [['PatientID', 'psa_velocity_diag_met']]
            )

            final_df = pd.merge(df, df_doubling, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, df_velocity, on = 'PatientID', how = 'left')

            if drop_dates:
                final_df = final_df.drop(columns = ['MetDiagnosisDate', 'DiagnosisDate', 'CRPCDate'])

            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset = ['PatientID'], keep = False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Enhanced_MetProstate.csv file with final shape: {final_df.shape} and unique PatientIDs: {(final_df['PatientID'].nunique())}")
            self.enhanced_df = final_df
            return final_df

        except Exception as e:
            logging.error(f"Error processing Enhanced_MetProstate.csv file: {e}")
            return None 
        
    def process_biomarkers(self,
                           file_path: str,
                           index_date_df: pd.DataFrame,
                           index_date_column: str, 
                           days_before: Optional[int] = None,
                           days_after: int = 0) -> Optional[pd.DataFrame]:
        """
        Processes Enhanced_MetPC_Biomarkers.csv by determining biomarker status for each patient within a specified time window relative to an index date. 

        Parameters
        ----------
        file_path : str
            Path to Enhanced_MetPC_Biomarkers.csv file
        index_date_df : pd.DataFrame
            DataFrame containing unique PatientIDs and their corresponding index dates. Only biomarker data for PatientIDs present in this DataFrame will be processed
        index_date_column : str
            Column name in index_date_df containing the index date
        days_before : int | None, optional
            Number of days before the index date to include. Must be >= 0 or None. If None, includes all prior results. Default: None
        days_after : int, optional
            Number of days after the index date to include. Must be >= 0. Default: 0
        
        Returns
        -------
        pd.DataFrame or None
            - PatientID : object
                unique patient identifier
            - BRCA_status : category
                positive if ever-positive, negative if only-negative, otherwise unknown

        Notes
        ------
        Biomarker cleaning processing: 
        - BRCA status is classifed according to these as:
            - 'positive' if any test result is positive (ever-positive)
            - 'negative' if any test is negative without positives (only-negative) 
            - 'unknown' if all results are indeterminate
        
        - Missing biomarker data handling:
            - All PatientIDs from index_date_df are included in the output
            - Patients without any biomarker tests will have NaN values for all biomarker columns
            - Missing ResultDate is imputed with SpecimenReceivedDate

        Output handling: 
        - Duplicate PatientIDs are logged as warnings if found but retained in output
        - Processed DataFrame is stored in self.biomarkers_df
        """
        # Input validation
        if not isinstance(index_date_df, pd.DataFrame):
            raise ValueError("index_date_df must be a pandas DataFrame")
        if 'PatientID' not in index_date_df.columns:
            raise ValueError("index_date_df must contain a 'PatientID' column")
        if not index_date_column or index_date_column not in index_date_df.columns:
            raise ValueError('index_date_column not found in index_date_df')
        if index_date_df['PatientID'].duplicated().any():
            raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
        
        if days_before is not None:
            if not isinstance(days_before, int) or days_before < 0:
                raise ValueError("days_before must be a non-negative integer or None")
        if not isinstance(days_after, int) or days_after < 0:
            raise ValueError("days_after must be a non-negative integer")
        
        index_date_df = index_date_df.copy()
        # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
        for col in index_date_df.columns:
            if col != 'PatientID':  # Keep PatientID unchanged for merging
                index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

        # Update index_date_column name
        index_date_column = f'imported_{index_date_column}'

        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Enhanced_MetPC_Biomarkers.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            df['ResultDate'] = pd.to_datetime(df['ResultDate'])
            df['SpecimenReceivedDate'] = pd.to_datetime(df['SpecimenReceivedDate'])

            # Impute missing ResultDate with SpecimenReceivedDate
            df['ResultDate'] = np.where(df['ResultDate'].isna(), df['SpecimenReceivedDate'], df['ResultDate'])

            index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])

            # Select PatientIDs that are included in the index_date_df the merge on 'left'
            df = df[df.PatientID.isin(index_date_df.PatientID)]
            df = pd.merge(
                    df,
                    index_date_df[['PatientID', index_date_column]],
                    on = 'PatientID',
                    how = 'left'
            )
            logging.info(f"Successfully merged Enhanced_MetPC_Biomarkers.csv df with index_date_df resulting in shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")
            
            # Create new variable 'index_to_result' that notes difference in days between resulted specimen and index date
            df['index_to_result'] = (df['ResultDate'] - df[index_date_column]).dt.days
            
            # Select biomarkers that fall within desired before and after index date
            if days_before is None:
                # Only filter for days after
                df_filtered = df[df['index_to_result'] <= days_after].copy()
            else:
                # Filter for both before and after
                df_filtered = df[
                    (df['index_to_result'] <= days_after) & 
                    (df['index_to_result'] >= -days_before)
                ].copy()

            # Process BRCA
            positive_values = {
                'BRCA1 mutation identified',
                'BRCA2 mutation identified',
                'Both BRCA1 and BRCA2 mutations identified',
                'BRCA mutation NOS' 
            }

            negative_values = {
                'No BRCA mutation',
                'Genetic Variant Favor Polymorphism',
            }

            brca_df = (
                df_filtered
                .query('BiomarkerName == "BRCA"')
                .groupby('PatientID')['BiomarkerStatus']
                .agg(lambda x: 'positive' if any(val in positive_values for val in x)
                    else ('negative' if any(val in negative_values for val in x)
                        else 'unknown'))
                .reset_index()
                .rename(columns={'BiomarkerStatus': 'BRCA_status'}) 
            )

            # Merge dataframes -- start with index_date_df to ensure all PatientIDs are included
            final_df = index_date_df[['PatientID']].copy()
            final_df = pd.merge(final_df, brca_df, on = 'PatientID', how = 'left')
            final_df['BRCA_status'] = final_df['BRCA_status'].astype('category')
            
            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset = ['PatientID'], keep = False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Enhanced_MetPC_Biomarkers.csv file with final shape: {final_df.shape} and unique PatientIDs: {(final_df['PatientID'].nunique())}")
            self.biomarkers_df = final_df
            return final_df

        except Exception as e:
            logging.error(f"Error processing Enhanced_MetPC_Biomarkers.csv file: {e}")
            return None
    
    def process_labs(self,
                     file_path: str,
                     index_date_df: pd.DataFrame,
                     index_date_column: str, 
                     additional_loinc_mappings: dict = None,
                     days_before: int = 90,
                     days_after: int = 0,
                     summary_lookback: int = 180) -> Optional[pd.DataFrame]:
        """
        Processes Lab.csv to determine patient lab values within a specified time window relative to an index date. Returns CBC, CMP, and PSA total values 
        nearest to index date, along with summary statistics (max, min, standard deviation, and slope) calculated over the summary period. PSA doubling time
        is also calculated over summary period. Additional lab tests can be included by providing corresponding LOINC code mappings.

        Parameters
        ----------
        file_path : str
            Path to Labs.csv file
        index_date_df : pd.DataFrame
            DataFrame containing unique PatientIDs and their corresponding index dates. Only labs for PatientIDs present in this DataFrame will be processed
        index_date_column : str
            Column name in index_date_df containing the index date
        additional_loinc_mappings : dict, optional
            Dictionary of additional lab names and their LOINC codes to add to the default mappings.
            Example: {'CEA': ['2039-6'], 'CA_15-3': ['6875-9'], 'CA_27-29': ['17842-6'], 'another_lab': ['12345', '67889-0']}
        days_before : int, optional
            Number of days before the index date to include for baseline lab values. Must be >= 0. Default: 90
        days_after : int, optional
            Number of days after the index date to include for baseline lab values. Also used as the end point for 
            summary statistics calculations. Must be >= 0. Default: 0
        summary_lookback : int, optional
            Number of days before index date to begin analyzing summary statistics. Analysis period extends 
            from (index_date - summary_lookback) to (index_date + days_after). Must be >= 0. Default: 180

        Returns
        -------
        pd.DataFrame or None
            - PatientID : object
                unique patient identifier

            Baseline values (closest to index date within days_before/days_after window):
            - hemoglobin : float, g/dL
            - wbc : float, K/uL
            - platelet : float, 10^9/L
            - creatinine : float, mg/dL
            - bun : float, mg/dL
            - sodium : float, mmol/L
            - chloride : float, mmol/L
            - bicarbonate : float, mmol/L
            - potassium : float, mmol/L
            - calcium : float, mg/dL
            - alp : float, U/L
            - ast : float, U/L
            - alt : float, U/L
            - total_bilirubin : float, mg/dL
            - albumin : float, g/L
            - psa : float, ug/L

            Summary statistics (calculated over period from index_date - summary_lookback to index_date + days_after):
            For each lab above, includes:
            - {lab}_max : float, maximum value
            - {lab}_min : float, minimum value
            - {lab}_std : float, standard deviation
            - {lab}_slope : float, rate of change over time (days)
            
            Bonus statistic calculated over the same period as the summary statistics:
            - psa_doubling_time : float, months

        Notes
        -----
        Data cleaning and processing: 
        - Imputation strategy for lab dates: missing ResultDate is imputed with TestDate
        - Imputation strategy for lab values:
            - For each lab, missing TestResultCleaned values are imputed from TestResult after removing flags (L, H, <, >)
            - Values outside physiological ranges for each lab are filtered out
        -Unit conversion corrections:
            - Hemoglobin: Values in g/uL are divided by 100,000 to convert to g/dL
            - WBC/Platelet: Values in 10*3/L are multiplied by 1,000,000; values in /mm3 or 10*3/mL are multiplied by 1,000
            - Creatinine/BUN/Calcium: Values in mg/L are multiplied by 10 to convert to mg/dL
            - Albumin: Values in mg/dL are multiplied by 1,000 to convert to g/L; values 1-6 are assumed to be g/dL and multiplied by 10
        - Lab value selection: 
            - Baseline lab value closest to index date is selected by minimum absolute day difference within window period of 
            (index_date - days_before) to (index_date + days_after)
            - Summary lab values are calculated within window period of (index_date - summary_lookback) to (index_date + days_after)
        For slope and PSA doubling time calculation:
            - Patient needs at least 2 valid measurements, at 2 valid time points, and time points must not be identical
            - PSA doubling formula: (ln(2)/PSA slope), measured in months
            - Where PSA slope = [ln(PSAMetDiagnosis) - ln(PSADiagnosis)]/[(MetDiagnosisDate - DiagnosisDate) in months]
        
        Output handling: 
        - All PatientIDs from index_date_df are included in the output and values are NaN for patients without lab values 
        - Duplicate PatientIDs are logged as warnings but retained in output 
        - Results are stored in self.labs_df attribute
        """
        # Input validation
        if not isinstance(index_date_df, pd.DataFrame):
            raise ValueError("index_date_df must be a pandas DataFrame")
        if 'PatientID' not in index_date_df.columns:
            raise ValueError("index_date_df must contain a 'PatientID' column")
        if not index_date_column or index_date_column not in index_date_df.columns:
            raise ValueError('index_date_column not found in index_date_df')
        if index_date_df['PatientID'].duplicated().any():
            raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
        
        if not isinstance(days_before, int) or days_before < 0:
            raise ValueError("days_before must be a non-negative integer")
        if not isinstance(days_after, int) or days_after < 0:
            raise ValueError("days_after must be a non-negative integer")
        if not isinstance(summary_lookback, int) or summary_lookback < 0:
            raise ValueError("summary_lookback must be a non-negative integer")
        
        index_date_df = index_date_df.copy()
        # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
        for col in index_date_df.columns:
            if col != 'PatientID':  # Keep PatientID unchanged for merging
                index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

        # Update index_date_column name
        index_date_column = f'imported_{index_date_column}'
        
        # Add user-provided mappings if they exist
        if additional_loinc_mappings is not None:
            if not isinstance(additional_loinc_mappings, dict):
                raise ValueError("Additional LOINC mappings must be provided as a dictionary")
            if not all(isinstance(v, list) for v in additional_loinc_mappings.values()):
                raise ValueError("LOINC codes must be provided as lists of strings")
                
            # Update the default mappings with additional ones
            self.LOINC_MAPPINGS.update(additional_loinc_mappings)

        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Lab.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            df['ResultDate'] = pd.to_datetime(df['ResultDate'])
            df['TestDate'] = pd.to_datetime(df['TestDate'])

            # Impute TestDate for missing ResultDate. 
            df['ResultDate'] = np.where(df['ResultDate'].isna(), df['TestDate'], df['ResultDate'])

            index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])

            # Select PatientIDs that are included in the index_date_df the merge on 'left'
            df = df[df.PatientID.isin(index_date_df.PatientID)]
            df = pd.merge(
                df,
                index_date_df[['PatientID', index_date_column]],
                on = 'PatientID',
                how = 'left'
            )
            logging.info(f"Successfully merged Lab.csv df with index_date_df resulting in shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")
            
            # Flatten LOINC codes 
            all_loinc_codes = sum(self.LOINC_MAPPINGS.values(), [])

            # Filter for LOINC codes 
            df = df[df['LOINC'].isin(all_loinc_codes)]

            # Map LOINC codes to lab names
            for lab_name, loinc_codes in self.LOINC_MAPPINGS.items():
                mask = df['LOINC'].isin(loinc_codes)
                df.loc[mask, 'lab_name'] = lab_name

            ## CBC PROCESSING ##
            
            # Hemoglobin conversion correction
            # TestResultCleaned incorrectly stored g/uL values 
            # Example: 12 g/uL was stored as 1,200,000 g/dL instead of 12 g/dL
            # Need to divide by 100,000 to restore correct value
            mask = (
                (df['lab_name'] == 'hemoglobin') & 
                (df['TestUnits'] == 'g/uL')
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] / 100000 

            # WBC and Platelet conversion correction
            # TestResultCleaned incorrectly stored 10*3/L values 
            # Example: 9 10*3/L was stored as 0.000009 10*9/L instead of 9 10*9/L
            # Need to multipley 1,000,000 to restore correct value
            mask = (
                ((df['lab_name'] == 'wbc') | (df['lab_name'] == 'platelet')) & 
                (df['TestUnits'] == '10*3/L')
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] * 1000000

           # WBC and Platelet conversion correction
            # TestResultCleaned incorrectly stored /mm3 and 10*3/mL values
            # Example: 9 /mm3 and 9 10*3/mL was stored as 0.009 10*9/L instead of 9 10*9/L
            # Need to multipley 1,000 to restore correct value
            mask = (
                ((df['lab_name'] == 'wbc') | (df['lab_name'] == 'platelet')) & 
                ((df['TestUnits'] == '/mm3') | (df['TestUnits'] == '10*3/mL'))
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] * 1000

            # Hemoglobin: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 3-20; and impute to TestResultCleaned
            mask = df.query('lab_name == "hemoglobin" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 3) & (x <= 20))
            )

            # WBC: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-40; and impute to TestResultCleaned
            mask = df.query('lab_name == "wbc" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 40))
            )
            
            # Platelet: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-1000; and impute to TestResultCleaned
            mask = df.query('lab_name == "platelet" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 1000))
            )

            ## CMP PROCESSING ##
            # Creatinine, BUN, and calcium conversion correction
            # TestResultCleaned incorrectly stored mg/L values 
            # Example: 1.6 mg/L was stored as 0.16 mg/dL instead of 1.6 mg/dL
            # Need to divide by 10 to restore correct value
            mask = (
                ((df['lab_name'] == 'creatinine') | (df['lab_name'] == 'bun') | (df['lab_name'] == 'calcium')) & 
                (df['TestUnits'] == 'mg/L')
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] * 10 

            # Albumin conversion correction
            # TestResultCleaned incorrectly stored mg/dL values 
            # Example: 3.7 mg/dL was stored as 0.037 g/L instead of 37 g/L
            # Need to multiply 1000 to restore correct value
            mask = (
                (df['lab_name'] == 'albumin') & 
                (df['TestUnits'] == 'mg/dL')
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] * 1000         

            # Creatinine: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-5; and impute to TestResultCleaned 
            mask = df.query('lab_name == "creatinine" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 5))
            )
            
            # BUN: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-100; and impute to TestResultCleaned 
            mask = df.query('lab_name == "bun" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 100))
            )

            # Sodium: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 110-160; and impute to TestResultCleaned 
            mask = df.query('lab_name == "sodium" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 110) & (x <= 160))
            )
            
            # Chloride: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 70-140; and impute to TestResultCleaned 
            mask = df.query('lab_name == "chloride" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 70) & (x <= 140))
            )
            
            # Bicarbonate: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 5-50; and impute to TestResultCleaned 
            mask = df.query('lab_name == "bicarbonate" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 5) & (x <= 50))
            )

            # Potassium: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 2-8; and impute to TestResultCleaned  
            mask = df.query('lab_name == "potassium" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 2) & (x <= 8))
            )
            
            # Calcium: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 5-15; and impute to TestResultCleaned 
            mask = df.query('lab_name == "calcium" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 5) & (x <= 15))
            )
            
            # ALP: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 20-3000; and impute to TestResultCleaned
            mask = df.query('lab_name == "alp" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 20) & (x <= 3000))
            )
            
            # AST: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 5-2000; and impute to TestResultCleaned
            mask = df.query('lab_name == "ast" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 5) & (x <= 2000))
            )
            
            # ALT: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 5-2000; and impute to TestResultCleaned
            mask = df.query('lab_name == "alt" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 5) & (x <= 2000))
            )
            
            # Total bilirubin: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-40; and impute to TestResultCleaned
            mask = df.query('lab_name == "total_bilirubin" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 40))
            )
            
            # Albumin
            mask = df.query('lab_name == "albumin" and TestResultCleaned.isna() and TestResult.notna()').index
            
            # First get the cleaned numeric values
            cleaned_alb_values = pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')

            # Identify which values are likely in which unit system
            # Values 1-6 are likely g/dL and need to be converted to g/L
            gdl_mask = (cleaned_alb_values >= 1) & (cleaned_alb_values <= 6)
            # Values 10-60 are likely already in g/L
            gl_mask = (cleaned_alb_values >= 10) & (cleaned_alb_values <= 60)

            # Convert g/dL values to g/L (multiply by 10)
            df.loc[mask[gdl_mask], 'TestResultCleaned'] = cleaned_alb_values[gdl_mask] * 10

            # Keep g/L values as they are
            df.loc[mask[gl_mask], 'TestResultCleaned'] = cleaned_alb_values[gl_mask]

            ## PSA PROCESSING ##
            # PSA conversion correction
            # TestResultCleaned incorrectly stored mg/dL values 
            # Example: 6.7 mg/dL was stored as 66,700 ug/L instead 6.7 ug/L
            # Need to divide by 1000 to restore correct value
            mask = (
                (df['lab_name'] == 'psa') & 
                (df['TestUnits'] == 'mg/dL')
            )
            df.loc[mask, 'TestResultCleaned'] = df.loc[mask, 'TestResultCleaned'] / 1000

            # PSA: Convert to TestResult to numeric after removing L, H, <, and >; filter for ranges from 0-1000; and impute to TestResultCleaned
            mask = df.query('lab_name == "psa" and TestResultCleaned.isna() and TestResult.notna()').index
            df.loc[mask, 'TestResultCleaned'] = (
                pd.to_numeric(df.loc[mask, 'TestResult'].str.replace(r'[LH<>]', '', regex=True).str.strip(), errors='coerce')
                .where(lambda x: (x >= 0) & (x <= 1000))
            ) 

            # Filter for desired window period for baseline labs after removing missing values after above imputation
            df = df.query('TestResultCleaned.notna()')
            df['index_to_lab'] = (df['ResultDate'] - df[index_date_column]).dt.days
            
            df_lab_index_filtered = df[
                (df['index_to_lab'] <= days_after) & 
                (df['index_to_lab'] >= -days_before)].copy()
            
            lab_df = (
                df_lab_index_filtered
                .assign(abs_index_to_lab = lambda x: abs(x['index_to_lab']))
                .sort_values('abs_index_to_lab')  
                .groupby(['PatientID', 'lab_name'])
                .first()  
                .reset_index()
                .pivot(index = 'PatientID', columns = 'lab_name', values = 'TestResultCleaned')
                .rename_axis(columns = None)
                .reset_index()
            )

            # Filter for desired window period for summary labs 
            df_lab_summary_filtered = df[
                (df['index_to_lab'] <= days_after) & 
                (df['index_to_lab'] >= -summary_lookback)].copy()
            
            max_df = (
                df_lab_summary_filtered
                .groupby(['PatientID', 'lab_name'])['TestResultCleaned'].max()
                .reset_index()
                .pivot(index = 'PatientID', columns = 'lab_name', values = 'TestResultCleaned')
                .rename_axis(columns = None)
                .rename(columns = lambda x: f'{x}_max')
                .reset_index()
            )
            
            min_df = (
                df_lab_summary_filtered
                .groupby(['PatientID', 'lab_name'])['TestResultCleaned'].min()
                .reset_index()
                .pivot(index = 'PatientID', columns = 'lab_name', values = 'TestResultCleaned')
                .rename_axis(columns = None)
                .rename(columns = lambda x: f'{x}_min')
                .reset_index()
            )
            
            std_df = (
                df_lab_summary_filtered
                .groupby(['PatientID', 'lab_name'])['TestResultCleaned'].std()
                .reset_index()
                .pivot(index = 'PatientID', columns = 'lab_name', values = 'TestResultCleaned')
                .rename_axis(columns = None)
                .rename(columns = lambda x: f'{x}_std')
                .reset_index()
            )
            
            slope_df = (
                df_lab_summary_filtered
                .groupby(['PatientID', 'lab_name'])[['index_to_lab', 'TestResultCleaned']]
                .apply(lambda x: np.polyfit(x['index_to_lab'],
                                            x['TestResultCleaned'],
                                            1)[0]                       # Extract slope coefficient with [0]
                    if (x['TestResultCleaned'].notna().sum() > 1 and    # Need at least 2 valid measurements
                        x['index_to_lab'].notna().sum() > 1 and         # Need at least 2 valid time points
                        len(x['index_to_lab'].unique()) > 1)            # Time points must not be identical
                    else np.nan)                                        # Return NaN if conditions for valid slope calculation aren't met
                .reset_index()
                .pivot(index = 'PatientID', columns = 'lab_name', values = 0)
                .rename_axis(columns = None)
                .rename(columns = lambda x: f'{x}_slope')
                .reset_index()
            )

            psa_doubling_df = (
                df_lab_summary_filtered.query('lab_name == "psa"')
                .groupby('PatientID')[['index_to_lab', 'TestResultCleaned']]
                .apply(lambda x: 
                    # Inner lambda: calculates slope, once passes data check
                    (lambda slope = np.polyfit(
                        x['index_to_lab']/30,                                       # /30 to convert days to months 
                        np.log(x['TestResultCleaned']),                             # log transform PSA values    
                        1)[0]:                                                      # Extract slope coefficient with [0]
                        
                        # Check if slope is positive 
                        math.log(2)/slope if slope > 0 else np.nan)()               # Return NaN if slope is negative or zero

                    # Outer lambda: data quality check     
                    if (x['TestResultCleaned'].notna().sum() > 1 and                # Need at least 2 valid measurements
                        x['index_to_lab'].notna().sum() > 1 and                     # Need at least 2 valid time points
                        len(x['index_to_lab'].unique()) > 1)                        # Time points must not be identical
                    else np.nan)
                .reset_index()
                .rename(columns={0: 'psa_doubling'})
            )
            
            # Merge dataframes - start with index_date_df to ensure all PatientIDs are included
            final_df = index_date_df[['PatientID']].copy()
            final_df = pd.merge(final_df, lab_df, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, max_df, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, min_df, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, std_df, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, slope_df, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, psa_doubling_df, on = 'PatientID', how = 'left')

            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset = ['PatientID'], keep = False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Lab.csv file with final shape: {final_df.shape} and unique PatientIDs: {(final_df['PatientID'].nunique())}")
            self.labs_df = final_df
            return final_df

        except Exception as e:
            logging.error(f"Error processing Lab.csv file: {e}")
            return None
        
    def process_diagnosis(self, 
                          file_path: str,
                          index_date_df: pd.DataFrame,
                          index_date_column: str,
                          days_before: Optional[int] = None,
                          days_after: int = 0) -> Optional[pd.DataFrame]:
        """
        Processes Diagnosis.csv by mapping ICD 9 and 10 codes to Elixhauser comorbidity index and calculates a van Walraven score. 
        It also determines site of metastases based on ICD 9 and 10 codes. 
        
        Parameters
        ----------
        file_path : str
            Path to Diagnosis.csv file
        index_date_df : pd.DataFrame
            DataFrame containing unique PatientIDs and their corresponding index dates. Only diagnoses for PatientIDs present in this DataFrame will be processed
        index_date_column : str
            Column name in index_date_df containing the index date
        days_before : int | None, optional
            Number of days before the index date to include for window period. Must be >= 0 or None. If None, includes all prior results. Default: None
        days_after : int, optional
            Number of days after the index date to include for window period. Must be >= 0. Default: 0
        
        Returns
        -------
        pd.DataFrame or None
            - PatientID : object unique patient identifier
            - chf : binary indicator for congestive heart failure
            - cardiac_arrhythmia : binary indicator for cardiac arrhythmias
            - valvular_disease : binary indicator for valvular disease
            - pulm_circulation : binary indicator for pulmonary circulation disorders
            - pvd : binary indicator for peripheral vascular disease
            - htn_uncomplicated : binary indicator for hypertension, uncomplicated
            - htn_complicated : binary indicator for hypertension, complicated
            - paralysis : binary indicator for paralysis
            - other_neuro : binary indicator for other neurological disorders
            - chronic_pulm_disease : binary indicator for chronic pulmonary disease
            - diabetes_uncomplicated : binary indicator for diabetes, uncomplicated
            - diabetes_complicated : binary indicator for diabetes, complicated
            - hypothyroid : binary indicator for hypothyroidism
            - renal_failure : binary indicator for renal failure
            - liver_disease : binary indicator for liver disease
            - PUD : binary indicator for peptic ulcer disease
            - aids_hiv : binary indicator for AIDS/HIV
            - lymphoma : binary indicator for lymphoma
            - rheumatic : binary indicator for rheumatoid arthritis/collagen vascular diseases
            - coagulopathy : binary indicator for coagulopathy
            - obesity : binary indicator for obesity
            - weight_loss : binary indicator for weight loss
            - fluid : binary indicator for fluid and electrolyte disorders
            - blood_loss_anemia : binary indicator for blood loss anemia
            - deficiency_anemia : binary indicator for deficiency anemia
            - alcohol_abuse : binary indicator for alcohol abuse
            - drug_abuse : binary indicator for drug abuse
            - psychoses : binary indicator for psychoses
            - depression : binary indicator for depression
            - van_walraven_score : weighted composite of the binary Elixhauser comorbidities 
            - lymph_met : binary indicator for lymph node metastasis
            - thoracic_met : binary indicator for thoracic metastasis (eg., lung, pleura, mediastinum, or other respiratory)
            - liver_met : binary indicator for liver metastasis
            - bone_met : binary indicator for bone metastasis 
            - brain_met : binary indicator for brain/CNS metastasis
            - other_gi_met : binary indicator for gastrointestinal tract metastasis excluding liver
            - other_met : binary indicator for other sites of metastasis

        Notes
        -----
        Mapping information: 
        - See "Coding algorithms for defining comorbidities in ICD-9-CM and ICD-10 administrative data" by Quan et al for details on ICD mapping to comorbidities. 
        For ICD-9 codes, the Enhanced ICD-9-CM by Quan was used for mapping.
        - See "A modification of the Elixhauser comorbidity measures into a point system for hospital death using administrative data" by van Walraven et al for 
        details on van Walraven score.
        - Metastatic cancer and tumor categories are excluded in the Elxihauser comorbidities and van Walraven score as this is intended for an advanced cancer population
        
        Output handling: 
        - All PatientIDs from index_date_df are included in the output and values will be set to 0 for patients with misisng Elixhauser comorbidities or metastasis sites, but NaN for missing van_walraven_score
        - Duplicate PatientIDs are logged as warnings but retained in output
        - Results are stored in self.diagnoses_df attribute
        """
        # Input validation
        if not isinstance(index_date_df, pd.DataFrame):
            raise ValueError("index_date_df must be a pandas DataFrame")
        if 'PatientID' not in index_date_df.columns:
            raise ValueError("index_date_df must contain a 'PatientID' column")
        if not index_date_column or index_date_column not in index_date_df.columns:
            raise ValueError('index_date_column not found in index_date_df')
        if index_date_df['PatientID'].duplicated().any():
            raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
        
        if days_before is not None:
            if not isinstance(days_before, int) or days_before < 0:
                raise ValueError("days_before must be a non-negative integer or None")
        if not isinstance(days_after, int) or days_after < 0:
            raise ValueError("days_after must be a non-negative integer")
        
        index_date_df = index_date_df.copy()
        # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
        for col in index_date_df.columns:
            if col != 'PatientID':  # Keep PatientID unchanged for merging
                index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

        # Update index_date_column name
        index_date_column = f'imported_{index_date_column}'

        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Diagnosis.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            df['DiagnosisDate'] = pd.to_datetime(df['DiagnosisDate'])
            index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])

            # Select PatientIDs that are included in the index_date_df the merge on 'left'
            df = df[df.PatientID.isin(index_date_df.PatientID)]
            df = pd.merge(
                df,
                index_date_df[['PatientID', index_date_column]],
                on = 'PatientID',
                how = 'left'
            )
            logging.info(f"Successfully merged Diagnosis.csv df with index_date_df resulting in shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            df['index_to_diagnosis'] = (df['DiagnosisDate'] - df[index_date_column]).dt.days
            
            # Select ICD codes that fall within desired before and after index date
            if days_before is None:
                # Only filter for days after
                df_filtered = df[df['index_to_diagnosis'] <= days_after].copy()
            else:
                # Filter for both before and after
                df_filtered = df[
                    (df['index_to_diagnosis'] <= days_after) & 
                    (df['index_to_diagnosis'] >= -days_before)
                ].copy()

            # Elixhauser comorbidities based on ICD-9 codes
            df9_elix = (
                df_filtered
                .query('DiagnosisCodeSystem == "ICD-9-CM"')
                .assign(diagnosis_code = lambda x: x['DiagnosisCode'].replace(r'\.', '', regex=True)) # Remove decimal points from ICD-9 codes to make mapping easier 
                .drop_duplicates(subset = ['PatientID', 'diagnosis_code'], keep = 'first')
                .assign(comorbidity=lambda x: x['diagnosis_code'].map(
                    lambda code: next((comorb for pattern, comorb in self.ICD_9_EXLIXHAUSER_MAPPING.items() 
                                    if re.match(pattern, code)), 'Other')))
                .query('comorbidity != "Other"') 
                .drop_duplicates(subset=['PatientID', 'comorbidity'], keep = 'first')
                .assign(value=1)  # Add a column of 1s to use for pivot
                .pivot(index = 'PatientID', columns = 'comorbidity', values = 'value')
                .fillna(0) 
                .astype('Int64')  
                .rename_axis(columns = None)
                .reset_index()
            )

            # Elixhauser comorbidities based on ICD-10 codes
            df10_elix = (
                df_filtered
                .query('DiagnosisCodeSystem == "ICD-10-CM"')
                .assign(diagnosis_code = lambda x: x['DiagnosisCode'].replace(r'\.', '', regex=True)) # Remove decimal points from ICD-10 codes to make mapping easier 
                .drop_duplicates(subset = ['PatientID', 'diagnosis_code'], keep = 'first')
                .assign(comorbidity=lambda x: x['diagnosis_code'].map(
                    lambda code: next((comorb for pattern, comorb in self.ICD_10_ELIXHAUSER_MAPPING.items() 
                                    if re.match(pattern, code)), 'Other')))
                .query('comorbidity != "Other"') 
                .drop_duplicates(subset=['PatientID', 'comorbidity'], keep = 'first')
                .assign(value=1)  # Add a column of 1s to use for pivot
                .pivot(index = 'PatientID', columns = 'comorbidity', values = 'value')
                .fillna(0) 
                .astype('Int64')
                .rename_axis(columns = None)
                .reset_index()  
            )

            all_columns_elix = ['PatientID'] + list(self.ICD_9_EXLIXHAUSER_MAPPING.values())
            
            # Reindex both dataframes to have all columns, filling missing ones with 0
            df9_elix_aligned = df9_elix.reindex(columns = all_columns_elix, fill_value = 0)
            df10_elix_aligned = df10_elix.reindex(columns = all_columns_elix, fill_value = 0)

            # Combine Elixhauser comorbidity dataframes for ICD-9 and ICD-10
            df_elix_combined = pd.concat([df9_elix_aligned, df10_elix_aligned]).groupby('PatientID').max().reset_index()

            # Calculate van Walraven score
            van_walraven_score = df_elix_combined.drop('PatientID', axis=1).mul(self.VAN_WALRAVEN_WEIGHTS).sum(axis=1)
            df_elix_combined['van_walraven_score'] = van_walraven_score

            # Metastatic sites based on ICD-9 codes 
            df9_mets = (
                df_filtered
                .query('DiagnosisCodeSystem == "ICD-9-CM"')
                .assign(diagnosis_code = lambda x: x['DiagnosisCode'].replace(r'\.', '', regex=True)) # Remove decimal points from ICD-9 codes to make mapping easier 
                .drop_duplicates(subset = ['PatientID', 'diagnosis_code'], keep = 'first')
                .assign(met_site=lambda x: x['diagnosis_code'].map(
                    lambda code: next((site for pattern, site in self.ICD_9_METS_MAPPING.items()
                                    if re.match(pattern, code)), 'no_met')))
                .query('met_site != "no_met"') 
                .drop_duplicates(subset=['PatientID', 'met_site'], keep = 'first')
                .assign(value=1)  # Add a column of 1s to use for pivot
                .pivot(index = 'PatientID', columns = 'met_site', values = 'value')
                .fillna(0) 
                .astype('Int64')  
                .rename_axis(columns = None)
                .reset_index()
            )

            # Metastatic sites based on ICD-10 codes 
            df10_mets = (
                df_filtered
                .query('DiagnosisCodeSystem == "ICD-10-CM"')
                .assign(diagnosis_code = lambda x: x['DiagnosisCode'].replace(r'\.', '', regex=True)) # Remove decimal points from ICD-9 codes to make mapping easier 
                .drop_duplicates(subset = ['PatientID', 'diagnosis_code'], keep = 'first')
                .assign(met_site=lambda x: x['diagnosis_code'].map(
                    lambda code: next((site for pattern, site in self.ICD_10_METS_MAPPING.items()
                                    if re.match(pattern, code)), 'no_met')))
                .query('met_site != "no_met"') 
                .drop_duplicates(subset=['PatientID', 'met_site'], keep = 'first')
                .assign(value=1)  # Add a column of 1s to use for pivot
                .pivot(index = 'PatientID', columns = 'met_site', values = 'value')
                .fillna(0) 
                .astype('Int64')  
                .rename_axis(columns = None)
                .reset_index()
            )

            all_columns_mets = ['PatientID'] + list(self.ICD_9_METS_MAPPING.values())
            
            # Reindex both dataframes to have all columns, filling missing ones with 0
            df9_mets_aligned = df9_mets.reindex(columns = all_columns_mets, fill_value = 0)
            df10_mets_aligned = df10_mets.reindex(columns = all_columns_mets, fill_value = 0)

            df_mets_combined = pd.concat([df9_mets_aligned, df10_mets_aligned]).groupby('PatientID').max().reset_index()

            # Start with index_date_df to ensure all PatientIDs are included
            final_df = index_date_df[['PatientID']].copy()
            final_df = pd.merge(final_df, df_elix_combined, on = 'PatientID', how = 'left')
            final_df = pd.merge(final_df, df_mets_combined, on = 'PatientID', how = 'left')

            binary_columns = [col for col in final_df.columns 
                    if col not in ['PatientID', 'van_walraven_score']]
            final_df[binary_columns] = final_df[binary_columns].fillna(0).astype('Int64')

            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset = ['PatientID'], keep = False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Diagnosis.csv file with final shape: {final_df.shape} and unique PatientIDs: {(final_df['PatientID'].nunique())}")
            self.diagnosis_df = final_df
            return final_df
        
        except Exception as e:
            logging.error(f"Error processing Diagnosis.csv file: {e}")
            return None
        
    def process_mortality(self,
                          file_path: str,
                          index_date_df: pd.DataFrame,
                          index_date_column: str,
                          visit_path: str = None, 
                          telemedicine_path: str = None, 
                          biomarkers_path: str = None, 
                          oral_path: str = None,
                          adt_path: str = None,
                          alpha_beta_emitters_path: str = None, 
                          primary_treatment_path: str = None,
                          provenge_path: str = None,
                          enhanced_path: str = None,
                          drop_dates: bool = True) -> Optional[pd.DataFrame]:
        """
        Processes Enhanced_Mortality_V2.csv by cleaning data types, calculating time from index date to death/censor, and determining mortality events. 

        Parameters
        ----------
        file_path : str
            Path to Enhanced_Mortality_V2.csv file
        index_date_df : pd.DataFrame
            DataFrame containing unique PatientIDs and their corresponding index dates. Only mortality data for PatientIDs present in this DataFrame will be processed
        index_date_column : str
            Column name in index_date_df containing the index date
        visit_path : str
            Path to Visit.csv file, using VisitDate to determine last EHR activity date for censored patients
        telemedicine_path : str
            Path to Telemedicine.csv file, using VisitDate to determine last EHR activity date for censored patients
        biomarkers_path : str
            Path to Enhanced_MetPC_Biomarkers.csv file, using SpecimenCollectedDate to determine last EHR activity date for censored patients
        oral_path : str
            Path to Enhanced_MetPC_Orals.csv file, using StartDate and EndDate to determine last EHR activity date for censored patients
        adt_path : str
            Path to Enhanced_MetPC_ADT.csv file, using StartDate and EndDate to determine last EHR activity date for censored patients 
        alpha_beta_emitters_path : str
            Path to Enhanced_MetPC_AlphaBetaEmitters.csv file, using AdministrationDate to determine last EHR activity date for censored patients 
        primary_treatment_path : str
            Path to Enhanced_MetPC_PrimaryTreatment.csv file, using TreatmentDate to determine last EHR activity date for censored patients 
        provenge_path : str
            Path to Enhanced_MetPC_Provenge.csv file, using StartDate to determine last EHR activity date for censored patients
        enhanced_path : str
            Path to Enhanced_MetProstate.csv file, using CRPCDate to determine last EHR activity date for censored patients
        drop_dates : bool, default = True
            If True, drops date columns (index_date_column, DateOfDeath, last_ehr_date)   
        
        Returns
        -------
        pd.DataFrame or None
            - PatientID : object
                unique patient identifier
            - duration : float
                days from index date to death or censor 
            - event : Int64
                mortality status (1 = death, 0 = censored)

        If drop_dates=False, the DataFrame will also include:
            - {index_date_column} : datetime64
                The index date for each patient
            - DateOfDeath : datetime64
                Date of death (if available)
            - last_ehr_activity : datetime64
                Most recent EHR activity date (if available from supplementary files)
                
        Notes
        ------
        Death date handling:
        - Known death date: 'event' = 1, 'duration' = days from index to death
        - No death date: 'event' = 0, 'duration' = days from index to last EHR activity
        
        Death date imputation for incomplete dates:
        - Missing day: Imputed to 15th of the month
        - Missing month and day: Imputed to July 1st of the year
    
        Censoring logic:
        - Patients without death dates are censored at their last EHR activity
        - Last EHR activity is determined as the maximum date across all provided
          supplementary files (visit, telemedicine, biomarkers, oral, ADT, alpha beta emitters, 
          primary treatment, provenge, and enhanced)
        - If no supplementary files are provided or a patient has no activity in 
          supplementary files, duration may be null for censored patients
        
        Output handling: 
        - Duplicate PatientIDs are logged as warnings if found but retained in output
        - Processed DataFrame is stored in self.mortality_df
        """
        # Input validation
        if not isinstance(index_date_df, pd.DataFrame):
            raise ValueError("index_date_df must be a pandas DataFrame")
        if 'PatientID' not in index_date_df.columns:
            raise ValueError("index_date_df must contain a 'PatientID' column")
        if not index_date_column or index_date_column not in index_date_df.columns:
            raise ValueError('index_date_column not found in index_date_df')
        if index_date_df['PatientID'].duplicated().any():
            raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
        
        index_date_df = index_date_df.copy()
        # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
        for col in index_date_df.columns:
            if col != 'PatientID':  # Keep PatientID unchanged for merging
                index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

        # Update index_date_column name
        index_date_column = f'imported_{index_date_column}'

        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Enhanced_Mortality_V2.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            # When only year is available: Impute to July 1st (mid-year)
            df['DateOfDeath'] = np.where(df['DateOfDeath'].str.len() == 4, df['DateOfDeath'] + '-07-01', df['DateOfDeath'])

            # When only month and year are available: Impute to the 15th day of the month
            df['DateOfDeath'] = np.where(df['DateOfDeath'].str.len() == 7, df['DateOfDeath'] + '-15', df['DateOfDeath'])

            df['DateOfDeath'] = pd.to_datetime(df['DateOfDeath'])

            # Process index dates and merge
            index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])
            df = pd.merge(
                index_date_df[['PatientID', index_date_column]],
                df,
                on = 'PatientID',
                how = 'left'
            )
            logging.info(f"Successfully merged Enhanced_Mortality_V2.csv df with index_date_df resulting in shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")
                
            # Create event column
            df['event'] = df['DateOfDeath'].notna().astype('Int64')

            # Initialize final dataframe
            final_df = df.copy()

            # Create a list to store all last activity date dataframes
            patient_last_dates = []

            # Determine last EHR data
            if all(path is None for path in [visit_path, telemedicine_path, biomarkers_path, oral_path, adt_path, alpha_beta_emitters_path, primary_treatment_path, provenge_path, enhanced_path]):
                logging.info("WARNING: At least one of visit_path, telemedicine_path, biomarkers_path, oral_path must, adt_path, alpha_beta_emitters_path, primary_treatment_path, provenge_path, or enhanced_path must be provided to calculate duration for those with a missing death date")
            else: 
                # Process visit and telemedicine data
                if visit_path is not None or telemedicine_path is not None:
                    visit_dates = []
                    try:
                        if visit_path is not None:
                            df_visit = pd.read_csv(visit_path)
                            df_visit['VisitDate'] = pd.to_datetime(df_visit['VisitDate'])
                            visit_dates.append(df_visit[['PatientID', 'VisitDate']])
                            
                        if telemedicine_path is not None:
                            df_tele = pd.read_csv(telemedicine_path)
                            df_tele['VisitDate'] = pd.to_datetime(df_tele['VisitDate'])
                            visit_dates.append(df_tele[['PatientID', 'VisitDate']])
                        
                        if visit_dates:
                            df_visit_combined = pd.concat(visit_dates)
                            df_visit_max = (
                                df_visit_combined
                                .query("PatientID in @index_date_df.PatientID")
                                .groupby('PatientID')['VisitDate']
                                .max()
                                .to_frame(name='last_visit_date')
                                .reset_index()
                            )
                            patient_last_dates.append(df_visit_max)
                    except Exception as e:
                        logging.error(f"Error processing Visit.csv or Telemedicine.csv: {e}")
                                            
                # Process biomarkers data
                if biomarkers_path is not None:
                    try: 
                        df_biomarkers = pd.read_csv(biomarkers_path)
                        df_biomarkers['SpecimenCollectedDate'] = pd.to_datetime(df_biomarkers['SpecimenCollectedDate'])

                        df_biomarkers_max = (
                            df_biomarkers
                            .query("PatientID in @index_date_df.PatientID")
                            .groupby('PatientID')['SpecimenCollectedDate']
                            .max()
                            .to_frame(name='last_biomarker_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_biomarkers_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_Biomarkers.csv file: {e}")

                # Process oral medication data
                if oral_path is not None:
                    try:
                        df_oral = pd.read_csv(oral_path)
                        df_oral['StartDate'] = pd.to_datetime(df_oral['StartDate'])
                        df_oral['EndDate'] = pd.to_datetime(df_oral['EndDate'])

                        df_oral_max = (
                            df_oral
                            .query("PatientID in @index_date_df.PatientID")
                            .assign(max_date=lambda x: x[['StartDate', 'EndDate']].max(axis=1))
                            .groupby('PatientID')['max_date']
                            .max()
                            .to_frame(name='last_oral_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_oral_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_Orals.csv file: {e}")
                
                # Process ADT data
                if adt_path is not None:
                    try:
                        df_adt = pd.read_csv(adt_path)
                        df_adt['StartDate'] = pd.to_datetime(df_adt['StartDate'])
                        df_adt['EndDate'] = pd.to_datetime(df_adt['EndDate'])

                        df_adt_max = (
                            df_adt
                            .query("PatientID in @index_date_df.PatientID")
                            .assign(max_date=lambda x: x[['StartDate', 'EndDate']].max(axis=1))
                            .groupby('PatientID')['max_date']
                            .max()
                            .to_frame(name='last_adt_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_adt_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_ADT.csv file: {e}")

                # Process alpha beta emitters data
                if alpha_beta_emitters_path is not None:
                    try:
                        df_abe = pd.read_csv(alpha_beta_emitters_path)
                        df_abe['AdministrationDate'] = pd.to_datetime(df_abe['AdministrationDate'])

                        df_abe_max = (
                            df_abe
                            .query("PatientID in @index_date_df.PatientID")
                            .groupby('PatientID')['AdministrationDate']
                            .max()
                            .to_frame(name='last_abe_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_abe_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_AlphaBetaEmitters.csv file: {e}")
                
                # Process primary treatment data
                if primary_treatment_path is not None:
                    try:
                        df_pt = pd.read_csv(primary_treatment_path)
                        df_pt['TreatmentDate'] = pd.to_datetime(df_pt['TreatmentDate'])

                        df_pt_max = (
                            df_pt
                            .query("PatientID in @index_date_df.PatientID")
                            .groupby('PatientID')['TreatmentDate']
                            .max()
                            .to_frame(name='last_primary_treatment_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_pt_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_PrimaryTreatment.csv file: {e}")

                # Process provenge data
                if provenge_path is not None:
                    try:
                        df_provenge = pd.read_csv(provenge_path)
                        df_provenge['StartDate'] = pd.to_datetime(df_provenge['StartDate'])

                        df_provenge_max = (
                            df_provenge
                            .query("PatientID in @index_date_df.PatientID")
                            .groupby('PatientID')['StartDate']
                            .max()
                            .to_frame(name='last_provenge_date')
                            .reset_index()
                        )
                        patient_last_dates.append(df_provenge_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetPC_Provenge.csv file: {e}")

                # Process enhanced data
                if enhanced_path is not None:
                    try:
                        df_enhanced = pd.read_csv(enhanced_path)
                        df_enhanced['CRPCDate'] = pd.to_datetime(df_enhanced['CRPCDate'])

                        df_enhanced_max = (
                            df_enhanced
                            .query("PatientID in @index_date_df.PatientID")
                            [['PatientID', 'CRPCDate']]
                        )
                        patient_last_dates.append(df_enhanced_max)
                    except Exception as e:
                        logging.error(f"Error reading Enhanced_MetProstate.csv file: {e}")

                # Combine all last activity dates
                if patient_last_dates:
                    # Start with the first dataframe
                    combined_dates = patient_last_dates[0]
                    
                    # Merge with any additional dataframes
                    for date_df in patient_last_dates[1:]:
                        combined_dates = pd.merge(combined_dates, date_df, on = 'PatientID', how = 'outer')
                    
                    # Calculate the last activity date across all columns
                    date_columns = [col for col in combined_dates.columns if col != 'PatientID']
                    if date_columns:
                        logging.info(f"The following columns {date_columns} are used to calculate the last EHR date")
                        combined_dates['last_ehr_activity'] = combined_dates[date_columns].max(axis=1)
                        single_date = combined_dates[['PatientID', 'last_ehr_activity']]
                        
                        # Merge with the main dataframe
                        final_df = pd.merge(final_df, single_date, on='PatientID', how='left')
     
            # Calculate duration
            if 'last_ehr_activity' in final_df.columns:
                final_df['duration'] = np.where(
                    final_df['event'] == 0, 
                    (final_df['last_ehr_activity'] - final_df[index_date_column]).dt.days, 
                    (final_df['DateOfDeath'] - final_df[index_date_column]).dt.days
                )
                
                # Drop date variables if specified
                if drop_dates:               
                    final_df = final_df.drop(columns=[index_date_column, 'DateOfDeath', 'last_ehr_activity'])
                       
            else: 
                final_df['duration'] = (final_df['DateOfDeath'] - final_df[index_date_column]).dt.days
            
                # Drop date variables if specified
                if drop_dates:               
                    final_df = final_df.drop(columns=[index_date_column, 'DateOfDeath'])
                
            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset=['PatientID'], keep=False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Enhanced_Mortality_V2.csv file with final shape: {final_df.shape} and unique PatientIDs: {final_df['PatientID'].nunique()}. There are {final_df['duration'].isna().sum()} out of {final_df['PatientID'].nunique()} patients with missing duration values")
            self.mortality_df = final_df
            return final_df

        except Exception as e:
            logging.error(f"Error processing Enhanced_Mortality_V2.csv file: {e}")
            return None
    
    def process_adt(self,
                    file_path: str,
                    index_date_df: pd.DataFrame,
                    index_date_column: str) -> Optional[pd.DataFrame]:
        """
        Processes Enhanced_MetPC_ADT.csv by determining if recipt of ADT occurred prior to index date of interest. 

        Parameters
        ----------
        file_path : str
            Path to Enhanced_MetPC_ADT.csv file
        index_date_df : pd.DataFrame
            DataFrame containing unique PatientIDs and their corresponding index dates. Only treatment data for PatientIDs present in this DataFrame will be processed
        index_date_column : str
            Column name in index_date_df containing the index date
        
        Returns
        -------
        pd.DataFrame or None
            - PatientID : object
                unique patient identifier
            - received_adt : Int64
                binary indicator (0/1) for recepit of ADT prior to index date 
            
        Notes
        ------
        - With >90% EndDate missing in the have a missing in the Enhanced_MetPC_ADT.csv file, predicting presence of ADT at time of index date or duration 
        of ADT is difficult and unreliable. Given this constraint, these calculations were not pursued since significant assumptions would need to be made. 
        - Missing StartDate was imputed with EndDate, if available.

        Output handling: 
        - Duplicate PatientIDs are logged as warnings if found but retained in output
        - Processed DataFrame is stored in self.adt_df
        """
        # Input validation
        if not isinstance(index_date_df, pd.DataFrame):
            raise ValueError("index_date_df must be a pandas DataFrame")
        if 'PatientID' not in index_date_df.columns:
            raise ValueError("index_date_df must contain a 'PatientID' column")
        if not index_date_column or index_date_column not in index_date_df.columns:
            raise ValueError('index_date_column not found in index_date_df')
        if index_date_df['PatientID'].duplicated().any():
            raise ValueError("index_date_df contains duplicate PatientID values, which is not allowed")
        
        index_date_df = index_date_df.copy()
        # Rename all columns from index_date_df except PatientID to avoid conflicts with merging and processing 
        for col in index_date_df.columns:
            if col != 'PatientID':  # Keep PatientID unchanged for merging
                index_date_df.rename(columns={col: f'imported_{col}'}, inplace=True)

        # Update index_date_column name
        index_date_column = f'imported_{index_date_column}'

        try:
            df_adt = pd.read_csv(file_path)
            df_adt['StartDate'] = pd.to_datetime(df_adt['StartDate'])
            df_adt['EndDate'] = pd.to_datetime(df_adt['EndDate'])

            # Impute EndDate for missing StartDate
            df_adt['StartDate'] = np.where(df_adt['StartDate'].isna(), df_adt['EndDate'], df_adt['StartDate'])

            index_date_df[index_date_column] = pd.to_datetime(index_date_df[index_date_column])

            # Select PatientIDs that are included in the index_date_df the merge on 'left'
            df_adt = df_adt[df_adt.PatientID.isin(index_date_df.PatientID)]
            df_adt = pd.merge(
                df_adt,
                index_date_df[['PatientID', index_date_column]], 
                on = 'PatientID',
                how = 'left'
            )

            # Create df for recepit of ADT relative to index date 
            ever_received_adt_df = (
                df_adt
                .assign(ever_received_adt = lambda x: np.where(x['StartDate'] <= x[index_date_column], 1, 0))
                .groupby('PatientID')['ever_received_adt']
                .max()
                .reset_index()
            )

            # Merge dataframes - start with index_date_df to ensure all PatientIDs are included
            final_df = index_date_df[['PatientID']].copy()
            final_df = pd.merge(final_df, ever_received_adt_df, on = 'PatientID', how = 'left')
            final_df['ever_received_adt'] = final_df['ever_received_adt'].astype('Int64')
          
            # Check for duplicate PatientIDs
            if len(final_df) > final_df['PatientID'].nunique():
                duplicate_ids = final_df[final_df.duplicated(subset=['PatientID'], keep=False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Enhanced_MetPC_ADT.csv file with final shape: {final_df.shape} and unique PatientIDs: {final_df['PatientID'].nunique()}")
            self.adt_df = final_df
            return final_df

        except Exception as e:
            logging.error(f"Error processing Enhanced_MetPC_ADT.csv file: {e}")
            return None