
import numpy as np


from SpectrumCommon.Const.TB import *
from SpectrumCommon.Const.GB import *
from SpectrumCommon.Const.DP import *

def TB_ARTCohortsDefaults(pop, base_year, final_year):
    

    dur_slice = slice(DP_D_ARTlt6m-1, DP_D_ARTgt12m)
    cd4_slice = slice(DP_CD4_GT500-1, DP_CD4_LT50)
    dim_sum   = (1,2,3,4)
    total = [
            pop[:, : , 10:15, cd4_slice, dur_slice].sum(axis=dim_sum),   # 0-9
            pop[:, : , 10:15, cd4_slice, dur_slice].sum(axis=dim_sum),   # 10-14
            pop[:, : , 15:,       cd4_slice, dur_slice].sum(axis=dim_sum)    # 15 plus 
    ]

    # dur_slice = slice(DP_D_ARTlt6m-1, DP_D_ART6to12m)
    # cd4_slice = slice(DP_CD4_100_199-1, DP_CD4_LT50)
    ART_with_serious_illnes = [
        np.full(len(total[0]), 0.05), #         pop[:, : , 10:15, cd4_slice, dur_slice].sum(axis=dim_sum)/total[TB_ART10_14],   # 0-9
        np.full(len(total[0]), 0.05), #         pop[:, : , 10:15, cd4_slice, dur_slice].sum(axis=dim_sum)/total[TB_ART10_14], # 10-14
        np.full(len(total[0]), 0.05) #         pop[:, : , 15:,       cd4_slice, dur_slice].sum(axis=dim_sum)/total[TB_ART15p]    # 15 plus 
    ]

    # ART_with_serious_illnes = np.full((3, final_year-base_year), 0.05)
    # ART_without_serious_illnes = np.full((3, final_year-base_year), 0.95)
    # cd4_slice = slice(DP_CD4_GT500-1, DP_CD4_200_249)
    ART_without_serious_illnes = [
            1-ART_with_serious_illnes[TB_ART0_9],   # 0-9
            1-ART_with_serious_illnes[TB_ART10_14], # 10-14
            1-ART_with_serious_illnes[TB_ART15p]    # 15 plus
        
    ]
    art_cohort = [[total[a].tolist(), ART_without_serious_illnes[a].tolist(), ART_with_serious_illnes[a].tolist()] for a in range(TB_ART0_9, TB_ART15p+1)]
    
        # pop_15_plus = 0
        # ART_without_serious_illnes = [0]*3
        # ART_with_serious_illnes = [0]*3

    if final_year>2030:
        for yr in range(2030+1, final_year+1):
            art_cohort[0][0].append(art_cohort[0][0][-1])
            art_cohort[0][1].append(art_cohort[0][1][-1])
            art_cohort[0][2].append(art_cohort[0][2][-1])
            art_cohort[1][0].append(art_cohort[1][0][-1])
            art_cohort[1][1].append(art_cohort[1][1][-1])
            art_cohort[1][2].append(art_cohort[1][2][-1])
            art_cohort[2][0].append(art_cohort[2][0][-1])
            art_cohort[2][1].append(art_cohort[2][1][-1])
            art_cohort[2][2].append(art_cohort[2][2][-1])

    return np.nan_to_num(art_cohort, copy=True, nan=0.0, posinf=None, neginf=None)

def TB_ARTCharacteristicDefaults(c_ep, year_len):

    art_characteristics = np.zeros((TB_ARTAgeLen, TB_ART_Charac_len, year_len))   #35, 10, 100-c_ep
    art_characteristics[TB_ART0_9   ][TB_ART_WithTBInfection] += 0.35
    art_characteristics[TB_ART10_14 ][TB_ART_WithTBInfection] += 0.35
    art_characteristics[TB_ART15p   ][TB_ART_WithTBInfection] += 0.35  

    art_characteristics[TB_ART0_9   ][TB_ART_WithTBDisease] += 0.01
    art_characteristics[TB_ART10_14 ][TB_ART_WithTBDisease] += 0.01
    art_characteristics[TB_ART15p   ][TB_ART_WithTBDisease] += 0.01

    art_characteristics[TB_ART0_9   ][TB_ART_PercPTB] +=  1-c_ep
    art_characteristics[TB_ART10_14 ][TB_ART_PercPTB] +=  1-c_ep
    art_characteristics[TB_ART15p   ][TB_ART_PercPTB] +=  1-c_ep
    return art_characteristics