from SpectrumCommon.Modvars.GB.GBDefs import *
from SpectrumCommon.Const.DP import *

from SpectrumCommon.Modvars.DP.DPIndices import *

sourceDictShape = {
    'Name' : 'str',
    'Summary' : 'str',
    'Source' : 'str',
    'Date' : 'str',
}

DP_modvar_shapes = {
    DP_DataSourceTag : gb_str_shape('DP_DataSourceTag'),
    # DP_DPSourcesTag : gb_shape('UPD Sources', type='dict', 
    #     innerShape={
    #         'Provider' : 'str',
    #         'VersionNum' : 'int',
    #         'Sources' : gb_shape('Data type', type='1d array of dict', indices=dp_upd_source_indices, 
    #             innerShape=gb_shape('Source type', type='dict',
    #                 innerShape=sourceDictShape
    #             )
    #         )
    #     }),
    DP_DefaultUPDLETag : gb_sex_year_shape('DP_DefaultUPDLETag'),
    DP_DefaultUPDIMRTag : gb_sex_year_shape('DP_DefaultUPDIMRTag'),
    DP_DefaultUPDCMRTag : gb_sex_year_shape('DP_DefaultUPDCMRTag'),
    DP_DefaultUPDSurvRateTag : gb_shape('DP_DefaultUPDSurvRateTag', type='3D float',
        dim=['surv_rate_ages', gb_sex_dim, gb_year_dim], 
        indices=[dp_surv_rate_indices, gb_sex_indices, gb_year_indices]),
    DP_TFRTag : gb_year_shape('Total fertility rate'),
    DP_ASFRTag : gb_5y_age_year_shape('DP_ASFRTag'),
    DP_UNPopASFRTag : gb_5y_age_year_shape('DP_UNPopASFRTag'),
    DP_ASFRTablesTag : gb_shape('DP_ASFRTablesTag', type='3D float',
        dim=['tables', 'cols', 'rows'], 
        indices=[dp_asfr_table_indices, dp_asfr_col_indices, dp_asfr_row_indices], has_default=False),
    DP_ASFRNumTag : gb_int_shape('DP_ASFRNumTag'),
    DP_ASFRCustomFlagTag : gb_bool_shape('DP_ASFRCustomFlagTag'),
    DP_SexBirthRatioTag : gb_year_shape('DP_SexBirthRatioTag'),
    DP_LETag : gb_sex_year_shape('Life expectancy'),
    DP_MigrRateTag : gb_sex_5y_age_year_shape('DP_MigrRateTag'),
    DP_MigrAgeDistTag : gb_sex_5y_age_year_shape('DP_MigrAgeDistTag'),
    DP_ModelLifeFileNameTag : gb_str_shape('DP_ModelLifeFileNameTag', has_default=False),
    DP_UseFYrSingleAgesTag : gb_bool_shape('DP_UseFYrSingleAgesTag', has_default=False),
    DP_RegionalAdjustPopDataTag : gb_sex_single_age_year_shape('DP_RegionalAdjustPopDataTag', has_default=False),
    DP_RegionalAdjustPopCBStateTag : gb_bool_shape('DP_RegionalAdjustPopCBStateTag', has_default=False),
    DP_RegionalAdjustPopFNameTag : gb_str_shape('DP_RegionalAdjustPopFNameTag',has_default=False),
    DP_RegionalAdjustPopFDateTag : gb_str_shape('DP_RegionalAdjustPopFDateTag', has_default=False),
    # DP_CustomFileYearsTag : gb_shape('DP_CustomFileYearsTag', type='dict',
    #     innerShape = {
    #         'DP_CPopFirstYr' : 'int',
    #         'DP_CPopFinalYr' : 'int'
    #     },
    #      has_default=False),
    DP_CustomPopStopRescalingYearTag : gb_int_shape('DP_CustomPopStopRescalingYearTag', has_default=False),
    DP_LifeTableNumTag : gb_int_shape('DP_LifeTableNumTag'),
    DP_SurvRateTag : gb_shape('DP_SurvRateTag', type='3D float',
        dim=['surv_rate_ages', gb_sex_dim, gb_year_dim], 
        indices=[dp_surv_rate_indices, gb_sex_indices, gb_year_indices]),
    
    DP_BigPopTag : gb_sex_single_age_year_shape('DP_BigPopTag'),
    DP_InitialConditionsTag : gb_shape('DP_InitialConditionsTag', type='dict', 
        innerShape={
            'pop1'      : gb_shape('pop1', '4D float',
                dim=[gb_sex_dim, gb_single_age_dim, 'hiv', 'dur'], 
                indices=[{'male' : 0, 'female' : 1}, 
                         gb_single_age_indices, 
                         {('hiv ' + str(a)) : a for a in GBRange(0, 7)},
                         {('dur ' + str(a)) : a for a in GBRange(0, 7)}
                ]
            ),
            'births'    : gb_shape('births', '2D float',
                dim=[gb_sex_dim, gb_single_age_dim], 
                indices=[{'male' : 0, 'female' : 1}, gb_single_age_indices]),
            'deaths'    : gb_shape('deaths', '2D float',
                dim=[gb_sex_dim, gb_single_age_dim], 
                indices=[{'male' : 0, 'female' : 1}, gb_single_age_indices])
        }),
    
    ##outputs## : gb_shape('##outputs##', type='incomplete', has_default=False),
    DP_BirthsTag : gb_sex_5y_age_year_shape('DP_BirthsTag', has_default=False),
    DP_DeathsByAgeTag : gb_sex_single_age_year_shape('DP_DeathsByAgeTag', has_default=False),
    DP_IMRTag : gb_sex_year_shape('DP_IMRTag'),
    DP_CMRTag : gb_sex_year_shape('DP_CMRTag'),
    DP_AdjLETag : gb_sex_year_shape('DP_AdjLETag', has_default=False),
    
    #DPTempModvars : gb_shape('#DPTempModvars', type='incomplete', has_default=False),
    
    DP_ResultAgeSpecificMortalityTag : gb_sex_single_age_year_shape('DP_ResultAgeSpecificMortalityTag', has_default=False),
    DP_ResultDeathsByAgeLastBirthdayTag : gb_sex_single_age_year_shape('DP_ResultDeathsByAgeLastBirthdayTag', has_default=False),
    DP_ResultCWRTag : gb_year_shape('DP_ResultCWRTag', has_default=False),
    DP_ResultGRRTag : gb_year_shape('DP_ResultGRRTag', has_default=False),
    DP_ResultMACTag : gb_year_shape('DP_ResultMACTag', has_default=False),
    DP_ResultNRRTag : gb_year_shape('DP_ResultNRRTag', has_default=False),
    DP_ResultIMRTag : gb_sex_year_shape('DP_ResultIMRTag', has_default=False),
    DP_ResultU5MRTag : gb_sex_year_shape('DP_ResultU5MRTag', has_default=False),
    DP_Result45q15Tag : gb_year_shape('DP_Result45q15Tag', has_default=False),
    DP_ResultCBRTag : gb_year_shape('DP_ResultCBRTag', has_default=False),
    DP_ResultCDRTag : gb_year_shape('DP_ResultCDRTag', has_default=False),
    DP_ResultDoublingTimeTag : gb_year_shape('DP_ResultDoublingTimeTag', has_default=False),
    DP_ResultGRTag : gb_year_shape('DP_ResultGRTag', has_default=False),
    DP_ResultRNITag : gb_year_shape('DP_ResultRNITag', has_default=False),
    DP_ResultMedianAgeTag : gb_sex_year_shape('DP_ResultMedianAgeTag', has_default=False),
    
    DP_UALowHighRatiosTag : gb_shape('DP_UALowHighRatiosTag', type='dictList',
        innerShape = {
            'isValidUAData' : 'bool',
            'mstID' : 'int',
            'values' : gb_shape('values', '3D Float', 
                dim=[gb_sex_dim, 'high/low ratios', gb_single_age_dim], 
                indices=[gb_sex_dim, {'dead space' : 0, 'low ratio' : DPUA_LowRatio, 'high ratio' : DPUA_HighRatio}, gb_year_indices]
            )
        }, 
        has_default=False),
    DP_AIDSDeathsFYearUATag : gb_int_shape('DP_AIDSDeathsFYearUATag', has_default=False),
    DP_UAHasRanTag : gb_bool_shape('DP_UAHasRanTag', has_default=False),
    # DP_UAInfoTag : gb_shape('DP_UAInfoTag', type='dict',
    #     innerShape = {
    #         'ProcessDate' : 'str',
    #         'ProcessDateDelphi' : 'int',
    #         'ProcessTime' : 'str',
    #         'NumIterations' : 'int',
    #         'AgDataYear' : 'int'
    #     }, 
    #     has_default=False),
}