
import numpy as np
from os import environ

from AvenirCommon.Database import GB_get_db_json, GB_file_exists
from AvenirCommon.Util import formatCountryFName, GBGetKeyValue

from SpectrumCommon.Const.GB import *
from SpectrumCommon.Const.DP import *
from SpectrumCommon.Const.TB import *
from SpectrumCommon.Const.GB.GBVersions import (TB_WHO_DB_VERSION,
                                               TB_DIAGNOSTICS_VERSION,
                                               TB_DYNAMICAL_MODEL_VERSION,
                                               TB_FORT_INPUT_VERSION,
                                               TB_FORT_OUTPUT_VERSION)
from SpectrumCommon.ServerUtils.ExceptionUtils import SpectrumException
from SpectrumCommon.Modvars.TB.TBModvDefs import *
from SpectrumCommon.Modvars.TB.TBNotifDefs import *
from SpectrumCommon.Modvars.TB.TBHouseholdDefs import *
from SpectrumCommon.Modvars.TB.TBHighRiskGrpDefs import *
from SpectrumCommon.Modvars.TB.TBResistanceDefs import *
from SpectrumCommon.Modvars.TB.TBPreventionDefs import *
from SpectrumCommon.Modvars.TB.TBARTDefs import *
from SpectrumCommon.Modvars.TB.TBScreeningDefs import *


