from SpectrumCommon.Const.DP import *
from SpectrumCommon.Modvars.DP.DPUtil import *
from SpectrumCommon.Const.GB import *
from SpectrumCommon.Modvars.GB.GBDefs import *
from AvenirCommon.Util import GBRange, getDataYear
from SpectrumCommon.Modvars.GB.GBUtil import  getYearIdx, getYear
from SpectrumCommon.Modvars.DP.DPUtil import sexMap

import numpy as np

# def EdAdSource_Init():
#     result = []
#     # (* Data Source Variables *)
#     # result : TstringList.Create
#     # 'Sorted : False
#     return result

def EdDmSource_Init():
    result = []
    # result : TstringList.Create
    # 'Sorted : False
    return result

def DPSources_Init(UPDdb):
    def getSourceDict():
        return {
            'Name' : 'N/A',
            'Summary' : 'N/A',
            'Source' : 'N/A',
            'Date' : 'N/A',
        }

    UPD = {
        'Provider' : '',
        'VersionNum' : 0,
        'Sources' : [getSourceDict() for i in GBRange(DP_UPD_Pop, DP_UPD_Migration)]
    }

    if 'general' in UPDdb:
        UPD['Provider'] = UPDdb['general']['Provider']
        UPD['VersionNum'] = UPDdb['general']['Version']
    
    if 'source' in UPDdb:   
        for src in GBRange(DP_UPD_Pop, DP_UPD_Migration):
            UPD['Sources'][src]['Name'] = UPDdb['source'][src]['Name']
            UPD['Sources'][src]['Summary'] = UPDdb['source'][src]['Summary']
            UPD['Sources'][src]['Source'] = UPDdb['source'][src]['Source']
            UPD['Sources'][src]['Date'] = UPDdb['source'][src]['RefDate']

    return UPD

def DefaultUPDLE_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = getList_sex_year(projInfo['finalYrIdx'])
    if 'lfts' in UPDdb:
        for sex in [GB_Male, GB_Female]:
            for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
                dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
                result[sex][t] = UPDdb['lfts'][dbyear][0][sex]['ex']

    return result

def DefaultUPDIMR_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = getList_sex_year(projInfo['finalYrIdx'])

    if 'lfts' in UPDdb:
        for sex in [GB_Male, GB_Female]:
            for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
                dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
                result[sex][t] = (1 - (UPDdb['lfts'][dbyear][1][sex]['lx'] / 100000)) * 1000

    return result

def DefaultUPDCMR_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = getList_sex_year(projInfo['finalYrIdx'])

    if 'lfts' in UPDdb:
        for sex in [GB_Male, GB_Female]:
            for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
                dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
                result[sex][t] = (1 - (UPDdb['lfts'][dbyear][5][sex]['lx'] / 100000)) * 1000
 
    return result

def CustomFilePYears_Init():
    result = {
        'DP_CPopFirstYr' : -1,
        'DP_CPopFinalYr' : -1,
    }
    return result

def CustomPopStopRescalingYear_Init(projInfo):
    result = projInfo['finalYr']
    return result

def DefaultUPDSurvRate_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = np.zeros((DP_MAX_SURVRATE + 1, GB_Female + 1, projInfo['finalYrIdx'] + 1)).tolist()

    if 'lfts' in UPDdb:
        for age in GBRange(0, DP_MaxSingleAges + 1):
            for sex in [GB_Male, GB_Female]:
                for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
                    dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
                    result[age + 1][sex][t] = UPDdb['lfts'][dbyear][age][sex]['sx']

    return result

def BigPop_Init(projInfo, UPDdb, InitialConditionsDB):
    result = getList_sex_singleAge_year(projInfo['finalYrIdx'])
    firstYear = projInfo['firstYr']
    firstYearIdx = getYearIdx(int(firstYear), projInfo['firstYr'])

    # if the first year is 1970, use the values directly from the UPD files.  
    # this is extra important for regeneration of the InitialConditionsDB so we don't compound any errors that could occur
    # in a closed loop system.  freshly seeding the base year data is good and stuff 
    if firstYear == GB_BaseYear:  
        if 'basepop' in UPDdb:                                          
            for age in GBRange(0, DP_MaxSingleAges):
                for sex in GBRange(GB_Male, GB_Female):
                            result[sex][age][firstYearIdx] = UPDdb['basepop'][age][sex]
    else: 
        Pop1DefaultDBArr = np.array(InitialConditionsDB['pop1'])
        for age in GBRange(0, DP_MaxSingleAges):
            for sex in GBRange(GB_Male, GB_Female):
                result[sex][age][firstYearIdx] = Pop1DefaultDBArr[firstYear - GB_BaseYear, sexMap[sex], age, :, :].sum() #t, S, A, CD4, Dur
    
    result[GB_BothSexes] = np.array(result)[GBRange(GB_Male, GB_Female)].sum(axis=0).tolist()

    return result

def FirstYearPop_Meta_Init(projInfo, UPDdb, InitialConditionsDB, subnat):
    return np.array(BigPop_Init(projInfo, UPDdb, InitialConditionsDB, subnat))[:, :, 0].tolist()
    
