
from os import environ

from AvenirCommon.Database import GB_get_db_json, GB_file_exists

from AvenirCommon.Util import formatCountryFName, GBGetKeyValue
from SpectrumCommon.Const.DP import *

from SpectrumCommon.Modvars.DP.DPModvarDef import *


def init_modvars(countryID, first_year, final_year, extra):
    subnatCode = GBGetKeyValue(extra, 'subNatCode', 0)
    
    connection = environ[GB_SPECT_MOD_DATA_CONN_ENV]
    UPDCountryName = formatCountryFName(countryID, DPUPDVersion)#, subnatCode)
    if GB_file_exists(connection, 'demproj', 'UPD/' + UPDCountryName):
        UPDdb = GB_get_db_json(connection, 'demproj', 'UPD/' + UPDCountryName) 
    else:
        UPDdb = {}

    initialConditionsCountryName = formatCountryFName(countryID, DPInitialConditionsDBVersion)#, subnatCode)
    if GB_file_exists(connection, 'demproj', DPInitialConditionsDBSubDir + initialConditionsCountryName):
        InitialConditionsDB = GB_get_db_json(connection, 'demproj', DPInitialConditionsDBSubDir + initialConditionsCountryName) 
    else:
        InitialConditionsDB = {}
        
        
    DPGlobalData = GB_get_db_json(connection, 'demproj',  'DP_Global_' + DPDatabaseVersion + '.JSON')    
   
    projInfo = {
        'firstYr'           : first_year,
        'finalYr'           : final_year,
        'firstYrIdx'        : 0,
        'finalYrIdx'        : final_year - first_year,
        'subnatCode'        : GBGetKeyValue(extra, 'subNatCode', 0),
        'region'            : GBGetKeyValue(extra, 'region', DP_Asia),
    }
    survRate = DefaultUPDSurvRate_Init(projInfo, UPDdb)
    bigPop = BigPop_Init(projInfo, UPDdb, InitialConditionsDB)


    migrRate = MigrRate_Init(projInfo, UPDdb)    
    return {
        DP_DataSourceTag :  '' if UPDdb == {} else UPDdb['general']['Provider'] + ' ' + UPDdb['general']['Version'],
        DP_DPSourcesTag :  DPSources_Init(UPDdb),
        DP_DefaultUPDLETag :  DefaultUPDLE_Init(projInfo, UPDdb),
        DP_DefaultUPDIMRTag :  DefaultUPDIMR_Init(projInfo, UPDdb),
        DP_DefaultUPDCMRTag :  DefaultUPDCMR_Init(projInfo, UPDdb),
        DP_DefaultUPDSurvRateTag :  DefaultUPDSurvRate_Init(projInfo, UPDdb),
        DP_TFRTag :  TFR_Init(projInfo, UPDdb),
        DP_ASFRTag :  ASFR_Init(projInfo, UPDdb),
        DP_UNPopASFRTag :  UNPopASFR_Init(projInfo, UPDdb),
        DP_ASFRTablesTag :  ASFRTables_Init(DPGlobalData),
        DP_ASFRNumTag :  DP_ASFR_UNPop if 'pasfrs' in UPDdb else DP_ASFR_Average,
        DP_ASFRCustomFlagTag :  False,
        DP_SexBirthRatioTag :  SexBirthRatio_Init(projInfo, UPDdb),
        DP_LETag :  LE_Init(projInfo, UPDdb),
        DP_MigrRateTag :  migrRate,
        DP_MigrAgeDistTag :  MigrAgeDist_Init(projInfo, migrRate),
        DP_ModelLifeFileNameTag :  'DP_stInputted',
        DP_UseFYrSingleAgesTag :  GB_FP not in extra['modules'],
        DP_RegionalAdjustPopDataTag :  getList_sex_singleAge_year(projInfo['finalYrIdx']),
        DP_RegionalAdjustPopCBStateTag :  False,
        DP_RegionalAdjustPopFNameTag :  '',
        DP_RegionalAdjustPopFDateTag :  '',
        DP_CustomFileYearsTag :  CustomFilePYears_Init(),
        DP_CustomPopStopRescalingYearTag :  CustomPopStopRescalingYear_Init(projInfo),
        DP_LifeTableNumTag :  DP_UNPopLT, 
        DP_SurvRateTag : survRate,
        
        DP_BigPopTag :  bigPop,
        DP_FirstYearPop_Meta_Tag : np.array(bigPop)[:, :, 0].tolist(),
        DP_InitialConditionsTag : DP_InitialConditions_Init(projInfo, InitialConditionsDB),

        ## outputs ##
        DP_BirthsTag : getList_sex_5yrAge_year(projInfo['finalYrIdx']),
        DP_DeathsByAgeTag : getList_sex_singleAge_year(projInfo['finalYrIdx']),
        DP_IMRTag : getList_sex_year(projInfo['finalYrIdx']),
        DP_CMRTag : getList_sex_year(projInfo['finalYrIdx']),
        DP_AdjLETag : getList_sex_year(projInfo['finalYrIdx']), 

        # DP Temp Modvars

        DP_ResultAgeSpecificMortalityTag :  getList_sex_singleAge_year(projInfo['finalYrIdx']),
        DP_ResultDeathsByAgeLastBirthdayTag :  getList_sex_singleAge_year(projInfo['finalYrIdx']),
        DP_ResultCWRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultGRRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultMACTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultNRRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultIMRTag :   getList_sex_year(projInfo['finalYrIdx']),
        DP_ResultU5MRTag :  getList_sex_year(projInfo['finalYrIdx']),
        DP_Result45q15Tag : getList_year(projInfo['finalYrIdx']),
        DP_ResultCBRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultCDRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultDoublingTimeTag : getList_year(projInfo['finalYrIdx']),
        DP_ResultGRTag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultRNITag :  getList_year(projInfo['finalYrIdx']),
        DP_ResultMedianAgeTag :  getList_sex_year(projInfo['finalYrIdx']),

        DP_UALowHighRatiosTag :  UALowHighRatios_Init(),
        DP_AIDSDeathsFYearUATag :  AIDSDeathsFYearUA_Init(),
        DP_UAHasRanTag :  UAHasRan_Init(),
        DP_UAInfoTag :  UAInfo_Init(),

    }
    
    
