import numpy as np

from AvenirCommon.Util import GBRange
import SpectrumCommon.Const.TB as tbc
import SpectrumCommon.External.TB.TBFortWrapper as fort



def create_notification_defaults_2022(ISO3, db_notification, Noti_Distribution, first_idx):

    notif_base = db_notification[first_idx]
    
    notif_base_split =  np.zeros((tbc.TB_NotificationCaseCount))
    notif_base_pct = np.zeros((tbc.TB_NotificationCaseCount))
    
    new_cases = Noti_Distribution['new_labconf_num']+Noti_Distribution['new_clindx_num']+Noti_Distribution['new_ep_num'] 
    relapse_cases = Noti_Distribution['ret_rel_labconf_num']+Noti_Distribution['ret_rel_clindx_num']+Noti_Distribution['ret_rel_ep_num']
    ret_nrel_cases = Noti_Distribution['ret_nrel_num'] 
    
    c_ep = (Noti_Distribution['new_ep_num']+Noti_Distribution['ret_rel_ep_num'])/(new_cases+relapse_cases) if (new_cases+relapse_cases)>0 else 0
    bac_pos_pct = (Noti_Distribution['new_labconf_num']+Noti_Distribution['ret_rel_labconf_num'])/(new_cases + relapse_cases)if (new_cases+relapse_cases)>0 else 0
    
    total_cases =  new_cases + relapse_cases + ret_nrel_cases

    notif_scale = total_cases/notif_base
    notif_base = notif_base*notif_scale
    
    # fort_outputs =  get_tb_fort_outputs(ISO3, TB_FinalHist_Year, notif[0], fort_vsn)
    # notif[tbc.TB_slice] = fort_outputs['nMid'][fort_slice]

    notif_HIV           = (total_cases)*Noti_Distribution['newrel_hivpos_perc']
    notif_HIVpNeonates  = notif_HIV*Noti_Distribution['newrel_014_perc']/15
    notif_HIVp0_14      = notif_HIV*Noti_Distribution['newrel_014_perc']-notif_HIVpNeonates
    notif_HIVp15plus    = notif_HIV*Noti_Distribution['newrel_15plus_perc']


    notif_NotOnART      = notif_HIV*(1-Noti_Distribution['newrel_art_perc'])
    notif_OnART         = notif_HIV-notif_NotOnART
    
    notif_OnARTNeoNates = notif_OnART*Noti_Distribution['newrel_014_perc']/15
    notif_OnART0_14     = notif_OnART*Noti_Distribution['newrel_014_perc']-notif_OnARTNeoNates
    notif_OnART15plus   = notif_OnART*Noti_Distribution['newrel_15plus_perc']
    
    notif_base_split[tbc.TB_NewCases                    ] = new_cases
    notif_base_split[tbc.TB_NewBacterConfirmCases       ] = Noti_Distribution['new_labconf_num']
    notif_base_split[tbc.TB_NewClinicDiagCases          ] = Noti_Distribution['new_clindx_num']
    notif_base_split[tbc.TB_NewExtrapulmonaryCases      ] = Noti_Distribution['new_ep_num']
    notif_base_split[tbc.TB_RelapseCases                ] = relapse_cases
    notif_base_split[tbc.TB_RelapseBacterConfirmCases   ] = Noti_Distribution['ret_rel_labconf_num']
    notif_base_split[tbc.TB_RelapseClinicDiagCases      ] = Noti_Distribution['ret_rel_clindx_num']
    notif_base_split[tbc.TB_RelapseExtrapulmonaryCases  ] = Noti_Distribution['ret_rel_ep_num']
    notif_base_split[tbc.TB_NewAndRelapsedCases  ]        = new_cases+ relapse_cases
    
    notif_base_split[tbc.TB_PrevTreatExcludeRelapseCases] = ret_nrel_cases
    notif_base_split[tbc.TB_PrevTreatIncludeRelapseCases] = relapse_cases+ret_nrel_cases
    notif_base_split[tbc.TB_NewRelapse0_9Cases  ] = Noti_Distribution['newrel_04_num']+Noti_Distribution['newrel_59_num']
    notif_base_split[tbc.TB_NewRelapse10_14Cases] = Noti_Distribution['newrel_1014_num']
    notif_base_split[tbc.TB_NewRelapse15pCases  ] = Noti_Distribution['newrel_15plus_num']
    notif_base_split[tbc.TB_HIVpCases        ] = notif_HIV 
    notif_base_split[tbc.TB_HIVp15pCases     ] = notif_HIVp15plus
    notif_base_split[tbc.TB_HIVp0_14Cases    ] = notif_HIVp0_14
    notif_base_split[tbc.TB_HIVpNeonateCases ] = notif_HIVpNeonates
    notif_base_split[tbc.TB_NotOnARTCases    ] = notif_NotOnART
    notif_base_split[tbc.TB_OnARTCases       ] = notif_OnART
    notif_base_split[tbc.TB_OnART15pCases    ] = notif_OnARTNeoNates
    notif_base_split[tbc.TB_OnARRT0_14Cases  ] = notif_OnART0_14
    notif_base_split[tbc.TB_OnARTNeonateCases] = notif_OnARTNeoNates
    notif_base_split[tbc.TB_HIVpExcludeART] = 0

    notif_base_pct[tbc.TB_NewCases                    ] = notif_base_split[tbc.TB_NewCases                    ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewBacterConfirmCases       ] = notif_base_split[tbc.TB_NewBacterConfirmCases       ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewClinicDiagCases          ] = notif_base_split[tbc.TB_NewClinicDiagCases          ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewExtrapulmonaryCases      ] = notif_base_split[tbc.TB_NewExtrapulmonaryCases      ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_RelapseCases                ] = notif_base_split[tbc.TB_RelapseCases                ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_RelapseBacterConfirmCases   ] = notif_base_split[tbc.TB_RelapseBacterConfirmCases   ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_RelapseClinicDiagCases      ] = notif_base_split[tbc.TB_RelapseClinicDiagCases      ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_RelapseExtrapulmonaryCases  ] = notif_base_split[tbc.TB_RelapseExtrapulmonaryCases  ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewAndRelapsedCases  ]        = notif_base_split[tbc.TB_NewAndRelapsedCases         ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_PrevTreatExcludeRelapseCases] = notif_base_split[tbc.TB_PrevTreatExcludeRelapseCases]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_PrevTreatIncludeRelapseCases] = notif_base_split[tbc.TB_PrevTreatIncludeRelapseCases]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewRelapse0_9Cases  ] = notif_base_split[tbc.TB_NewRelapse0_9Cases  ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewRelapse10_14Cases] = notif_base_split[tbc.TB_NewRelapse10_14Cases]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NewRelapse15pCases  ] = notif_base_split[tbc.TB_NewRelapse15pCases  ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_HIVpCases        ] = notif_base_split[tbc.TB_HIVpCases        ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_HIVp15pCases     ] = notif_base_split[tbc.TB_HIVp15pCases     ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_HIVp0_14Cases    ] = notif_base_split[tbc.TB_HIVp0_14Cases    ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_HIVpNeonateCases ] = notif_base_split[tbc.TB_HIVpNeonateCases ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_NotOnARTCases    ] = notif_base_split[tbc.TB_NotOnARTCases    ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_OnARTCases       ] = notif_base_split[tbc.TB_OnARTCases       ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_OnART15pCases    ] = notif_base_split[tbc.TB_OnART15pCases    ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_OnARRT0_14Cases  ] = notif_base_split[tbc.TB_OnARRT0_14Cases  ]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_OnARTNeonateCases] = notif_base_split[tbc.TB_OnARTNeonateCases]/total_cases if total_cases >0 else 0
    notif_base_pct[tbc.TB_HIVpExcludeART   ] = notif_base_split[tbc.TB_NotOnARTCases    ]/total_cases if total_cases >0 else 0

    # notif_split = apply_notif_split(notif_split, notif, notif_pct)

    return notif_base_pct, notif_scale, c_ep, bac_pos_pct