#@autoreload
def init_modvars(countryID, first_year, final_year,  extra):
    
    # year options
   
    # base_year = TB_Base_Year if first_year!=2021 else 2021
    if (first_year<TB_Min_First_Year):
        raise SpectrumException(GB_SW_BADPROJPARAM, message=f'TB first year is less than {str(TB_Min_First_Year)}')
    
    if ( first_year>TB_Max_First_Year):
        raise SpectrumException(GB_SW_BADPROJPARAM, message=f'TB first year is greater than {str(TB_Max_First_Year)}')

    base_year = first_year
    # who_db_vsn = '_V6' # if first_year!=2021 else '_V1'
    # fort_vsn = '_V7' #if first_year!=2021 else '_V2'
    dp_t = base_year-GB_BaseYear
    subnatCode = GBGetKeyValue(extra, 'subNatCode', 0)
    connection = environ['AVENIR_SW_DEFAULT_DATA_CONNECTION']

    year_len = final_year-base_year + 1

    fort_slice = slice(base_year-TB_FirstHist_Year, min(TB_Fort_Final_Year, final_year) -TB_FirstHist_Year+1)
    tb_slice  = slice(0, min(TB_Fort_Final_Year, final_year)-base_year+1)

    ###############################################################################################
    ############################## Database defaults ##############################################
    ###############################################################################################
    
    db =  GB_get_db_json(connection, 
                         'tuberculosis', 
                         f'countries\\{countryID}_{TB_WHO_DB_VERSION}.JSON')
    diagnostic_db = GB_get_db_json(connection, 
                                   'tuberculosis', 
                                   f'diagnosticDefaults_{TB_DIAGNOSTICS_VERSION}.JSON')
    # screening_db  = GB_get_db_json(connection, 'tuberculosis', 'screeningDefaults_V1.JSON')

    
    first_idx = base_year-db['projections']['startYear']
    final_idx = final_year-db['projections']['startYear']
    burden_first_idx = base_year-2000
    burden_final_idx = final_year-2000

        # TB_HisIncidenceTag  : hist_inci,
        # TB_HisMortalityTag  : ,
    # flatline defaults 
    TB_flatlineDB(db, 'notification', final_idx)
    TB_flatlineDB(db, 'prevalencePct', final_idx)
    TB_flatlineDB(db, 'prevalence', final_idx)
    TB_flatlineDB(db, 'pop', final_idx)
    TB_flatlinBurden(db, 'e_inc_num', final_idx)
    TB_flatlinBurden(db, 'e_mort_num', final_idx)
    
    #subnat notification
    pop, pop_15_plus = TB_getPopDefaults(connection, countryID, subnatCode, base_year, final_year)

    ###############################################################################################
    ###############################################################################################
    ###############################################################################################


    hh_prev_methods, hh_prev_method_IDs = HouseholdPreventionMethodsUsed()
    ART_prev_methods, ART_prev_method_IDs  = ARTPreventionMethodsUsed()
    HRG_prev_methods, HRG_prev_method_IDs  = HRGPreventionMethodsUsed()

    fort_inputs = fort.get_tb_fort_inputs(countryID, final_year, tb_slice, fort_slice, None, None, None)
    notif_base_pct, notif_scale, c_ep, bac_pos_pct = create_notification_defaults_2022(countryID, db['projections']['notification'], db['Noti_Distribution'], TB_FinalHist_Year-db['projections']['startYear'])
    notif, notif_split, notif_pct, fort_outputs = create_notification_defaults(countryID, db['projections']['notification'], db['Noti_Distribution'], first_year, first_idx, final_idx, year_len, 
                                                                               notif_base_pct, tb_slice, fort_slice, TB_FORT_OUTPUT_VERSION, notif_scale, 1)

    rr_profile, rp_Rif = TB_RRProfileDefaults(db, notif_pct, year_len)

    ## Screening Algorithm secions
    screen_map = TB_screening_sens_spec_map()

    patient_screen_selected = TB_screening_alg_patient_defaults()
    patient_sensitivity = [screen_map[screen_id][TB_Sensitivity_index] for screen_id in patient_screen_selected]
    patient_specificity = [screen_map[screen_id][TB_Specificity_index] for screen_id in patient_screen_selected]

    household_screen_selected = TB_screening_alg_household_defaults()
    household_sensitivity = [screen_map[screen_id][TB_Sensitivity_index] for screen_id in household_screen_selected]
    household_specificity = [screen_map[screen_id][TB_Specificity_index] for screen_id in household_screen_selected]

    art_screen_selected = TB_screening_alg_art_defaults()

    art_not_serious_sensitivity = [screen_map[screen_id][TB_Sensitivity_index] for screen_id in art_screen_selected[TB_ARTNotServere]]
    art_not_serious_specificity = [screen_map[screen_id][TB_Specificity_index] for screen_id in art_screen_selected[TB_ARTNotServere]]
    art_serious_sensitivity = [screen_map[screen_id][TB_Sensitivity_index] for screen_id in art_screen_selected[TB_ARTSevereIllness]]
    art_serious_specificity = [screen_map[screen_id][TB_Specificity_index] for screen_id in art_screen_selected[TB_ARTSevereIllness]]
    art_sensitivity = [ art_not_serious_sensitivity, art_serious_sensitivity]
    art_specificity = [ art_not_serious_specificity, art_serious_specificity]

    hrg_screen_selected = TB_screening_alg_hrg_defaults()
    hrg_sensitivity = [screen_map[screen_id][TB_Sensitivity_index] for screen_id in hrg_screen_selected]
    hrg_specificity = [screen_map[screen_id][TB_Specificity_index] for screen_id in hrg_screen_selected]

    incidence = np.zeros(year_len)
    incidence[tb_slice] = fort_outputs['iMid'][fort_slice]
    # flatline no longer needed, taking all of forts values
    # incidence[TB_Fort_Final_Year-base_year:] = incidence[TB_Fort_Final_Year-base_year] # flatline incidence from 2030 onwards
    
    mortality = np.zeros(year_len)
    mortality[tb_slice] = fort_outputs['mMid'][fort_slice]

    # prevalence = calc_prevalence_from_notif_inci(notif, incidence)
    prevalence = np.zeros(year_len)
    prevalence[tb_slice] = fort_outputs['pMid'][fort_slice]
    prev_pct = prevalence/db['projections']['pop'][first_idx:final_idx+1]

    if GB_file_exists(connection, 'tuberculosis', f'dynamicalmodel\\countries\\{countryID}_{TB_DYNAMICAL_MODEL_VERSION}.JSON'):
        dyn_db =  GB_get_db_json(connection, 'tuberculosis', f'dynamicalmodel\\countries\\{countryID}_{TB_DYNAMICAL_MODEL_VERSION}.JSON')
    else:
        dyn_db =  GB_get_db_json(connection, 'tuberculosis', f'dynamicalmodel\\countries\\NGA_{TB_DYNAMICAL_MODEL_VERSION}.JSON')
    # p = dyn_db['p']
    # r = dyn_db['r']
    # dyn_db['prm']['r'] = r
    # dyn_db['prm']['p'] = p
    prm = dyn_db['prm']
    # p = {		
    #     'rfbeta_mdr'  : 0,
    #     'rfbeta_hiv'  : 0,
    #     'imm'         : np.zeros((TB_DynHIVARTLen)).tolist(),
    #     'HIV_relrate' : 0, 
    #     'pu'          : 0,
    #     'Dx'          : np.zeros((TB_DynSectorLen)).tolist(),
    #     'Tx_init'     : np.zeros((TB_DynSectorLen)).tolist(),
    #     'TX_complete' : np.zeros((TB_DynSectorLen)).tolist(),
    #     'cure'        : np.zeros((TB_DynSectorLen)).tolist(),
    #     'MDR_rec2015' : np.zeros((TB_DynSectorLen)).tolist(),
    #     'MDR_rec2022' : np.zeros((TB_DynSectorLen)).tolist(),
    #     'Tx_init2'    : np.zeros((TB_DynSectorLen)).tolist(),
    #     'SL_trans'    : np.zeros((TB_DynSectorLen)).tolist(), 
    #     'tsrsl'       : np.zeros((TB_DynSectorLen)).tolist(),
    #     'cure2'       : np.zeros((TB_DynSectorLen)).tolist(),
    #     'PT_PLHIV'    : np.zeros((TB_DynHIVARTLen)).tolist(),
    #     'VE'          : np.zeros((TB_DynExposureLen)).tolist(),    
    # }
    # r = {
    #     'beta'         : 0,
    #     'reactivation' :  np.zeros((TB_DynAgeLen, TB_DynHIVARTLen)).tolist(),
    #     'progression'  :  np.zeros((TB_DynAgeLen, TB_DynHIVARTLen)).tolist(),
    #     'LTBI_stabil'  : np.zeros((TB_DynHIVARTLen)).tolist(),
    #     'relapse'      : np.zeros((TB_DynRelapseLen)).tolist(),
    #     'self_cure'    : np.zeros((TB_DynHIVARTLen)).tolist(),
    #     'mort_TB'      : np.zeros((TB_DynHIVLen)).tolist(),
    #     'HIV_mort'     : 0,
    #     'sym'          : 0,
    #     'cs'           :  np.zeros((TB_DynAgeLen)).tolist(),
    #     'cs2'          :  np.zeros((TB_DynAgeLen)).tolist(),
    #     'cs3'          :  np.zeros((TB_DynAgeLen)).tolist(),
    #     'Dx'           : 0,
    #     'Tx'           : 0,
    #     'Tx2'          : 0,
    #     'default'      : np.zeros((TB_DynSectorLen)).tolist(),
    #     'MDR_acqu'     : 0,
    #     'default2'     : np.zeros((TB_DynSectorLen)).tolist(),
    #     # HIV related
    #     'ART_init'     : 0,
    #     # Demographics
    #     'mort'         : np.zeros((TB_DynAgeLen)).tolist(),
    #     'ageing'        : np.zeros((TB_DynAgeLen)).tolist(),
    #     'vacc'         : 0,
    #     'waning'       : 0, 
    # }
    used_HIV_incd=0
    # prm = {
    #     'cm': np.zeros((TB_DynAgeLen,TB_DynAgeLen)).tolist(),
    #     'rHIV': 0,
    #     'ART_start':0,
    #     'r': r, 
    #     'p': p,
    #     'usedHIVIncd' : used_HIV_incd
    # }
    # prm = TB_PRM_Default_Nigeria
    prm['usedHIVIncd' ] = used_HIV_incd
    
    r_include = {
                'beta'         : True,
                'reactivation' : False,
                'progression'  : False,
                'LTBI_stabil'  : False,
                'relapse'      : False,
                'self_cure'    : True,
                'mort_TB'      : True,
                'HIV_mort'     : True,
                'sym'          : True,
                'cs'           : True,
                'cs2'          : True,
                'cs3'          : True,
                'Dx'           : False,
                'Tx'           : False,
                'Tx2'          : False,
                'default'      : False,
                'MDR_acqu'     : True,
                'default2'     : False,
                # HIV related
                'ART_init'     : True,
                # Demographics
                'mort'         : False,
                'ageing'       : False,
                'vacc'         : False,
                'waning'       : False, 
    }
    p_include = {		
        'rfbeta_mdr'  : True,
        'rfbeta_hiv'  : True,
        'imm'         : False,
        'HIV_relrate' : True,
        'pu'          : True,
        'Dx'          : True,
        'Tx_init'     : False,
        'TX_complete' : True,
        'cure'        : False,
        'MDR_rec2015'  : False,
        'MDR_rec2022'  : True,
        'Tx_init2'    : False,
        'SL_trans'    : False,
        'tsrsl'       : False,
        'cure2'       : False,
        'PT_PLHIV'    : False,
        'VE'          : False,
    }
    used_HIV_incd_include = False
    prm_include = {
                'cm': False,
                'rHIV': False,
                'ART_start':False,
                'r' : r_include,
                'p' : p_include,
                'UsedHIVIncd' : used_HIV_incd_include
    }
    
    r_range = {
        'beta'         : [0, 40],
        'reactivation' : [],
        'progression'  : [],
        'LTBI_stabil'  : [],
        'relapse'      : [],
        'self_cure'    : [0.85/6, 1.15/6], # 1/6*[0.85 - 1.15]
        'mort_TB'      : [[0, 2/6], [0, 100/6]],  
        'HIV_mort'     : [0, 10],
        'sym'          : [0.1, 100.0],
        'cs'           : [0.1, 100.0],
        'cs2'          : [1, 24],
        'cs3'          : [1, 24],
        'Dx'           : [],
        'Tx'           : [],
        'Tx2'          : [],
        'default'      : [],
        'MDR_acqu'     : [0.0, 0.06],
        'default2'     : [],
        # HIV related
        'ART_init'     : [0, 10],
        # Demographics
        'mort'         : [],
        'ageing'       : [],
        'vacc'         : [],
        'waning'       : [],
    }
    used_HIV_incd_range = []
    p_range = {		
        'rfbeta_mdr'  : [0.0, 1.0],
        'rfbeta_hiv'  : [0.0, 1.0],
        'imm'         : [],
        'HIV_relrate' : [1, 100],
        'pu'          : [],
        'Dx'          : [[0.75, 0.9], [0.1, 0.3]],
        'Tx_init'     : [],
        'TX_complete' : [[0.75, 0.95], [0.4, 0.8]],
        'cure'        : [],
        'MDR_rec2015' : [],
        'MDR_rec2022' : [[0, 1], [0, 1]],
        'Tx_init2'    : [],
        'SL_trans'    : [],
        'tsrsl'       : [],
        'cure2'       : [],
        'PT_PLHIV'    : [],
        'VE'          : [],
    }
    prm_range = {
        'cm'        : [],
        'rHIV'      : [],
        'ART_start' : [],
        'p':p_range,
        'r':r_range,
        'usedHIVIncd':used_HIV_incd_range
    }
    #modvar creation list
    return {
        TB_BaseYearTag: base_year,
        TB_DataSourceTag: extra['countryDetails']['sources']['TBSource'],
        TB_UseOldFailsafeTag: False, #only true for old 2021 projections 
        TB_UseOldIPTag: False, #only true for old 2022 projections 

        TB_HistNotifStartYearTag   : db['notifications']['startYear'],
        TB_HistNotifTag            : {key:db['notifications'][key] for key in TB_NOTIF_DB_KEYS},
        TB_HistNotiDistributionTag : {key:db['Noti_Distribution'][key] for key in TB_NOTIF_DIST_KEYS},
        
        TB_ScreeningAlgMapTag              : screen_map,

        # From the wireframes
        TB_PulmonaryAgeTag                 : TB_Pulmonary_list,
        TB_PrevalenceTag                   : np.full((TB_Pulmonary_count, year_len), 0.10),
        TB_PatientScreenSensitivityTag     : patient_sensitivity, # np.full((TB_Pulmonary_count), 0.90),
        TB_PatientScreenSpecificityTag     : patient_specificity, # np.full((TB_Pulmonary_count), 0.99),
        TB_PatientScreenAlgOptionsTag      : TB_screening_alg_patient_options(),
        TB_PatientScreenAlgSelectedTag     : patient_screen_selected,

        TB_NumClinicallyAssessedTag        : np.zeros((TB_Pulmonary_count, year_len)),
        TB_NumReferredTag                  : np.zeros((TB_Pulmonary_count, year_len)),
        TB_NumScreenedToReferTag           : np.zeros((TB_Pulmonary_count, year_len)),
        TB_PercTrueCasesTag                : np.zeros((TB_Pulmonary_count, year_len)),

        TB_PatientDiagCaseValuesTag            : diagnostic_db['patients']['values'],
        TB_PatientDiagCaseIDsTag               : diagnostic_db['patients']['IDs'],
        TB_PatientDiagCaseMethodsTag           : diagnostic_db['patients']['methods'],
        TB_PatientDiagCaseSensitivityTag       : diagnostic_db['patients']['sensitivity'],
        TB_PatientDiagCaseSpecificityTag       : diagnostic_db['patients']['specificity'],

        TB_PatientDiagSensitivityTag       : np.full((TB_Pulmonary_count), 0.84), # db['Sensitivity']['Patient'], # 
        TB_PatientDiagSpecificityTag       : np.full((TB_Pulmonary_count), 0.97), # db['Specificity']['Patient'], # 
        TB_NumDiagnosedPTBTag              : np.zeros((TB_Pulmonary_count, year_len)),
        TB_NumConfirmedPTBTag              : np.zeros((TB_Pulmonary_count, year_len)),
        TB_NumDiagnosedTBTag               : np.zeros((TB_Pulmonary_count, year_len)),
        TB_NumInitTreatmentTag             : np.zeros((TB_Pulmonary_count, year_len)),

        TB_PatientScreeningOutputTag       : np.zeros((TB_Pulmonary_count, TB_ScreenOutputLen, year_len)),
        TB_PatientDiagnosticOutputTag      : np.zeros((TB_Pulmonary_count, TB_DiagnosisLen, year_len)),
        TB_PatientTargetPopTag             : np.zeros((TB_Pulmonary_count, TB_PatientTargPopLen, year_len)),
                    
        TB_PatientCustStepMenu : TB_Cust_patient_menu,
        TB_PatientCustStep1 : TB_Cust_patient_default,
        TB_PatientCustStep2 : TB_Cust_patient_default,
        TB_PatientCustParallel : np.full((TB_Pulmonary_count), True).tolist(),
        TB_ScreenCustomTag : {}, #empty dict of custom ids
        TB_DiagnosticCustomTag : {},
        

        # provider households
        
        # Household 
        ## Household demographics
        ### Demographics
        TB_HHSizeTag: db['estimates'].get('e_hh_size', [0])[0],# Average household size (#)
        TB_HHPropU5Tag: 0.15,#  Proportion of household that is u5 (# )
        TB_HHProp5_14Tag: 0.10,# Proportion of household that is 5-14 yrs 
        TB_HHCasesinIndexTag: 1.2,# Average number of TB cases in household with at least one case (#)

        ### TB infection characteristics
        TB_HHPropU5LTBITag: np.full((year_len), 0.35), # Proportion of u5s with TB infection in household with TB infection (# )
        TB_HHPropO5LTBITag: np.full((year_len), 0.45), # Proportion of other household members with TB infection in household with active case (# )

        #### TB disease characteristics
        TB_HHPropU5TBTag:  np.full((year_len), 0.05),  # Proportion of�u5s with TB disease in household with TB disease (# )
        TB_HHPropU5PTBTag:  np.full((year_len), 1-c_ep),  # , where t is final year of WHO data
        TB_HHPropO5TBTag:  np.full((year_len), 0.05),  # Proportion of other household members with TB disease in household with TB disease (# )
        TB_HHPropO5PTBTag: np.full((year_len), 1-c_ep), # Proportion with pulmonary TB, PTB (# )

        ##     Index cases
        TB_HHNumNewBactPosTag : notif*bac_pos_pct, # Number of new bacteriologically positive TB cases (adult), total notification from database 
        TB_HHNumNewBactNegTag : notif*(1-bac_pos_pct),# Number of new bacteriologically negative TB cases (adult),  total notification from database x 1- self.Noti_Data["new_labconf" + "ret_rel_labconf" ][TB_noti_perc][t]

        ## household contact tracing coverage
        TB_HHPIndexCasesScreenedTag: np.full((year_len), 1.00), # Proportion of households of index cases screened/evaluated
        TB_HHPropBACnScreenedTag: np.full((year_len), 0.20),  # Proportion of contacts of new bacteriologically negative cases screened/evaluated
        
        ## Tab Household contacts by age
        TB_HHContacts                : np.zeros((TB_HHAgeLen, TB_ContactLen, year_len)),   
        
        ## Tab Household screening
        ### Household proportion
        TB_HHProportionTag           : np.full((TB_HHAgeLen, year_len), 0.5), 

        ### Household sensitivity amd specificity
        TB_HHScreenSensitivityTag : household_sensitivity, # np.full((TB_HHAgeLen), 0.90),
        TB_HHScreenSpecificityTag : household_specificity, # np.full((TB_HHAgeLen), 0.99),  
        TB_HHScreenAlgOptionsTag  : TB_screening_alg_household_options(),
        TB_HHScreenAlgSelectedTag : household_screen_selected,


        ### Household screening outputs
        TB_HHNumScreenedTag           : np.zeros((TB_HHAgeLen, year_len)),
        TB_HHNumReferredTag           : np.zeros((TB_HHAgeLen, year_len)),
        TB_HHNumScreenedToReferTag    : np.zeros((TB_HHAgeLen, year_len)),
        TB_HHPercTrueCasesTag         : np.zeros((TB_HHAgeLen, year_len)),     

        ## Tab Household diagnostic evaluation
        ### TB Case type split by treatment and sputum status
        TB_HHDiagCaseValuesTag               : diagnostic_db['households']['values'],
        TB_HHDiagCaseIDsTag                  : diagnostic_db['households']['IDs'],
        TB_HHDiagCaseMethodsTag              : diagnostic_db['households']['methods'],
        TB_HHDiagCaseSensitivityTag          : diagnostic_db['households']['sensitivity'],
        TB_HHDiagCaseSpecificityTag          : diagnostic_db['households']['specificity'],

        ### Household sensitivity and specificity

        TB_HHDiagSensitivityTag : np.full((TB_HHAgeLen),0.84), # db['Sensitivity']['HouseHold'], # np.full((TB_HHAgeLen), 0.90),
        TB_HHDiagSpecificityTag : np.full((TB_HHAgeLen),0.97), # db['Specificity']['HouseHold'], # np.full((TB_HHAgeLen), 0.99),  

        ### Diagnosis outputs
        TB_HHNumDiagnosedPTBTag       : np.zeros((TB_HHAgeLen, year_len)),             
        TB_HHNumConfirmedPTBTag       : np.zeros((TB_HHAgeLen, year_len)),  

        ## Tab Household disease and treatment
        TB_HHPropLinkedTreatTag       : np.full((TB_HHAgeLen, year_len), 1.00),           
        TB_HHNumInitTreatmentTag      : np.zeros((TB_HHAgeLen, year_len)),   
        
        ## Tab Household prevention
        TB_HHPrvntInActiveEligibleTag  : np.zeros((TB_HHAgeLen, year_len)),     
        TB_HHPrvntReceivePTPrevTag    : np.full((TB_HHAgeLen, year_len), 0.00),   
        TB_HHPropPrvntPresumpTag       : TB_HHPropPrvntPresumpDefaults(year_len), # np.full((TB_HHAgeLen, year_len), 0.00),  # 100 only fot children   
        TB_HHPropPrvntTestedTag        : TB_HHPropPrvntTested(year_len),  # np.full((TB_HHAgeLen, year_len), 0),    
        
        TB_HHPrvntMethodsTag           : hh_prev_methods,    
        TB_HHPrvntMethodIDsTag         : hh_prev_method_IDs, 
        TB_HHPropEligiblePrvntLinkTag  : np.full((TB_HHAgeLen, year_len), 1.00),  
        
        TB_HHNumberOnTPTTag         : np.zeros((TB_HHAgeLen, year_len)),  
        TB_HHTPTCompleteRateTag     : np.full((TB_HHAgeLen, year_len), 1.0),
        TB_HHNumberCompleteTPTTag   : np.zeros((TB_HHAgeLen, year_len)),  

        ## Tab Household output
        TB_HHScreeningOutputTag  : np.zeros((TB_HHAgeLen, TB_ScreenOutputLen, year_len)),   
        TB_HHDiagnosticOutputTag : np.zeros((TB_HHAgeLen, TB_DiagnosisLen, year_len)),   
        TB_HHTargetPopsTag       : np.zeros((TB_HHAgeLen, TB_TargetPopLen, year_len)),   
   

        TB_HHCustStepMenu : TB_Cust_HH_menu,
        TB_HHCustStep1 : TB_Cust_HH_default,
        TB_HHCustStep2 : TB_Cust_HH_default,
        TB_HHCustParallel : np.full((TB_HHAgeLen), True).tolist(),
        
        ## Tab ART cohort sizes and TB characteristics
        ### ART cohort sizes
        TB_ARTCohortSizesTag         : TB_ARTCohortsDefaults(pop, base_year, final_year),

        ### TB characteristics
        TB_ARTCharacteristicsTag     : TB_ARTCharacteristicDefaults(c_ep, year_len),
        
        ## Tab ART screening
        ### ART proportion
        TB_ARTProportionTag           : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 0.50), 

        ### ART sensitivity amd specificity
        TB_ARTScreenSensitivityTag : art_sensitivity, # np.full((TB_ARTAgeLen), 0.90),
        TB_ARTScreenSpecificityTag : art_specificity, # np.full((TB_ARTAgeLen), 0.99),  
        TB_ARTScreenAlgOptionsTag  : TB_screening_alg_art_options(),
        TB_ARTScreenAlgSelectedTag : art_screen_selected,

        ### ART screening outputs
        TB_ARTNumScreenedTag           : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),
        TB_ARTNumReferredTag           : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),
        TB_ARTNumScreenedToReferTag    : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),
        TB_ARTPercTrueCasesTag         : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),     

        ## Tab ART diagnostic evaluation
        ### TB Case type split by treatment and sputum status
        TB_ARTDiagCaseValuesTag               : [diagnostic_db['art']['values']]*TB_ARTSevereLen,
        TB_ARTDiagCaseIDsTag                  : [diagnostic_db['art']['IDs']]*TB_ARTSevereLen,
        TB_ARTDiagCaseMethodsTag              : [diagnostic_db['art']['methods']]*TB_ARTSevereLen,
        TB_ARTDiagCaseSensitivityTag          : [diagnostic_db['art']['sensitivity']]*TB_ARTSevereLen,
        TB_ARTDiagCaseSpecificityTag          : [diagnostic_db['art']['specificity']]*TB_ARTSevereLen,

        ### ART sensitivity amd specificity
        TB_ARTDiagSensitivityTag : np.full((TB_ARTSevereLen, TB_ARTAgeLen),0.84), # np.full((TB_ARTAgeLen), 0.90),
        TB_ARTDiagSpecificityTag : np.full((TB_ARTSevereLen, TB_ARTAgeLen),0.97), # np.full((TB_ARTAgeLen), 0.99), 

        ### Diagnosis outputs
        TB_ARTNumDiagnosedPTBTag       : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),             
        TB_ARTNumConfirmedPTBTag       : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),  

        ## Tab ART disease and treatment
        TB_ARTPropLinkedTreatTag       : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 1.00),           
        TB_ARTNumInitTreatmentTag      : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),   
        
        ## Tab ART prevention
        TB_ARTPrvntInActiveEligibleTag  : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),   
        TB_ARTPrvntReceivePTPrevTag     : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 0.00),     
        TB_ARTPropPrvntPresumpTag       : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 0.50),     
        TB_ARTPropPrvntTestedTag        : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 0),    
        
        TB_ARTPrvntMethodsTag           : ART_prev_methods,    
        TB_ARTPrvntMethodIDsTag         : ART_prev_method_IDs, 
        TB_ARTPropEligiblePrvntLinkTag  : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 1.00),  
        
        TB_ARTNumberOnTPTTag         : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),  
        TB_ARTTPTCompleteRateTag     : np.full((TB_ARTSevereLen, TB_ARTAgeLen, year_len), 1.0),
        TB_ARTNumberCompleteTPTTag   : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, year_len)),  
        
        ## Tab ART output
        TB_ARTScreeningOutputTag  : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, TB_ScreenOutputLen, year_len)),   
        TB_ARTDiagnosticOutputTag : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, TB_DiagnosisLen, year_len)),   
        TB_ARTTargetPopsTag             : np.zeros((TB_ARTSevereLen, TB_ARTAgeLen, TB_TargetPopLen, year_len)), 
        
        TB_ARTCustStepMenu : TB_Cust_ART_menu,
        TB_ARTCustStep1 : TB_Cust_ART_default,
        TB_ARTCustStep2 : TB_Cust_ART_default,
        TB_ARTCustParallel : np.full((TB_ARTSevereLen, TB_ARTAgeLen), True).tolist(),  

        ## Tab HRG demographics
        TB_HRGDemographicsTag      : TB_HighRiskGroupDemographicDefaults(pop_15_plus, c_ep, db['projections']['prevalencePct'][first_idx:final_idx+1], year_len), 
        ## Tab HRG screening
        ### HRG proportion
        TB_HRGProportionTag        : np.full((TB_HRG_Len, year_len), 0.00), 

        ### HRG sensitivity amd specificity
        TB_HRGScreenSensitivityTag : hrg_sensitivity,
        TB_HRGScreenSpecificityTag : hrg_specificity, 
        TB_HRGScreenAlgOptionsTag  : TB_screening_alg_hrg_options(),
        TB_HRGScreenAlgSelectedTag : hrg_screen_selected, 

        ### HRG screening outputs
        TB_HRGNumScreenedTag           : np.full((TB_HRG_Len, year_len), 12.3),
        TB_HRGNumReferredTag           : np.full((TB_HRG_Len, year_len), 1.2),
        TB_HRGNumScreenedToReferTag    : np.full((TB_HRG_Len, year_len), 11.1),
        TB_HRGPercTrueCasesTag         : np.full((TB_HRG_Len, year_len), 0.90),     

        ## Tab HRG diagnostic evaluation
        ### TB Case type split by treatment and sputum status
        TB_HRGDiagCaseValuesTag               : [diagnostic_db['highriskgroups']['values'][0]]*TB_HRG_Len,
        TB_HRGDiagCaseIDsTag                  : [diagnostic_db['highriskgroups']['IDs'][0]]*TB_HRG_Len,
        TB_HRGDiagCaseMethodsTag              : [diagnostic_db['highriskgroups']['methods'][0]]*TB_HRG_Len,
        TB_HRGDiagCaseSensitivityTag          : [diagnostic_db['highriskgroups']['sensitivity'][0]]*TB_HRG_Len,
        TB_HRGDiagCaseSpecificityTag          : [diagnostic_db['highriskgroups']['specificity'][0]]*TB_HRG_Len,

        ### HRG sensitivity amd specificity
        TB_HRGDiagSensitivityTag : np.full((TB_HRG_Len),0.84), # np.full((TB_HRG_Len),  db['Sensitivity']['HRG'][0]),
        TB_HRGDiagSpecificityTag : np.full((TB_HRG_Len),0.97), # np.full((TB_HRG_Len),  db['Specificity']['HRG'][0]),  

        ### Diagnosis outputs
        TB_HRGNumDiagnosedPTBTag       : np.full((TB_HRG_Len, year_len), 1.3),             
        TB_HRGNumConfirmedPTBTag       : np.full((TB_HRG_Len, year_len), 39.3),  

        ## Tab HRG disease and treatment
        TB_HRGPropLinkedTreatTag       : np.full((TB_HRG_Len, year_len), 1.00),           
        TB_HRGNumInitTreatmentTag      : np.full((TB_HRG_Len, year_len), 1.3),   
        
        ## Tab HRG prevention
        TB_HRGPrvntInActiveEligibleTag  : np.full((TB_HRG_Len, year_len), 0.497),   
        TB_HRGPrvntReceivePTPrevTag     : np.full((TB_HRG_Len, year_len), 0.00),   
        TB_HRGPropPrvntPresumpTag       : np.full((TB_HRG_Len, year_len), 0.00),     
        TB_HRGPropPrvntTestedTag        : np.full((TB_HRG_Len, year_len), 1),    
        
        TB_HRGPrvntMethodsTag           : HRG_prev_methods,    
        TB_HRGPrvntMethodIDsTag         : HRG_prev_method_IDs, 
        TB_HRGPropEligiblePrvntLinkTag  : np.full((TB_HRG_Len, year_len), 1.00),  
        
        TB_HRGNumberOnTPTTag         : np.zeros((TB_HRG_Len, year_len)),  
        TB_HRGTPTCompleteRateTag     : np.full((TB_HRG_Len,  year_len), 1.0),
        TB_HRGNumberCompleteTPTTag   : np.zeros((TB_HRG_Len, year_len)), 
        
        ## Tab HRG output
        TB_HRGScreeningOutputTag       : np.zeros((TB_HRG_Len, TB_ScreenOutputLen, year_len)),   
        TB_HRGDiagnosticOutputTag      : np.zeros((TB_HRG_Len, TB_DiagnosisLen, year_len)), 
        TB_HRGTargetPopsTag            : np.zeros((TB_HRG_Len, TB_TargetPopLen, year_len)), 

        TB_HRGCustStepMenu : TB_Cust_HRG_menu,
        TB_HRGCustStep1 : TB_Cust_HRG_default,
        TB_HRGCustStep2 : TB_Cust_HRG_default,
        TB_HRGCustParallel : np.full((TB_HRG_Len), True).tolist(), 


        TB_ResistProfileTag  : rr_profile,
        TB_ResistPropNonSevereTag  : np.full((TB_NewRetreatLen, TB_ResistAgeLen, year_len), 0.95),
        TB_ResistPropLungSurgeryTag  : np.full((TB_NewRetreatLen, TB_ResistAgeLen, year_len), 0.05),
        TB_ResistMeningitisTag  : np.full((TB_NewRetreatLen, TB_ResistAgeLen, year_len), 0.05),
        TB_ResistPropTreatExhaustTag  : np.full((TB_NewRetreatLen, TB_ResistAgeLen, year_len), 0.05),

        TB_PatientTestNumPTB  : np.full((year_len), 95),
        TB_PatientTestNumXDR  : np.full((year_len), 5),
        TB_PatientTestNumEPTB : np.full((year_len), 10), 

        TB_CoverageTag : TB_ResistanceCoverageDefaults(year_len),
        TB_ResistNumToBeTestedTag : np.zeros((TB_NumTestLen, year_len)), 
        TB_ResistNumTestedTag : np.zeros((TB_NumTestLen, year_len)), 
        TB_ResistNumConfirmedTag : np.full((TB_NumConfResistLen, TB_ResistAgeLen, year_len), 50), 
        TB_ResistNumNotConfirmedTag : np.full((TB_NumConfResistLen, TB_ResistAgeLen, year_len), 0), 


        ## Prevention of TB
        TB_PropPreventionTag : TB_PreventionDefaults(rp_Rif ,year_len),

        ## Vaccination
        TB_BCGImmunizationTag : np.full((year_len), 1.00),
        TB_InfantVaccineTag : np.full((year_len), 1.00),
        TB_AdolAdultVaccineTag : np.full((year_len), 1.00),

        ## Palliative care
        TB_PalliativeCareTag : np.zeros((TB_ResistAgeLen, year_len)),

        TB_BCGTargetPopTag          : np.full((year_len), 0),
        TB_InfantVaccTargPopTag     : np.full((year_len), 0),
        TB_AdultVaccTargPopTag      : np.full((year_len), 0),
        TB_PalliativeCareTargPopTag : np.full((TB_ResistAgeLen, year_len), 0),




        #Extra target populations
        TB_PreventionTargetTag : np.zeros((TB_ResistanceLen, TB_ResistAgeLen, year_len)),
        
        TB_RR_DST_PTBTag     : np.zeros(year_len),
        TB_RR_DST_ETBTag     : np.zeros(year_len),
        TB_INH_DST_PTBTag    : np.zeros(year_len),
        TB_INH_DST_ETBTag    : np.zeros(year_len),
        TB_BDQ_DST_PTBTag    : np.zeros(year_len),
        TB_BDQ_DST_ETBTag    : np.zeros(year_len),
        TB_FQ_DST_PTBTag     : np.zeros(year_len),
        TB_FQ_DST_ETBTag     : np.zeros(year_len),
        TB_BDQ_FQ_DST_PTBTag : np.zeros(year_len),
        TB_BDQ_FQ_DST_ETBTag : np.zeros(year_len),
        TB_LZD_RR_DST_PTBTag : np.zeros(year_len),
        TB_LZD_RR_DST_ETBTag : np.zeros(year_len),
        TB_LZD_FQ_DST_PTBTag : np.zeros(year_len),
        TB_LZD_FQ_DST_ETBTag : np.zeros(year_len),
        TB_PZA_DSTTag        : np.zeros(year_len),
        TB_ETO_DSTTag        : np.zeros(year_len),
        TB_EMC_DSTTag        : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_U15Tag         : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_15PLusTag      : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_PTB_U15Tag     : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_PTB_15PLusTag  : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_ETB_U15Tag     : np.zeros(year_len),
        TB_BDQ_LZD_NumDx_ETB_15PLusTag  : np.zeros(year_len),



        TB_AmikacynRifampicinResistant     : np.zeros((TB_PulLen, year_len)), 
        TB_EthionamideRifampicinResistant  : np.zeros((TB_PulLen, year_len)), 
        TB_PyrazinamideRifampicinResistant : np.zeros((TB_PulLen, year_len)), 

        TB_Rif_INH_Sen_SevereTag     : np.zeros((TB_ChildAdultLen, year_len)),
        TB_Rif_INH_Sen_NotSevereTag  : np.zeros((TB_ChildAdultLen, year_len)),
        TB_Rif_INH_Sen_MeningitisTag : np.zeros((TB_ChildAdultLen, year_len)),
        TB_Rif_INH_ResTag            : np.zeros((TB_ChildAdultLen, year_len)),
        TB_RR_FQ_SenTag              : np.zeros((TB_ChildAdultLen, year_len)),
        TB_Rif_FQ_Res_PlusTag        : np.zeros((TB_ChildAdultLen, year_len)),
        TB_BQD_LZD_Tag               : np.zeros((TB_ChildAdultLen, year_len)),
        TB_TxFL_Tag                  : np.zeros((TB_ChildAdultLen, year_len)),
        TB_TxSL_Tag                  : np.zeros((TB_ChildAdultLen, year_len)),
        TB_TxPreXDR_Tag              : np.zeros((TB_ChildAdultLen, year_len)),
        TB_LungSurgeryTag            : np.zeros((TB_ChildAdultLen, year_len)),
        

        TB_lPTB_ReferredDxTag           : np.zeros((TB_ChildAdultLen, year_len)),

        TB_HIVnClinicallyAssessedTag : np.zeros((TB_PulLen, TB_ChildAdultLen, year_len)),
        TB_HIVAnyReferredDxTag       : np.zeros((TB_PulLen, TB_ChildAdultLen, year_len)),

        # Notification
        TB_HisNotificationTag  : notif,
        TB_TargNotificationTag : notif,
        TB_ResNotificationTag  : np.array([notif, notif, notif*0]),

        # TB_ResPrevalenceTag    : ,# might come from pete dodds model,# might come from pete dodds model
        TB_ResHIVTBCoInfTag    : np.zeros((TB_HIVCoInfLen, year_len)),

        TB_HisIncidenceTag  : db['burden']['e_inc_num'][burden_first_idx:burden_final_idx+1],
        TB_HisMortalityTag  : db['burden']['e_mort_num'][burden_first_idx:burden_final_idx+1],
        TB_RetNrelPercTag   : db['Noti_Distribution']['ret_nrel_perc'],
        TB_ResIncidenceTag  : np.array([incidence, incidence, incidence*0]), 
        TB_ResMortalityTag  : np.array([mortality, mortality, mortality*0]), 
        TB_ResPrevalenceTag : np.array([prevalence, prevalence, prevalence*0]), 
        TB_UserEnteredIncidenceTag  : incidence,
        TB_UserEnteredMortalityTag  : mortality,
        TB_AlternativeProjBurdenTag : False,

        TB_FortIncidenceTag    : np.zeros((TB_ResultBounds, year_len)),
        TB_FortNotificationTag : np.zeros((TB_ResultBounds, year_len)),
        TB_FortPrevalenceTag   : np.zeros((TB_ResultBounds, year_len)),
        TB_FortMortalityTag    : np.zeros((TB_ResultBounds, year_len)),
        
        TB_SubnationalRatioTag  : 1.0, # 1.0 means it is a national model
        

        TB_ResPrevalencePctTag : db['projections']['prevalencePct'][first_idx:final_idx+1],

        TB_CounterfactualScenarioTag :  np.array([
                                    [0.76]*year_len,
                                    [0.00]*year_len,
                                    [0.84]*year_len,
                                ]
                                ),
        TB_CurrentScenarioTag        :  np.array([
                                    [0.76]*year_len,
                                    [0.30]*year_len,
                                    [0.84]*year_len,
                                ]
                                ),
                                
        TB_CounterfactualNotifTag: notif,
        
        TB_Baseline_HRi_allTag : np.full(year_len, 0.00),
        TB_Scaleup_HRi_allTag : np.full(year_len, 0.30),

        TB_NotificationByCaseTypeTag :  notif_split, # np.full((TB_NotificationCaseCount, year_len), 0),
        TB_PercNotificationTag       :  notif_pct,
        
        TB_FortVsnTag : '',
    
        TB_DynRatesTag : prm['r'],
        
        TB_DynPRMTag : prm,
        TB_DynUsedHIVIncdTag : used_HIV_incd,
        TB_DynProportionsTag: prm['p'],

        TB_DynUsedFitParamTag : {
            TB_DynRatesTag : r_include,
            TB_DynPRMTag : prm_include,
            TB_DynUsedHIVIncdTag : used_HIV_incd_include,
            TB_DynProportionsTag: p_include,
        },
        TB_DynRangesTag: {
            TB_DynRatesTag : r_range,
            TB_DynPRMTag : prm_range,
            TB_DynUsedHIVIncdTag : used_HIV_incd_range,
            TB_DynProportionsTag : p_range,
        },
                
        # TB_DynXSTag       : dyn_db['xs'],
        TB_DynPopDataTag  : dyn_db['data'],
        TB_DynamicalFirstYearTag : first_year,
        TB_DynamicalFinalYearTag : final_year,
        
        TB_FinalHistoricalYearTag: TB_FinalHist_Year, 
        TB_FortInputVersionTag: TB_FORT_INPUT_VERSION,
        TB_FortInputTag : fort_inputs,
        TB_FortInputScaleupTag : fort_inputs, # Default to same as inputs
        TB_FortOutputVersionTag : TB_FORT_OUTPUT_VERSION,
        TB_FortOutputTag: fort_outputs,

    }

