from os import environ
import numpy as np

from AvenirCommon.Database import GB_get_db_json

from SpectrumCommon.Const.DP import DP_NaomiEstimateChoice, DP_NaomiFinalAgeCount, DP_BigPopTag, DP_MigrRateTag
from SpectrumCommon.Const.GB import GB_SPECT_MOD_DATA_CONN_ENV
from SpectrumCommon.Const.PJ.PJNTags import PJN_ISO3AlphaTag

def dp_apply_naomi_subnat(projection, subnat_region):

    
    subnat = GB_get_db_json(environ[GB_SPECT_MOD_DATA_CONN_ENV], 'demproj', f'subnationals\\{subnat_region}_V1.JSON')
    
    
    nat = GB_get_db_json(environ[GB_SPECT_MOD_DATA_CONN_ENV], 'demproj', f'subnationals\\{projection[PJN_ISO3AlphaTag]}_V1.JSON')
    ratio_2023 = np.array(subnat['population'][DP_NaomiEstimateChoice])[:, DP_NaomiFinalAgeCount].sum()/np.array(nat['population'][DP_NaomiEstimateChoice])[:, DP_NaomiFinalAgeCount].sum()
    projection['ratio_2023'] = ratio_2023
    
    # survRate = projection[DP_SurvRateTag]
    projection[DP_BigPopTag] = projection[DP_BigPopTag]*ratio_2023
    projection[DP_MigrRateTag] = projection[DP_MigrRateTag]*0.0
    
    # subnat_year_ratio = 
    
    # popFiveYearAgeGroups = np.zeros((18, 3))
    # # popFiveYearAgeGroups[1:18, 1:] = np.array(subnat['population'][DP_NaomiEstimateChoice])[:, :DP_NaomiFinalAgeCount].transpose()
    # popFiveYearAgeGroups[1:18, 1:] = (np.array(nat['population'][DP_NaomiEstimateChoice])[:, :DP_NaomiFinalAgeCount]*nat_ratio).transpose()
    # convertFiveYearAgeGroupPopsToSingleYearAgePops({DP_SurvRateTag:survRate, DP_BigPopTag:bigPop}, popFiveYearAgeGroups)

    return projection