import numpy as np

from SpectrumCommon.Const.TB import *


def TB_RRProfileDefaults(db, notif_pct, year_len):
    rr_profile = np.zeros((TB_NewRetreatLen, TB_ResistAgeLen, TB_ResistProfileLen, year_len))
    db_RR_Profile = np.nan_to_num(np.array(db['RR_Profile']), copy=True, nan=0.0, posinf=None, neginf=None)
    db['RR_Profile']  = db_RR_Profile.tolist()
    for t in range(0, year_len):
        rr_profile[TB_NewCases, TB_Child, :, t] = db_RR_Profile[TB_NewCases].copy() #Child
        rr_profile[TB_NewCases, TB_Adult, :, t] = db_RR_Profile[TB_NewCases].copy() #Adult

        rr_profile[TB_RetreatCases, TB_Child, :, t] = db_RR_Profile[TB_RetreatCases].copy() #Child
        rr_profile[TB_RetreatCases, TB_Adult, :, t] = db_RR_Profile[TB_RetreatCases].copy() #Adult

    rp_newU15 = rr_profile[TB_NewCases][TB_Child][TB_Rif_Res]
    rp_new15p = rr_profile[TB_NewCases][TB_Adult] [TB_Rif_Res]

    rp_retU15 = rr_profile[TB_RetreatCases][TB_Child][TB_Rif_Res]
    rp_ret15p = rr_profile[TB_RetreatCases][TB_Adult] [TB_Rif_Res]

    pRet      = notif_pct[TB_PrevTreatIncludeRelapseCases]
    pNew      = 1-pRet

    rp_Rif = np.zeros((TB_ResistanceLen, TB_ResistAgeLen, year_len))
    rp_Rif[TB_MDR][TB_ResistAge0_14] = pNew*(rp_newU15) + (pRet)*(rp_retU15)    
    rp_Rif[TB_MDR][TB_ResistAge15p]  = pNew*(rp_new15p) + (pRet)*(rp_ret15p)  

    rp_Rif[TB_SEN][TB_ResistAge0_14]=(1-rp_Rif[TB_MDR][TB_ResistAge0_14] )
    rp_Rif[TB_SEN][TB_ResistAge15p]=(1-rp_Rif[TB_MDR][TB_ResistAge15p] )

    return rr_profile, rp_Rif

def TB_ResistanceCoverageDefaults(year_len):
    coverages  = np.full((TB_CoverageLen, year_len), 0.50)

    coverages[0]  = 1 # db['coverage']['DST']
    coverages[1]  = 1 # db['coverage']['DST']
    coverages[2]  = 1 # db['coverage']['DST']
    coverages[3]  = 1 # db['coverage']['DST']
    return coverages