def create_notification_defaults(ISO3, db_notification, Noti_Distribution, first_year, first_idx, final_idx, year_len, notif_base_pct, tb_slice, fort_slice, fort_vsn, notif_scale, subnat_ratio):
    notif = np.array(db_notification[first_idx:final_idx+1])*notif_scale
    
    notif_split =  np.full((tbc.TB_NotificationCaseCount, year_len), 0)
    notif_pct = np.zeros((tbc.TB_NotificationCaseCount))
    
    fort_outputs =  fort.get_tb_fort_outputs(ISO3, fort_vsn)
    notif[tb_slice] = [v*subnat_ratio for v in fort_outputs['nMid'][fort_slice]]
    last_default_year = tbc.TB_Base_Year if first_year<=tbc.TB_Base_Year else first_year
    for yr in range(first_year, last_default_year+1):
        t = yr-first_year
        total_cases = notif[t]
        
        notif_split[tbc.TB_NewCases                    ][t] = notif_base_pct[tbc.TB_NewCases]*total_cases
        notif_split[tbc.TB_NewBacterConfirmCases       ][t] = notif_base_pct[tbc.TB_NewBacterConfirmCases]*total_cases
        notif_split[tbc.TB_NewClinicDiagCases          ][t] = notif_base_pct[tbc.TB_NewClinicDiagCases]*total_cases
        notif_split[tbc.TB_NewExtrapulmonaryCases      ][t] = notif_base_pct[tbc.TB_NewExtrapulmonaryCases]*total_cases
        notif_split[tbc.TB_RelapseCases                ][t] = notif_base_pct[tbc.TB_RelapseCases]*total_cases
        notif_split[tbc.TB_RelapseBacterConfirmCases   ][t] = notif_base_pct[tbc.TB_RelapseBacterConfirmCases]*total_cases
        notif_split[tbc.TB_RelapseClinicDiagCases      ][t] = notif_base_pct[tbc.TB_RelapseClinicDiagCases]*total_cases
        notif_split[tbc.TB_RelapseExtrapulmonaryCases  ][t] = notif_base_pct[tbc.TB_RelapseExtrapulmonaryCases]*total_cases
        notif_split[tbc.TB_NewAndRelapsedCases  ][t]        = notif_base_pct[tbc.TB_NewAndRelapsedCases]*total_cases
        
        notif_split[tbc.TB_PrevTreatExcludeRelapseCases][t] = notif_base_pct[tbc.TB_PrevTreatExcludeRelapseCases]*total_cases
        notif_split[tbc.TB_PrevTreatIncludeRelapseCases][t] = notif_base_pct[tbc.TB_PrevTreatIncludeRelapseCases]*total_cases
        notif_split[tbc.TB_NewRelapse0_9Cases  ][t] = notif_base_pct[tbc.TB_NewRelapse0_9Cases  ]*total_cases
        notif_split[tbc.TB_NewRelapse10_14Cases][t] = notif_base_pct[tbc.TB_NewRelapse10_14Cases]*total_cases
        notif_split[tbc.TB_NewRelapse15pCases  ][t] = notif_base_pct[tbc.TB_NewRelapse15pCases  ]*total_cases
        notif_split[tbc.TB_HIVpCases        ][t] = notif_base_pct[tbc.TB_HIVpCases        ]*total_cases
        notif_split[tbc.TB_HIVp15pCases     ][t] = notif_base_pct[tbc.TB_HIVp15pCases     ]*total_cases
        notif_split[tbc.TB_HIVp0_14Cases    ][t] = notif_base_pct[tbc.TB_HIVp0_14Cases    ]*total_cases
        notif_split[tbc.TB_HIVpNeonateCases ][t] = notif_base_pct[tbc.TB_HIVpNeonateCases ]*total_cases
        notif_split[tbc.TB_NotOnARTCases    ][t] = notif_base_pct[tbc.TB_NotOnARTCases    ]*total_cases
        notif_split[tbc.TB_OnARTCases       ][t] = notif_base_pct[tbc.TB_OnARTCases       ]*total_cases
        notif_split[tbc.TB_OnART15pCases    ][t] = notif_base_pct[tbc.TB_OnART15pCases    ]*total_cases
        notif_split[tbc.TB_OnARRT0_14Cases  ][t] = notif_base_pct[tbc.TB_OnARRT0_14Cases  ]*total_cases
        notif_split[tbc.TB_OnARTNeonateCases][t] = notif_base_pct[tbc.TB_OnARTNeonateCases]*total_cases
        notif_split[tbc.TB_HIVpExcludeART   ][t] = notif_base_pct[tbc.TB_HIVpExcludeART]*total_cases

        new_cases = notif_split[tbc.TB_NewCases][t] 
        relapse_cases = notif_split[tbc.TB_RelapseCases][t]
        notif_HIV = notif_split[tbc.TB_HIVpCases][t]

        notif_pct[tbc.TB_NewCases                    ] = notif_split[tbc.TB_NewCases                    ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_NewBacterConfirmCases       ] = notif_split[tbc.TB_NewBacterConfirmCases       ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_NewClinicDiagCases          ] = notif_split[tbc.TB_NewClinicDiagCases          ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_NewExtrapulmonaryCases      ] = notif_split[tbc.TB_NewExtrapulmonaryCases      ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_RelapseCases                ] = notif_split[tbc.TB_RelapseCases                ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_RelapseBacterConfirmCases   ] = notif_split[tbc.TB_RelapseBacterConfirmCases   ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_RelapseClinicDiagCases      ] = notif_split[tbc.TB_RelapseClinicDiagCases      ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_RelapseExtrapulmonaryCases  ] = notif_split[tbc.TB_RelapseExtrapulmonaryCases  ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_NewAndRelapsedCases  ]        = notif_split[tbc.TB_NewAndRelapsedCases         ][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_PrevTreatExcludeRelapseCases] = notif_split[tbc.TB_PrevTreatExcludeRelapseCases][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_PrevTreatIncludeRelapseCases] = notif_split[tbc.TB_PrevTreatIncludeRelapseCases][t]/total_cases if total_cases > 0 else 0
        notif_pct[tbc.TB_NewRelapse0_9Cases  ] = notif_split[tbc.TB_NewRelapse0_9Cases  ][t]/(new_cases+relapse_cases) if (new_cases+relapse_cases) > 0 else 0
        notif_pct[tbc.TB_NewRelapse10_14Cases] = notif_split[tbc.TB_NewRelapse10_14Cases][t]/(new_cases+relapse_cases) if (new_cases+relapse_cases) > 0 else 0
        notif_pct[tbc.TB_NewRelapse15pCases  ] = notif_split[tbc.TB_NewRelapse15pCases  ][t]/(new_cases+relapse_cases) if (new_cases+relapse_cases) > 0 else 0
        notif_pct[tbc.TB_HIVpCases        ] = notif_split[tbc.TB_HIVpCases        ][t]/(new_cases+relapse_cases) if (new_cases+relapse_cases) > 0 else 0
        notif_pct[tbc.TB_HIVp15pCases     ] = notif_split[tbc.TB_HIVp15pCases     ][t]/notif_split[tbc.TB_HIVpCases][t] if notif_split[tbc.TB_HIVpCases][t] > 0 else 0
        notif_pct[tbc.TB_HIVp0_14Cases    ] = notif_split[tbc.TB_HIVp0_14Cases    ][t]/notif_split[tbc.TB_HIVpCases][t] if notif_split[tbc.TB_HIVpCases][t] > 0 else 0
        notif_pct[tbc.TB_HIVpNeonateCases ] = notif_split[tbc.TB_HIVpNeonateCases ][t]/notif_split[tbc.TB_HIVpCases][t] if notif_split[tbc.TB_HIVpCases][t] > 0 else 0
        notif_pct[tbc.TB_NotOnARTCases    ] = notif_split[tbc.TB_NotOnARTCases    ][t]/notif_split[tbc.TB_HIVpCases][t] if notif_split[tbc.TB_HIVpCases][t] > 0 else 0
        notif_pct[tbc.TB_OnARTCases       ] = notif_split[tbc.TB_OnARTCases       ][t]/notif_split[tbc.TB_HIVpCases][t] if notif_split[tbc.TB_HIVpCases][t] > 0 else 0
        notif_pct[tbc.TB_OnART15pCases    ] = 0.01 # notif_split[tbc.TB_OnART15pCases    ][0]/notif_split[tbc.TB_OnARTCases][0]
        notif_pct[tbc.TB_OnARRT0_14Cases  ] = 0.01 # notif_split[tbc.TB_OnARRT0_14Cases  ][0]/notif_split[tbc.TB_OnARTCases][0]
        notif_pct[tbc.TB_OnARTNeonateCases] = 0.01 # notif_split[tbc.TB_OnARTNeonateCases][0]/notif_split[tbc.TB_OnARTCases][0]
        notif_pct[tbc.TB_HIVpExcludeART   ] = notif_split[tbc.TB_NotOnARTCases    ][0]/(notif_split[tbc.TB_NotOnARTCases    ][0]+(total_cases-notif_HIV)) if (notif_split[tbc.TB_NotOnARTCases    ][0]+(total_cases-notif_HIV)) > 0 else 0

    notif_split = apply_notif_split(first_year, notif_split, notif, notif_pct)

    return notif, notif_split, notif_pct, fort_outputs

def apply_notif_split(first_year, notif_split, notif, notif_pct):
    base_t = tbc.TB_Base_Year-first_year+1
    for i in GBRange(0, tbc.TB_NotOnARTCases):
        for t in GBRange(base_t, len(notif_split[i])-1):
            notif_split[i][t] = notif[t]*notif_pct[i]

    for i in GBRange(tbc.TB_HIVpCases+1, tbc.TB_HIVpExcludeART):
        for t in GBRange(base_t, len(notif_split[i])-1):
            notif_split[i][t] = notif[t]*notif_pct[i]*notif_pct[tbc.TB_HIVpCases]

    return notif_split
