import math
import numpy

from AvenirCommon.Util import GBRange,  GBEqual
from SpectrumCommon.Modvars.GB.GBUtil import getYearIdx
from SpectrumCommon.Const.GB import *
from SpectrumCommon.Const.DP import *
from SpectrumCommon.Const.AM.AMTags import *
from SpectrumCommon.Const.DP.DPTags import *
from SpectrumCommon.Const.PJ.PJNTags import *
from SpectrumCommon.ServerUtils.FilterUtils import get_all_modvars

from SpectrumCommon.Util.AM.AMUtil import GetAM_HIV, getChildNumRecART

sexMap = numpy.zeros(GB_Female + 1).tolist()
sexMap[GB_Male] = 0
sexMap[GB_Female] = 1      

def singleAgeTo5YearAG(age):
    return min(math.floor(age / 5) + 1, DP_MAX_AGE)

def runDPDataChecker(projID, dataYear):
    # dataYear = 2023

    modvars = get_all_modvars(projID)

    firstYear = modvars[PJN_FirstYearTag]
    finalYear = modvars[PJN_FinalYearTag]

    yearIdx = getYearIdx(dataYear, firstYear)
    
    DDC = {
        'PMTCT' : False,
        'ARTPed' : False,
        'ARTMale' : False,
        'ARTFemale' : False,
        'KOSChild' : False,
        'KOSMale' : False,
        'KOSFemale' : False,
        'VSChild' : False,
        'VSMale' : False,
        'VSFemale' : False,
        'ChildARTInterrupt' : False,
        'AdultARTInterrupt' : False,
        'ChildProgression' : False,
        'ChildDistribution' : False,
        'ChildMort' : False,
        'ChildARTMort' : False,
        'AdultProg' : False,
        'AdultDist' : False,
        'AdultMort' : False,
        'AdultARTMort' : False,
        'FRRCD4' : False,
        'FRRAge' : False,
        'CdArt' : False,
        'MTCTRate' : False,
        'AllocDef' : False,
        'Shiny90Valid' : False,
        'UAValid' : False,
        'AdultARTCov' : False,
        'ChildARTCov' : False,
        'PMTCTCov' : False,
        'EPPMaxAdj' : False,
    }

    # year check
    if (firstYear >= dataYear) or (finalYear <= dataYear):
        return DDC

    ARVRegimen = modvars[AM_ARVRegimenTag] 
    DDC['PMTCT'] = bool((ARVRegimen[DP_PrenatalProphylaxis, DP_TripleARTBefPreg, DP_Number, yearIdx] +
        ARVRegimen[DP_PrenatalProphylaxis, DP_TripleARTDurPreg, DP_Number, yearIdx] +
        ARVRegimen[DP_PrenatalProphylaxis, DP_TripleARTDurPreg_Late, DP_Number, yearIdx]) > 100)


    ChildTreatInputs = modvars[AM_ChildTreatInputsTag] 
    ChildARTByAgeGroupPerNum = modvars[AM_ChildARTByAgeGroupPerNumTag]
    if (ChildTreatInputs[DP_PerChildHIVRecART, yearIdx] != DPNotAvail):
        DDC['ARTPed'] = bool(((ChildTreatInputs[DP_PerChildHIVRecART, yearIdx] == DP_Number) and
            (ChildTreatInputs[DP_PerChildHIVRecART, yearIdx] > 0)))
    else:
        DDC['ARTPed'] = bool((ChildTreatInputs[DP_PerChildHIVRecART0_4, yearIdx] +
            ChildTreatInputs[DP_PerChildHIVRecART5_9, yearIdx] +
            ChildTreatInputs[DP_PerChildHIVRecART10_14, yearIdx]) > 0)


    HAARTBySex = modvars[AM_HAARTBySexTag]  
    HAARTBySexPerNum = modvars[AM_HAARTBySexPerNumTag]
    DDC['ARTMale'] = bool((HAARTBySexPerNum[GB_Male, yearIdx] == DP_Number) and (HAARTBySex[GB_Male, yearIdx] > 0))
    DDC['ARTFemale'] = bool((HAARTBySexPerNum[GB_Female, yearIdx] == DP_Number) and (HAARTBySex[GB_Female, yearIdx] > 0))


    KnowledgeOfStatusInput = modvars[AM_KnowledgeOfStatusInputTag]  
    DDC['KOSChild'] = bool(KnowledgeOfStatusInput[DP_Cas_AG_Child, yearIdx] > 0)

    DDC['KOSMale'] = bool(KnowledgeOfStatusInput[DP_Cas_AG_AdultMale, yearIdx] > 0)

    DDC['KOSFemale'] = bool(KnowledgeOfStatusInput[DP_Cas_AG_AdultFemale, yearIdx] > 0)


    ViralSuppressionInput = modvars[AM_ViralSuppressionInputTag] 
    DDC['VSChild'] = bool(((ViralSuppressionInput[DP_Cas_AG_Child, DP_VS_NumTested, yearIdx] >= 0) and 
        (ViralSuppressionInput[DP_Cas_AG_Child, DP_VS_PLHIVSuppressed, yearIdx] >= 0)))

    DDC['VSMale'] = bool(((ViralSuppressionInput[DP_Cas_AG_AdultMale, DP_VS_NumTested, yearIdx] >= 0) and 
        (ViralSuppressionInput[DP_Cas_AG_AdultMale, DP_VS_PLHIVSuppressed, yearIdx] >= 0))) 
    
    DDC['VSFemale'] = bool((ViralSuppressionInput[DP_Cas_AG_AdultFemale, DP_VS_NumTested, yearIdx] >= 0) and 
        (ViralSuppressionInput[DP_Cas_AG_AdultFemale, DP_VS_PLHIVSuppressed, yearIdx] >= 0))
    
    # DDC['ChildARTInterrupt'] = TODO
    # DDC['AdultARTInterrupt'] = TODO
   
    ChildAnnRateProgressLowerCD4 = modvars[AM_ChildAnnRateProgressLowerCD4Tag] 
    DDC['ChildProgression'] = True
    for sd in GBRange(DP_CD4_Per_GT30, DP_CD4_Per_5_10):
        for s in GBRange(GB_Male, GB_Female):
            for a in GBRange(DP_A0t2, DP_A3t4):
                if not GBEqual(ChildAnnRateProgressLowerCD4[DP_Data, s, a, sd], ChildAnnRateProgressLowerCD4[DP_Default, s, a, sd]):
                    DDC['ChildProgression'] = False
    

    ChildDistNewInfectionsCD4 = modvars[AM_ChildDistNewInfectionsCD4Tag] 
    DDC['ChildDistribution'] = True
    for c in GBRange(DP_CD4_Per_GT30, DP_CD4_Per_LT5):
        if not GBEqual(ChildDistNewInfectionsCD4[DP_Data, c], ChildDistNewInfectionsCD4[DP_Default, c]):
            DDC['ChildDistribution'] = False

    ChildMortByCD4NoART = modvars[AM_ChildMortByCD4NoARTTag] 
    DDC['ChildMort'] = True
    for a in GBRange(DP_A0t2, DP_A5t14):
        for b in GBRange(DP_P_Perinatal, DP_P_BF12):
            for c in GBRange(DP_CD4_Per_GT30, DP_CD4_Per_LT5):
                if not GBEqual(ChildMortByCD4NoART[DP_Data, a, b, c], ChildMortByCD4NoART[DP_Default, a, b, c]):
                    DDC['ChildMort'] = False


    DDC['ChildARTMort'] = True
    ChildMortByCD4WithART0to6 = modvars[AM_ChildMortByCD4WithART0to6Tag] 
    ChildMortByCD4WithART7to12 = modvars[AM_ChildMortByCD4WithART7to12Tag] 
    ChildMortByCD4WithARTGT12 = modvars[AM_ChildMortByCD4WithARTGT12Tag] 
    for s in GBRange(GB_Male, GB_Female):
        for a in GBRange(DP_A0, DP_A3t4):
            for c in GBRange(DP_CD4_Per_GT30, DP_CD4_Per_LT5):
                if not GBEqual(ChildMortByCD4WithART0to6[DP_Data, s, a, c], ChildMortByCD4WithART0to6[DP_Default, s, a, c]):
                    DDC['ChildARTMort'] = False
                if not GBEqual(ChildMortByCD4WithART7to12[DP_Data, s, a, c], ChildMortByCD4WithART7to12[DP_Default, s, a, c]):
                    DDC['ChildARTMort'] = False
                if not GBEqual(ChildMortByCD4WithARTGT12[DP_Data, s, a, c], ChildMortByCD4WithARTGT12[DP_Default, s, a, c]):
                    DDC['ChildARTMort'] = False

    ChildMortalityRates = modvars[AM_ChildMortalityRatesTag] 
    for a in GBRange(DP_CD4_0t4, DP_CD4_5t14):
        for b in GBRange(DP_MortRates_LT12Mo, DP_MortRates_GT12Mo):
            for t in GBRange(getYearIdx(firstYear, firstYear), getYearIdx(finalYear, firstYear)):
                if not GBEqual(ChildMortalityRates[DP_Data, a, b, t], ChildMortalityRates[DP_Default, a, b, t]):
                    DDC['ChildARTMort'] = False


    DDC['AdultProg'] = True
    AdultAnnRateProgressLowerCD4 = modvars[AM_AdultAnnRateProgressLowerCD4Tag] 
    for s in GBRange(GB_Male, GB_Female):
        for a in GBRange(DP_CD4_15_24, DP_CD4_45_54):
            for c in GBRange(DP_CD4_GT500, DP_CD4_50_99):
                if not GBEqual(AdultAnnRateProgressLowerCD4[DP_Data, s, a, c], AdultAnnRateProgressLowerCD4[DP_Default, s, a, c]):
                    DDC['AdultProg'] = False


    DDC['AdultDist'] = True
    AdultDistNewInfectionsCD4 = modvars[AM_AdultDistNewInfectionsCD4Tag] 
    for s in GBRange(GB_Male, GB_Female):
        for a in GBRange(DP_CD4_15_24, DP_CD4_45_54):
            for c in GBRange(DP_CD4_GT500, DP_CD4_LT50):
                if not GBEqual(AdultDistNewInfectionsCD4[DP_Data, s, a, c], AdultDistNewInfectionsCD4[DP_Default, s, a, c]):
                    DDC['AdultDist'] = False


    DDC['AdultMort'] = True
    AdultMortByCD4NoART = modvars[AM_AdultMortByCD4NoARTTag] 
    for s in GBRange(GB_Male, GB_Female):
        for a in GBRange(DP_CD4_15_24, DP_CD4_45_54):
            for c in GBRange(DP_CD4_GT500, DP_CD4_LT50):
                if not GBEqual(AdultMortByCD4NoART[DP_Data, s, a, c], AdultMortByCD4NoART[DP_Default, s, a, c]):
                    DDC['AdultMort'] = False


    DDC['AdultARTMort'] = True
    AdultMortByCD4WithART0to6 = modvars[AM_AdultMortByCD4WithART0to6Tag] 
    AdultMortByCD4WithART7to12 = modvars[AM_AdultMortByCD4WithART7to12Tag] 
    AdultMortByCD4WithARTGT12 = modvars[AM_AdultMortByCD4WithARTGT12Tag] 
    for s in GBRange(GB_Male, GB_Female):
        for a in GBRange(DP_CD4_15_24, DP_CD4_45_54):
            for c in GBRange(DP_CD4_GT500, DP_CD4_LT50):
                if not GBEqual(AdultMortByCD4WithART0to6[DP_Data, s, a, c], AdultMortByCD4WithART0to6[DP_Default, s, a, c]):
                    DDC['AdultARTMort'] = False
                if not GBEqual(AdultMortByCD4WithART7to12[DP_Data, s, a, c], AdultMortByCD4WithART7to12[DP_Default, s, a, c]):
                    DDC['AdultARTMort'] = False
                if not GBEqual(AdultMortByCD4WithARTGT12[DP_Data, s, a, c], AdultMortByCD4WithARTGT12[DP_Default, s, a, c]):
                    DDC['AdultARTMort'] = False

    MortalityRates = modvars[AM_MortalityRatesTag] 
    for b in GBRange(DP_MortRates_LT12Mo, DP_MortRates_GT12Mo):
        for t in GBRange(getYearIdx(firstYear, firstYear), getYearIdx(finalYear, firstYear)):
            if not GBEqual(MortalityRates[DP_Data, b, t], MortalityRates[DP_Default, b, t]):
                DDC['AdultARTMort'] = False


    DDC['FRRCD4'] = True
    FertCD4Discount = modvars[AM_FertCD4DiscountTag] 
    for c in GBRange(DP_CD4_GT500, DP_CD4_LT50):
        if not GBEqual(FertCD4Discount[DP_Data, c], FertCD4Discount[DP_Default, c]):
            DDC['FRRCD4'] = False


    DDC['FRRAge'] = True
    HIVTFR = modvars[AM_HIVTFRTag] 
    for a in GBRange(DP_A20_24, DP_A45_49):
        if not GBEqual(HIVTFR[DP_Data, a, yearIdx], HIVTFR[DP_Default, a, yearIdx]):
            DDC['FRRAge'] = False

    TFR15_19DefaultData = modvars[AM_TFR15_19DefaultDataTag] 
    TFRRegionMeta = modvars[AM_TFRRegion_MetaTag]['default'] 
    if not GBEqual(TFR15_19DefaultData[TFRRegionMeta], HIVTFR[DP_Default, DP_A15_19, yearIdx]):
        DDC['FRRAge'] = False


    DDC['CdArt'] = True
    RatioWomenOnART = modvars[AM_RatioWomenOnARTTag] 
    for a in GBRange(DP_A15_19, DP_A45_49):
        if not GBEqual(RatioWomenOnART[DP_Data, a], RatioWomenOnART[DP_Default, a]):
            DDC['CdArt'] = False


    DDC['MTCTRate'] = True
    TransEffAssump = modvars[AM_TransEffAssumpTag] 
    for regimen in GBRange(DP_NoProphExistInfCD4LT200, DP_ARTStartDurPreg_Late):
        for stage in GBRange(DP_Perinatal, DP_BreastfeedingGE350):
            if not GBEqual(TransEffAssump[DP_Data, regimen, stage], TransEffAssump[DP_Default, regimen, stage]):
                DDC['MTCTRate'] = False


    DDC['AllocDef'] = True
    NewARTPatAlloc = modvars[AM_NewARTPatAllocTag] 
    NewARTPatAllocMeta = modvars[AM_NewARTPatAlloc_MetaTag]['default']      
    for i in GBRange(DP_AdvOpt_ART_ExpMort, DP_AdvOpt_ART_PropElig):
        if not GBEqual(NewARTPatAlloc[i], NewARTPatAllocMeta[i]):
            DDC['AllocDef'] = False


    Shiny90AIDSDeathsFYear = modvars[AM_Shiny90AIDSDeathsFYearTag] 
    AIDSDeaths = modvars[AM_AIDSDeathsByAgeTag] 

    # AIDSDeathsSum = 0
    # for a in GBRange(0, 80):
    AIDSDeathsSum = AIDSDeaths[GB_BothSexes, 0:80, getYearIdx(finalYear, firstYear)].sum()
    DDC['Shiny90Valid'] = GBEqual(math.trunc(AIDSDeathsSum), math.trunc(Shiny90AIDSDeathsFYear))


    AIDSDeathsFYearUA = modvars[DP_AIDSDeathsFYearUATag] 
    AIDSDeaths = modvars[AM_AIDSDeathsByAgeTag] 

    # AIDSDeathsSum = 0
    # for a in GBRange(0, 80):
    AIDSDeathsSum = AIDSDeaths[GB_BothSexes, 0:80, getYearIdx(finalYear, firstYear)].sum()
    DDC['UAValid'] = GBEqual(math.trunc(AIDSDeathsSum), math.trunc(int(AIDSDeathsFYearUA)))


    DDC['AdultARTCov'] = True
    HAARTBySex = modvars[AM_HAARTBySexTag]  
    for s in GBRange(GB_Male, GB_Female):
        for t in GBRange(getYearIdx(firstYear, firstYear), getYearIdx(finalYear, firstYear)):
            if HAARTBySex[s, t] > GetAM_HIV(modvars, t, GB_BothSexes, 15, 80):
                DDC['AdultARTCov'] = False


    DDC['ChildARTCov'] = True
    for t in GBRange(getYearIdx(firstYear, firstYear), getYearIdx(finalYear, firstYear)):
        if getChildNumRecART(modvars, t) > GetAM_HIV(modvars, t, GB_BothSexes, 0, 14):
            DDC['ChildARTCov'] = False


    ChildOnPMTCT = modvars[AM_ChildOnPMTCTTag] 
    ChildNeedPMTCT = modvars[AM_ChildNeedPMTCTTag] 
    for t in GBRange(getYearIdx(firstYear, firstYear), getYearIdx(finalYear, firstYear)):
        if ChildOnPMTCT[t] > ChildNeedPMTCT[t]:
            DDC['PMTCTCov'] = False

    # DDC['EPPMaxAdj'] = (modvars[AM_IncidenceOptionsTag] == DP_EPPIncOpt) and (modvars[AM_EPPMaxAdjFactorTag] == 10)
    DDC['EPPMaxAdj'] = modvars[AM_EPPPrevAdjTag] and GBEqual(modvars[AM_EPPMaxAdjFactorTag], 10)

    return DDC