def DP_InitialConditions_Init(projInfo, InitialConditionsDB):
    firstYear = projInfo['firstYr']
    result = {
        'pop1'  : InitialConditionsDB['pop1'][firstYear - GB_BaseYear] if (not (firstYear == GB_BaseYear)) else 'why are you here',
        'births'  : InitialConditionsDB['births'][firstYear - GB_BaseYear] if (not (firstYear == GB_BaseYear)) else 'why are you here',
        'deaths'  : InitialConditionsDB['deaths'][firstYear - GB_BaseYear] if (not (firstYear == GB_BaseYear)) else 'why are you here',
    }

    return result

def TFR_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear
 
    result = getList_year(projInfo['finalYrIdx'])

    if 'tfr' in UPDdb:
        for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
            dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
            result[t] = UPDdb['tfr'][dbyear]

    return result

def ASFR_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear
    
    result = getList_5yrAge_year(projInfo['finalYrIdx'])

    if 'pasfrs' in UPDdb:
        for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
            dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
            sum = 0
            for age in GBRange(15, 49):
                result[singleAgeTo5YearAG(age)][t] += UPDdb['pasfrs'][dbyear][age] * 100
                sum += UPDdb['pasfrs'][dbyear][age] * 100
            result[DP_AllAges][t] = sum

    return result

def UNPopASFR_Init(projInfo, UPDdb):   
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear
    
    result = getList_5yrAge_year(projInfo['finalYrIdx'])

    if 'pasfrs' in UPDdb:
        for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
            dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
            sum = 0
            for age in GBRange(15, 49):
                result[singleAgeTo5YearAG(age)][t] += UPDdb['pasfrs'][dbyear][age] * 100
                sum += UPDdb['pasfrs'][dbyear][age] * 100
            result[DP_AllAges][t] = sum

    return result

def ASFRTables_Init(DPGlobalData):

    result = np.zeros((DP_ASFR_Average + 1, DP_ASFR_NumCols + 1, DP_ASFR_NumRows + 1)).tolist()

    for table in GBRange(DP_ASFR_UNAfrica, DP_ASFR_Average):
        tableName = ASFRTables[table]
        for row in GBRange(1, DP_ASFR_NumRows):
            for col in GBRange(1, DP_ASFR_NumCols): 
                result[table][col][row] = DPGlobalData['ASFRTables'][tableName][row - 1][col - 1]

    return result

def SexBirthRatio_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = getList_year(projInfo['finalYrIdx'])

    if 'srb' in UPDdb:
        for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
            dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
            result[t] = UPDdb['srb'][dbyear]

    return result

def LE_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear

    result = getList_sex_year(projInfo['finalYrIdx'])

    if 'lfts' in UPDdb:
        for sex in [GB_Male, GB_Female]:
            for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
                dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
                result[sex][t] = UPDdb['lfts'][dbyear][0][sex]['ex']

    return result

def MigrRate_Init(projInfo, UPDdb):
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear
    
    result = getList_sex_5yrAge_year(projInfo['finalYrIdx'])

    if 'migration' in UPDdb:
        for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
            dbyear = str(getDataYear(dataFirstYear, dataFinalYear, getYear(t, projInfo['firstYr'])))   #for flatlining
            for sex in GBRange(GB_Male, GB_Female):
                sum = 0
                for age in GBRange(0, 80):
                    result[sex][singleAgeTo5YearAG(age)][t] += UPDdb['migration'][dbyear][age][sex]
                    sum += UPDdb['migration'][dbyear][age][sex]
                result[sex][DP_AllAges][t] = sum

    return result

def MigrAgeDist_Init(projInfo, migrRate):
    result = getList_sex_5yrAge_year(projInfo['finalYrIdx'])

    for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
        for sex in GBRange(GB_Male, GB_Female):
            for age in GBRange(DP_A0_4, DP_MAX_AGE):
                if migrRate[sex][DP_AllAges][t] == 0:
                    result[sex][age][t] = 0
                else:           
                    result[sex][age][t] = migrRate[sex][age][t] / migrRate[sex][DP_AllAges][t] * 100

    return result

def COVID19DeathAgeDist_Init(projInfo, DPGlobalData):

    result = getList_sex_5yrAge_year(projInfo['finalYrIdx'])

    for t in GBRange(DP_MIN_YEAR, projInfo['finalYrIdx']):
        for sex in GBRange(GB_Male, GB_Female):
            for data in [DP_Data, DP_Default]:
                result[data][sex][DP_AUnder1][t] = DPGlobalData['Covid19AgeDist'][DP_A0_4] * 0.67
                result[data][sex][DP_A1_4][t] = DPGlobalData['Covid19AgeDist'][DP_A0_4] * 0.33

                for age in GBRange(DP_A5_9, DP_MAX_AGE):
                        result[data][sex][age][t] = DPGlobalData['Covid19AgeDist'][age]

    return result

def UALowHighRatios_Init():
    return []

def AIDSDeathsFYearUA_Init():
    return 0

def UAHasRan_Init():
    return False

def UAInfo_Init():
    return {
        'ProcessDate'   : '',
        'ProcessDateDelphi' : 0.0,      #dies with delphi desktop
        'ProcessTime'   : '', 
        'NumIterations' : DPUA_Default_Num_Iterations,
        'AgDataYear'    : 2022,
    }