import streamlit as st
from analytics.reports.utils.p_and_c_variables import elb_hcpi_speciality
from ai.sql.executor import BigQueryExecution
from core.config import config

bq = BigQueryExecution(project_id=config.GCP_PROJECT)

@st.cache_data(ttl=3600*24, show_spinner=False)  # Cache for 24 hours
def initialize_provider_lists():
    """
    Initialize provider and payer lists when the application starts.
    This is cached for 24 hours to balance freshness and performance.
    """
    try:
        elb_hcpi_speciality_string = "'" + "','".join(elb_hcpi_speciality) + "'"
        hcpi_speciality_class_query = f'''select distinct Specialty 
            from ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "individual"
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''
        
        hcpi_speciality_group_query = f'''select distinct sProv_Specialty_Desc
            FROM ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "individual"
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        hcpi_state_query = f'''select distinct State
            FROM ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "individual" and State is not null
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        hcpo_speciality_class_query = f'''select distinct Specialty
            from ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "organization"
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        hcpo_speciality_group_query = f'''select distinct sProv_Specialty_Desc
            from ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "organization"
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        hcpo_state_query = f'''select distinct State
            FROM ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "organization" and State is not null
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        hcpo_name_query = f'''select distinct sProv_Name
            FROM ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Provider
            where sProv_Entity_Type = "organization"
            and Specialty in ({elb_hcpi_speciality_string})
            order by 1'''

        payers_name_query = '''select distinct Plan_Type
            FROM ASF_Med_AI_CLARIVATE_2023_12_29_Partitioned.Medication
            where Plan_Type is not null
            order by 1'''

        lists_dict = {
            'HCPI Speciality Class': bq.run_query(hcpi_speciality_class_query)['Specialty'].values.tolist(),
            'HCPI Speciality Group': bq.run_query(hcpi_speciality_group_query)['sProv_Specialty_Desc'].values.tolist(),
            'HCPI State': bq.run_query(hcpi_state_query)['State'].values.tolist(),
            'HCPO Speciality Class': bq.run_query(hcpo_speciality_class_query)['Specialty'].values.tolist(),
            'HCPO Speciality Group': bq.run_query(hcpo_speciality_group_query)['sProv_Specialty_Desc'].values.tolist(),
            'HCPO State': bq.run_query(hcpo_state_query)['State'].values.tolist(),
            'HCPO Name': bq.run_query(hcpo_name_query)['sProv_Name'].values.tolist(),
            'Payer Name': bq.run_query(payers_name_query)['Plan_Type'].values.tolist()
        }
        
        return lists_dict
    except Exception as e:
        st.error(f"An error occurred: {e}")
        return {
            'HCPI Specialty Class': [],
            'HCPI Specialty Group': [],
            'HCPI State': [],
            'HCPO Specialty Class': [],
            'HCPO Specialty Group': [],
            'HCPO State': [],
            'HCPO Name': [],
            'Payer Name': []
        }

@st.cache_data(ttl=3600*24, show_spinner=False)
def get_dx_rx_lists():
    try:
        dx_group_query = ('select distinct Diagnosis_ICD10_Group_Desc\n'
                          'from ASF_Med_AI_CLARIVATE_P_AND_C_3.Diagnosis_Aggregated\n'
                          'where Diagnosis_ICD10_Group_Desc != ""\n'
                          'order by 1\n')

        dx_code_query = ('select distinct concat(Diagnosis_ICD10," ", Diagnosis_Desc) as Diagnosis_ICD10_Desc\n'
                         'from ASF_Med_AI_CLARIVATE_P_AND_C_3.Diagnosis_Aggregated\n'
                         'where Diagnosis_ICD10 != "Total"\n'
                         'order by 1\n')
        

        rx_class_query = ('select distinct Medication_Class\n'
                          'from ASF_Med_AI_CLARIVATE_P_AND_C_3.Medication_Aggregated\n'
                          'order by 1\n')

        rx_generic_query = ('select distinct Medication_Generic\n'
                            'from ASF_Med_AI_CLARIVATE_P_AND_C_3.Medication_Aggregated\n'
                            'where Medication_Generic != "Total"\n'
                            'order by 1\n')

        rx_brand_query = ('select distinct Medication_Brand\n'
                          'from ASF_Med_AI_CLARIVATE_P_AND_C_3.Medication_Aggregated\n'
                          'where Medication_Brand not in ("Total", "Sold As Generic")\n'
                          'order by 1\n')

        dx_group_df = bq.run_query(dx_group_query)
        dx_group_list = dx_group_df['Diagnosis_ICD10_Group_Desc'].values.tolist()

        dx_code_df = bq.run_query(dx_code_query)
        dx_code_list = dx_code_df['Diagnosis_ICD10_Desc'].values.tolist()

        rx_class_df = bq.run_query(rx_class_query)
        rx_class_list = rx_class_df['Medication_Class'].values.tolist()

        rx_generic_df = bq.run_query(rx_generic_query)
        rx_generic_list = rx_generic_df['Medication_Generic'].values.tolist()

        rx_brand_df = bq.run_query(rx_brand_query)
        rx_brand_list = rx_brand_df['Medication_Brand'].values.tolist()

        lists_dict = {
            'Diagnosis code group': dx_group_list,
            'Diagnosis code': dx_code_list,
            'Medication class': rx_class_list,
            'Medication generic name': rx_generic_list,
            'Medication brand name': rx_brand_list
        }
        
        return lists_dict
    except Exception as e:
        st.error(f"An error occurred: {e}")
        return {
            'Diagnosis code group': [],
            'Diagnosis code': [],
            'Medication class': [],
            'Medication generic name': [],
            'Medication brand name': []
        }