import numpy as np

from SpectrumCommon.Const.TB import *


def TB_HighRiskGroupDemographicDefaults(pop_15_plus, c_ep, prev_pct, year_len):
    tb_demographics = np.zeros((TB_HRG_Len, TG_HRG_DemoGraphics_Len, year_len))
    tb_demographics[TB_HRG_Prisoners        ][TB_HRG_Dem_RelRisk] += 004.01
    tb_demographics[TB_HRG_Miners           ][TB_HRG_Dem_RelRisk] += 001.92
    tb_demographics[TB_HRG_SeekingCare      ][TB_HRG_Dem_RelRisk] += 003.0
    tb_demographics[TB_HRG_StructuralRFs    ][TB_HRG_Dem_RelRisk] += 002.5
    tb_demographics[TB_HRG_HighTBPrev       ][TB_HRG_Dem_RelRisk] += 002.5
    tb_demographics[TB_HRG_UntreatedFibrosis][TB_HRG_Dem_RelRisk] += 002.5

    for h in range(0, TB_HRG_Len):
        tb_demographics[h][TG_HRG_Dem_TotalPop] += pop_15_plus  # this is total 15+ pop and should come from DP. We can also create a database entry for it.
        # Proportion of population 15+ in HRG
        tb_demographics[h][TG_HRG_Dem_Prop15p] += 0.001 if h == TB_HRG_HighTBPrev else 0# 0% default for all groups except "General population in settings with ≥ 0.5% general prevalence" which defaults to 3.5%. Al values can be edited by user.
        # Proportion of�ART patients 15 years and older with TB infection (# )
        tb_demographics[h][TG_HRG_Dem_TBInfect] += 0.40# from editor

        # tb_demographics[h][TB_HRG_Dem_RelRisk] += 0.025
        # Proportion of�ART patients 15 years and older with TB disease (# )
        tb_demographics[h][TB_HRG_Dem_TBDisease] = prev_pct*tb_demographics[h][TB_HRG_Dem_RelRisk]

        # Proportion with pulmonary TB, PTB (# )
        tb_demographics[h][TB_HRG_Dem_PTB] += 1-c_ep# from editor

    tb_demographics[TB_HRG_SeekingCare][TB_HRG_Dem_TBDisease] [tb_demographics[TB_HRG_SeekingCare][TB_HRG_Dem_TBDisease] <0.1] = 0.1
    return tb_demographics