
from AvenirCommon.Util import GBRange
from SpectrumCommon.Const.TB import *
from SpectrumCommon.Const.PJ.PJNTags import PJN_ISO3AlphaTag, PJN_FinalYearTag
from SpectrumCommon.Modvars.TB.TBNotifDefs import apply_notif_split
import SpectrumCommon.External.TB.TBFortWrapper as fort



def delete_modvar(projection, tag):
    if tag in projection:
        del projection[tag]

def TB_PreUpdateModvars(projection, created_projection):
    if TB_DataSourceTag not in projection:
        # Projections created prior to the addition of this modvar should all have the same TB data source (WHO 2023)
        projection[TB_DataSourceTag] = 'WHO 2023'
        
    if TB_FinalHistoricalYearTag not in projection:
        projection[TB_FinalHistoricalYearTag] = TB_FinalHist_Year

    if TB_FortInputVersionTag not in projection:
        if (projection[TB_UseOldFailsafeTag]):
            projection[TB_FortInputVersionTag] = 'V1'
            projection[TB_FortOutputVersionTag] = 'V2'
        elif  projection[TB_UseOldIPTag]:
            projection[TB_FortInputVersionTag] = 'V2'
            projection[TB_FortOutputVersionTag] = 'V3'
        else: 
            projection[TB_FortInputVersionTag] = 'V6'
            projection[TB_FortOutputVersionTag] = 'V6'

        ISO3 = projection[PJN_ISO3AlphaTag]
        finalYear = projection[PJN_FinalYearTag]
        base_year = projection[TB_BaseYearTag]# deafult to TB_FirstHist_Year
        fort_base_index = max(2015, base_year)-TB_FirstHist_Year
        fort_slice = slice(fort_base_index, min(TB_Fort_Final_Year, finalYear) -TB_FirstHist_Year+1)
        tb_slice  = slice(max(2015-base_year, 0), min(TB_Fort_Final_Year, finalYear)-base_year+1)

        projection[TB_FortInputTag] = fort.get_tb_fort_inputs(ISO3, finalYear, tb_slice, fort_slice, 
                                                              None, None, None, 
                                                              fort_vsn=projection[TB_FortInputVersionTag])
        projection[TB_FortOutputTag] = fort.get_tb_fort_outputs(ISO3, projection[TB_FortOutputVersionTag])
        if projection[TB_FortInputVersionTag]  in ['V1', 'V2']:
            projection[TB_FortInputTag]["sEp"] = []
            projection[TB_FortInputTag]["tXf"] = []
            projection[TB_FortInputTag]["pHat"] = []
            for yr in GBRange(projection[TB_FortInputTag]["year"][0], projection[TB_FortInputTag]["year"][-1]):
                projection[TB_FortInputTag]["sEp"].append(None)
                projection[TB_FortInputTag]["pHat"].append(None)
                projection[TB_FortInputTag]["tXf"].append(0.035)

            for yr in GBRange(projection[TB_FortInputTag]["year"][-1]+1, TB_Fort_Final_Year):
                projection[TB_FortInputTag]["year"].append(yr)    
                for key in ["iHat", "sEi", "nHat", "sEn", "mHat", "sEm", "pHat", "sEp"]:
                    projection[TB_FortInputTag][key].append(None)
                projection[TB_FortInputTag]["tXf"].append(0.035)
            projection[TB_FortInputTag]["hRd"] = [1]*len(projection[TB_FortInputTag]["year"]) 
            projection[TB_FortInputTag]["hRi"] = [1]*len(projection[TB_FortInputTag]["year"]) 
            projection[TB_FortInputTag]["oRt"] = [1]*len(projection[TB_FortInputTag]["year"])
        pass

    return projection

def TB_UpdateModvars(projection, created_projection):

    # if TB_OldTag in projection:
    #     projection[TB_NewOldTag] = new_modvars[TB_NewOldTag].copy()
    #     del projection[TB_OldTag]
    if TB_NotificationByCaseTypeTag_V1 in projection:
        notif = projection[TB_TargNotificationTag]
        notif_pct   = created_projection[TB_PercNotificationTag]
        notif_split = created_projection[TB_NotificationByCaseTypeTag]
        notif_split = apply_notif_split(created_projection[TB_BaseYearTag], notif_split, notif, notif_pct)
        projection[TB_NotificationByCaseTypeTag] = notif_split
    # if TB_PercNotificationTag_V1 in projection:
    if TB_BaseYear_V1 in projection:
        projection[TB_UseOldFailsafeTag] = projection[TB_BaseYear_V1] == 2021
        projection[TB_UseOldIPTag] = projection[TB_BaseYear_V1] == 2022
    if TB_HistoricalTag_V1 in projection:
        projection[TB_HistNotifStartYearTag] = projection[TB_HistoricalTag_V1] ['notifications']['startYear']
        for key in TB_NOTIF_DB_KEYS:
            projection[TB_HistNotifTag][key] = projection[TB_HistoricalTag_V1] ['notifications'][key]
        for key in TB_NOTIF_DIST_KEYS:
            projection[TB_HistNotifTag][key] = projection[TB_HistoricalTag_V1] ['Noti_Distribution'][key]
    if TB_HHPrvntMethodsTag_V1 in projection:
            projection[TB_HHPrvntMethodsTag][TB_HH5_14] = projection[TB_HHPrvntMethodsTag_V1][TB_HH5_14] 
            projection[TB_HHPrvntMethodsTag][TB_HH15p] = projection[TB_HHPrvntMethodsTag_V1][TB_HH15p] 
    
    # delete tags
    for tag in TB_TagsForDeletion:
        if tag in projection:
            del projection[tag]

    return projection

TB_ARTDiagCaseValuesTag_V1          = '<TB_ARTDiagCaseValues_V1>'                      
TB_ARTDiagCaseIDsTag_V1             = '<TB_ARTDiagCaseIDs_V1>'                    
TB_ARTDiagCaseMethodsTag_V1         = '<TB_ARTDiagCaseMethodIDs_V1>'   
TB_ARTDiagCaseSensitivityTag_V1     = '<TB_ARTDiagCaseSensitivity_V1>'
TB_ARTDiagCaseSpecificityTag_V1     = '<TB_ARTDiagCaseSpecificity_V1>'  
TB_NotificationByCaseTypeTag_V1     = '<TB_NotificationByCaseType_V1>'
TB_PercNotificationTag_V1           = '<TB_PercentNotification_V1>'      
TB_HHPrvntMethodsTag_V1           = '<TB_HHPrvntMethods_V1>'  
TB_HHPrvntMethodIDsTag_V1         = '<TB_HHPrvntMethodIDs_V1>'   

TB_BaseYear_V1         = '<TB_BaseYear_V1>' 
TB_HistoricalTag_V1    = '<TB_HistoricalData_V1>'

TB_TagsForDeletion = [
    TB_ARTDiagCaseValuesTag_V1, 
    TB_ARTDiagCaseIDsTag_V1,    
    TB_ARTDiagCaseMethodsTag_V1,
    TB_ARTDiagCaseSensitivityTag_V1,
    TB_ARTDiagCaseSpecificityTag_V1,
    TB_NotificationByCaseTypeTag_V1,
    TB_PercNotificationTag_V1,
    TB_BaseYear_V1,
    TB_HistoricalTag_V1,
    TB_HHPrvntMethodsTag_V1,
    TB_HHPrvntMethodIDsTag_V1

]