from os import environ
import numpy as np

from AvenirCommon.Database import GB_get_db_json, GB_file_exists

from SpectrumCommon.Const.PJ.PJNTags import PJN_ISO3AlphaTag, PJN_FirstYearTag, PJN_FinalYearTag
from SpectrumCommon.Const.DP import DP_NaomiEstimateChoice, DP_NaomiFinalAgeCount
from SpectrumCommon.Const.TB import *
from SpectrumCommon.Modvars.TB.TBNotifDefs import *
from SpectrumCommon.Modvars.TB.TBModvars import TB_flatlineDB, TB_flatlinBurden, TB_getPopDefaults
from SpectrumCommon.Const.GB.GBVersions import (TB_WHO_DB_VERSION,
                                               TB_FORT_OUTPUT_VERSION)

def tb_apply_naomi_subnat(projection, subnat_region):
    countryID = projection[PJN_ISO3AlphaTag]
    
    connection = environ['AVENIR_SW_DEFAULT_DATA_CONNECTION']
    db =  GB_get_db_json(connection, 
                         'tuberculosis', 
                         f'countries\\{countryID}_{TB_WHO_DB_VERSION}.JSON')

    base_year = projection[PJN_FirstYearTag]
    final_year = projection[PJN_FinalYearTag]
    year_len = projection[PJN_FinalYearTag]-projection[PJN_FirstYearTag] + 1

    fort_slice = slice(max(2024, base_year)-TB_FirstHist_Year, min(TB_Fort_Final_Year, final_year) -TB_FirstHist_Year+1)
    tb_slice  = slice(max(2024-base_year, 0), min(TB_Fort_Final_Year, final_year)-base_year+1)

    
    first_idx = projection[PJN_FirstYearTag]-db['projections']['startYear']
    final_idx = projection[PJN_FinalYearTag]-db['projections']['startYear']

        # TB_HisIncidenceTag  : hist_inci,
        # TB_HisMortalityTag  : ,
    # flatline defaults 
    TB_flatlineDB(db, 'notification', final_idx)
    TB_flatlineDB(db, 'prevalencePct', final_idx)
    TB_flatlineDB(db, 'prevalence', final_idx)
    TB_flatlineDB(db, 'pop', final_idx)
    TB_flatlinBurden(db, 'e_inc_num', final_idx)
    TB_flatlinBurden(db, 'e_mort_num', final_idx)
    
    #subnat notification
    pop, pop_15_plus = TB_getPopDefaults(environ["AVENIR_SW_DEFAULT_DATA_CONNECTION"], countryID, subnat_region, projection[PJN_FirstYearTag], projection[PJN_FinalYearTag])
    DBName = f"subnationals\\{countryID}_subnat_V1.JSON"
    if GB_file_exists(environ["AVENIR_SW_DEFAULT_DATA_CONNECTION"], "tuberculosis", DBName):
        subnat = GB_get_db_json(environ["AVENIR_SW_DEFAULT_DATA_CONNECTION"], "tuberculosis", DBName)
        subnat = subnat[subnat_region]
        subnat_notif_ratio = subnat['notifProp']
        db['projections']['notification'] = [v*subnat['notifProp'] for v in db['projections']['notification'] ]
        db['projections']['prevalence'] = [v*subnat['notifProp'] for v in db['projections']['notification'] ]
        db['projections']['pop'] = [v*subnat['notifProp'] for v in db['projections']['notification'] ]
        db['projections']['e_inc_num'] = [v*subnat['notifProp'] for v in db['projections']['notification'] ]
        db['projections']['e_mort_num'] = [v*subnat['notifProp'] for v in db['projections']['notification'] ]
        for v in db['Noti_Distribution']:
            if '_perc' not in v:
                db['Noti_Distribution'][v] = db['Noti_Distribution'][v]*subnat['notifProp'] 
        demproj_subnat = GB_get_db_json(environ["AVENIR_SW_DEFAULT_DATA_CONNECTION"], 'demproj', f'subnationals\\{subnat_region}_V1.JSON')
        
        pop_ratio = np.array(demproj_subnat['population'][DP_NaomiEstimateChoice])[:,0:DP_NaomiFinalAgeCount].sum()/(pop[:, :, 0:, :, :].sum())
        pop_15_plus = np.array(demproj_subnat['population'][DP_NaomiEstimateChoice])[:,3:DP_NaomiFinalAgeCount].sum()
        pop = pop*pop_ratio
        pass
        projection[TB_SubnationalRatioTag]  = subnat_notif_ratio
        
        projection[TB_HisIncidenceTag] = [v*subnat_notif_ratio for v in db['burden']['e_inc_num'][first_idx:final_idx+1]]
        projection[TB_HisMortalityTag] = [v*subnat_notif_ratio for v in db['burden']['e_mort_num'][first_idx:final_idx+1]]
        
        
        notif_base_pct, notif_scale, c_ep, bac_pos_pct = create_notification_defaults_2022(countryID, db['projections']['notification'], db['Noti_Distribution'], 2024-db['projections']['startYear'])
        
        notif, notif_split, notif_pct, fort_outputs = create_notification_defaults(countryID, db['projections']['notification'], db['Noti_Distribution'], 
                                                                                   projection[PJN_FirstYearTag], 
                                                                                   first_idx, final_idx, year_len, notif_base_pct, tb_slice, fort_slice, 
                                                                                   TB_FORT_OUTPUT_VERSION, notif_scale, subnat_notif_ratio)
        
        projection[TB_HisNotificationTag] = notif
        projection[TB_TargNotificationTag] = notif
        projection[TB_ResNotificationTag] = np.array([notif, notif, notif*0])
        
        

        year_len = projection[PJN_FinalYearTag]-projection[PJN_FirstYearTag] + 1
        fort_slice = slice(max(2024, projection[PJN_FirstYearTag])-TB_FirstHist_Year, min(TB_Fort_Final_Year, projection[PJN_FinalYearTag]) -TB_FirstHist_Year+1)
        tb_slice  = slice(max(2024-projection[PJN_FirstYearTag], 0), min(TB_Fort_Final_Year, projection[PJN_FinalYearTag])-projection[PJN_FirstYearTag]+1)

        incidence = np.zeros(year_len)
        incidence[tb_slice] = [v*subnat_notif_ratio for v in fort_outputs['iMid'][fort_slice]]
        # flatline no longer needed, taking all of forts values
        # incidence[TB_Fort_Final_Year-base_year:] = incidence[TB_Fort_Final_Year-base_year] # flatline incidence from 2030 onwards
        
        mortality = np.zeros(year_len)
        mortality[tb_slice] = [v*subnat_notif_ratio for v in fort_outputs['mMid'][fort_slice]]

        # prevalence = calc_prevalence_from_notif_inci(notif, incidence)
        prevalence = np.zeros(year_len)
        prevalence[tb_slice] = [v*subnat_notif_ratio for v in fort_outputs['pMid'][fort_slice]]
        projection[TB_ResIncidenceTag ] =  np.array([incidence, incidence, incidence*0])
        projection[TB_ResMortalityTag ] =  np.array([mortality, mortality, mortality*0])
        projection[TB_ResPrevalenceTag] =  np.array([prevalence, prevalence, prevalence*0])
    return projection