from SpectrumCommon.Const.VW.VWConst import *
from SpectrumCommon.Const.DP import *
from SpectrumCommon.Const.AM import *
from SpectrumCommon.Const.FP import *
from SpectrumCommon.Const.CS import *
from AvenirCommon.Util import GBRange

def VW_TotalPop_Init(projection, SubnatData_DB):
    result = np.zeros((VW_Subnat+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        result[i] = 0
        for s in GBRange(VW_Male, VW_Female):
            for a in GBRange(0, 80):
                result[i] += projection[DP_BigPopTag][s][a][0]

    if 'PercPop' in SubnatData_DB:
        result[VW_Subnat] *= float(SubnatData_DB['PercPop'])

    return result

def VW_TFR_Init(projection, SubnatData_DB, source):
    result = np.zeros((VW_Subnat+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        result[i] = projection[DP_TFRTag][0]

    source = SubnatData_DB[source]
    
    if str(CS_MstTFR) in source:
        record = source[str(CS_MstTFR)]
        if 'value' in record:
            tfr_rate = float(record['value'])
            result[VW_Subnat] = min(tfr_rate * result[VW_National], 100)

    return result

def VW_CPR_Init(projection, SubnatData_DB, source):
    result = np.zeros((VW_Subnat+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        result[i] = projection[FP_PrevalenceTag][FP_All_Ages][FP_all_need][0]

    source = SubnatData_DB[source]
    
    if str(CS_MstCPR) in source:
        record = source[str(CS_MstCPR)]
        if 'value' in record:
            cpr_rate = float(record['value'])
            result[VW_Subnat] = min(cpr_rate * result[VW_National], 100)

    return result

def VW_Incidence_Init(projection):
    result = np.zeros((VW_Subnat+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        IncidenceOptions = projection[AM_IncidenceOptionsTag]
        result[i] = projection[AM_IncidenceByFitTag][IncidenceOptions][0]

    return result

def VW_PMTCT_Init(projection, t):
    result = np.zeros((VW_Subnat+1, VW_PostnatalProphylaxis+1)).tolist()
    
    ARVRegimen = projection[AM_ARVRegimenTag]
    
    for i in GBRange(VW_National, VW_Subnat):
        
        value = 0
        for r in GBRange(DP_SingleDoseNevir, DP_TripleARTDurPreg_Late):
            value += ARVRegimen[DP_PrenatalProphylaxis][r][DP_Percent][t]
        result[i][VW_PrenatalProphylaxis] = value
        
        value = 0
        for r in GBRange(DP_OptA, DP_OptB):
            value += ARVRegimen[DP_PostnatalProphylaxis][r][DP_Percent][t]
        result[i][VW_PostnatalProphylaxis] = value

    return result

def VW_Cotrim_Init(projection, t):
    result = np.zeros((VW_Subnat+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        result[i] = projection[AM_ChildTreatInputsTag][DP_PerChildHIVPosCot][t]

    return result

def VW_AdultART_Init(projection, t):
    result = np.zeros((VW_Subnat+1, VW_Female+1)).tolist()
    
    for i in GBRange(VW_National, VW_Subnat):
        for s in GBRange(VW_Male, VW_Female):
            result[i][s] = projection[AM_HAARTBySexTag][s][t]

    return result