# Helper function that will flatline the values from the database
# to the final year of the projection
def TB_flatlineDB(db, name, final_idx):
    if final_idx >= len(db['projections'][name]):#'notification'
        for i in range(len(db['projections'][name]), final_idx+1):
             db['projections'][name].append(db['projections'][name][-1])
             
def TB_flatlinBurden(db, name, final_idx):
    if final_idx >= len(db['burden'][name]):#'notification'
        for i in range(len(db['burden'][name]), final_idx+1):
             db['burden'][name].append(db['burden'][name][-1])


def TB_getPopDefaults(connection, countryID, subnatCode, base_year, final_year):    
    pop1CountryName = formatCountryFName(countryID, DPInitialConditionsDBVersion)#, subnatCode)
    #pop1CountryName = formatCountryFName(countryID, DPUPDVersion)#, subnatCode) 
    if GB_file_exists(connection, 'demproj', DPInitialConditionsDBSubDir + pop1CountryName):
        pop1 = np.array(GB_get_db_json(connection, 'demproj', DPInitialConditionsDBSubDir + pop1CountryName)['pop1']) 
        yr_slice = slice(base_year-GB_BaseYear, final_year-GB_BaseYear+1)

        pop = pop1[yr_slice] 
        pop_15_plus = pop[:, :, 15:, :, :].sum()
        return pop, pop_15_plus
    else:
        raise Exception('Country does not exist for TB Pop')
    

def calc_prevalence_from_notif_inci(notification, incidence):
    tx_mid = (2+0.2)/2
    ut_mid = (4+1)/2
    return  notification*tx_mid + (incidence - notification)*ut_mid
    