from copy import deepcopy
from io import BytesIO

from AvenirCommon.Util import GBRange

from SpectrumCommon.Const.AM import *
from SpectrumCommon.Const.DP import *
from SpectrumCommon.Const.PJ import *

from SpectrumCommon.Modvars.GB.GBUtil import getYearIdx
from Tools.AM.EPP.EPPDataRW import ReadSPTFile, convert_epp_modvars_to_list, convert_epp_modvars_to_numpy

def AM_UpdateModvars(projection, created_projection):

    finalYrIdx = getYearIdx(projection[PJN_FinalYearTag], projection[PJN_FirstYearTag])

    # Default should alwauys come from created projection
    projection[AM_FRRbyLocationTag][DP_Default] = created_projection[AM_FRRbyLocationTag][DP_Default]

    # split cotrim data up into separate modvars, out of conglomerate MV
    if AM_ChildARTCalcTag in projection:
        projection[AM_ChildNeedCotrimTag] = projection[AM_ChildARTCalcTag][DP_NumNeedCot].copy()
        projection[AM_ChildOnCotrimTag] = projection[AM_ChildARTCalcTag][DP_NumOnCot].copy()

    # from module-specific to projection-wide sources
    if AM_EdAdSourceTag in projection:
        for src in projection[AM_EdAdSourceTag]:
            projection[PJN_UserSourceTag].append(src)

    if AM_PercLostFollowupTag in projection:
        projection[AM_PercInterruptedTag] = projection[AM_PercLostFollowupTag].copy()

    if AM_PercLostFollowupChildTag in projection:
        projection[AM_PercInterruptedChildTag] = projection[AM_PercLostFollowupChildTag].copy()   
    
    if AM_PreventionNeedsPrEPTagOld in projection:
        for i, pop_type in enumerate(created_projection[AM_PreventionNeedsPrEPTag]['data']):
            prev = pop_type['Prev']
            projection[AM_PreventionNeedsPrEPTagOld]['data'][i]['Prev'] = prev
        
        
        projection[AM_PreventionNeedsPrEPTag] = projection[AM_PreventionNeedsPrEPTagOld].copy() # copy old tag to new tag with updated data
        del projection[AM_PreventionNeedsPrEPTagOld] # remove old tag so this code doesn't run again 

    # if the flag still exists, then check if there is zero data.
    # if the data is zero then copy latest default data into prev need modvars
    if AM_PreventionNeedsFirstTimeTag in projection:
        sum = 0
        for i, pop_type in enumerate(projection[AM_PreventionNeedsPrEPTag]['data']):
            sum += pop_type['PopFirstYear']
        if sum==0:
            for tag in [AM_PreventionNeedsShowVMMCTabTag, AM_PreventionNeedsCondomsTag, 
                        AM_PreventionNeedsVMMCTag, AM_PreventionNeedsKeyPopsTag, AM_PreventionNeedsPrEPTag]:
                projection[tag] = created_projection[tag].copy() # copy old tag to new tag with updated data
            projection[AM_PreventionNeedsFirstTimeTag]

        del projection[AM_PreventionNeedsFirstTimeTag] # remove old tag so this code doesn't run again
        

    if AM_ValidationWaterfallTag in projection:

        if type(projection[AM_ValidationWaterfallTag]) == dict:
            result = [projection[AM_ValidationWaterfallTag] for i in GBRange(0, finalYrIdx)]
            projection[AM_ValidationWaterfallTag] = result

        for yearVals in projection[AM_ValidationWaterfallTag]:
            if 'Disengaged' in yearVals:
                yearVals['Interrupted'] = yearVals['Disengaged']             
                yearVals.pop('Disengaged')
                
            if 'TransferredIn' in yearVals:         
                yearVals.pop('TransferredIn')      
            if 'SitesAdded' in yearVals:         
                yearVals.pop('SitesAdded')  
            if 'TransferredOut' in yearVals:         
                yearVals.pop('TransferredOut')   
            if 'SitesDropped' in yearVals:         
                yearVals.pop('SitesDropped')

                
    if AM_EPPEpidemicTagV1 in projection:  
        if projection['spt'] is not None:
            buffer = BytesIO(projection['spt'])
            buffer.seek(0)
            convert_epp_modvars_to_numpy(projection)
            ReadSPTFile(buffer, projection)
            convert_epp_modvars_to_list(projection)
            pass
        #  = '<AM_EPPEpidemic_V2>'
    
    for t, val in enumerate(projection[AM_HIVTestingTag]):
        created_projection[AM_HIVTestingTag][t].update(val)
    projection[AM_HIVTestingTag] = deepcopy(created_projection[AM_HIVTestingTag])

    def fixChildMortCD4ART(projection, tag):
        if tag in projection:
            if type(projection[tag][0]) == int:
                projection[tag][0] = np.zeros_like(projection[tag][DP_Default]).tolist()
                
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithART0to6Tag)
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithART0to6PercTag)
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithART7to12Tag)
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithART7to12PercTag)
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithARTGT12Tag)
    fixChildMortCD4ART(projection, AM_ChildMortByCD4WithARTGT12PercTag)

    AM_CSAVR_UpdateModvars(projection, created_projection)

    AM_Shiny90_UpdateModvars(projection, created_projection)

    created_projection[AM_InitialConditionsTag].update(projection[AM_InitialConditionsTag])
    projection[AM_InitialConditionsTag] = deepcopy(created_projection[AM_InitialConditionsTag])

    return projection


def AM_Shiny90_UpdateModvars(projection, created_projection):
    def transferDataFromBadTag(toTag, fromTag):
        if fromTag in projection:
            projection[toTag] = projection[fromTag]
           
    transferDataFromBadTag(AM_Shiny90PopTag,    AM_Shiny90PopTagOld)
    transferDataFromBadTag(AM_Shiny90NumTestsTag, AM_Shiny90NumTestsTagOld)
    transferDataFromBadTag(AM_Shiny90NumTested12MTag, AM_Shiny90NumTested12MTagOld)
    transferDataFromBadTag(AM_Shiny90TotalTestsTag, AM_Shiny90TotalTestsTagOld)
    transferDataFromBadTag(AM_Shiny90EverTestedTag, AM_Shiny90EverTestedTagOld)
    transferDataFromBadTag(AM_Shiny90DiagnosedTag, AM_Shiny90DiagnosedTagOld)
    transferDataFromBadTag(AM_Shiny90NumAwareTag, AM_Shiny90NumAwareTagOld)
    transferDataFromBadTag(AM_Shiny90NumDiagnosesTag, AM_Shiny90NumDiagnosesTagOld)
    transferDataFromBadTag(AM_Shiny90NumDiagnosesModTag, AM_Shiny90NumDiagnosesModTagOld)
    transferDataFromBadTag(AM_Shiny90NumLateDiagnosesTag, AM_Shiny90NumLateDiagnosesTagOld)
    transferDataFromBadTag(AM_Shiny90NotDiagHIV1YrTag, AM_Shiny90NotDiagHIV1YrTagOld)
    transferDataFromBadTag(AM_Shiny90IsFittedTag, AM_Shiny90IsFittedTagOld)



def AM_CSAVR_UpdateModvars(projection, created_projection):
    def transferFitData(toTag, fromTag):
        if fromTag in projection:
            projection[toTag][:DP_rLogistic+1] = projection[fromTag].copy()   

    transferFitData(AM_CSAVRAdjustIRRsTag,                  AM_CSAVRAdjustIRRsTagV1)
    # transferFitData(AM_CSAVRFitMethodTag,                   AM_FitIncidenceFitMethodTag)
    transferFitData(AM_CSAVRMetaDataTag,                    AM_CSAVRMetaDataTagV1)
    transferFitData(AM_CSAVRMeanCD4atDiagnosisTag,          AM_MeanCD4atDiagnosisTagV1)
    transferFitData(AM_CSAVRMeanCD4atDiagnosisByAgeSexTag,  AM_MeanCD4atDiagnosisByAgeSexTagV1)
    transferFitData(AM_CSAVRTimeInfToDiagTag,               AM_TimeInfToDiagTagV1)
    transferFitData(AM_CSAVRPropOfDiagnosedTag,             AM_PropOfDiagnosedTagV1)
    transferFitData(AM_CSAVRPropOfDiagnosedNoARTTag,        AM_PropOfDiagnosedNoARTTagV1)
    transferFitData(AM_CSAVRPropOfDiagnosedByAgeSexCD4Tag,  AM_PropOfDiagnosedByAgeSexCD4TagV1)
    transferFitData(AM_CSAVRNumPLHIVTag,                    AM_CSAVRNumPLHIVTagV1)
    transferFitData(AM_CSAVRNumPLHIVByAgeSexCD4Tag,         AM_CSAVRNumPLHIVByAgeSexCD4TagV1)
    transferFitData(AM_CSAVRNumDiagnosedTag,                AM_CSAVRNumDiagnosedTagV1)
    transferFitData(AM_CSAVRNumDiagnosedByAgeSexCD4Tag,     AM_CSAVRNumDiagnosedByAgeSexCD4TagV1)
    transferFitData(AM_CSAVRAIDSDeathsTag,                  AM_CSAVRAIDSDeathsTagV1)
    transferFitData(AM_CSAVRAIDSDeathsByAgeSexTag,          AM_CSAVRAIDSDeathsByAgeSexTagV1)
    transferFitData(AM_CSAVRNumNewInfectionsTag,            AM_CSAVRNumNewInfectionsTagV1)
    transferFitData(AM_CSAVRUncertaintyParamsTag,           AM_FitIncidenceUncertaintyParamsTag)
    transferFitData(AM_CSAVRHIVSexRatioTag,                 AM_CSAVRHIVSexRatioTagV1)
    transferFitData(AM_CSAVRDistOfHIVTag,                   AM_CSAVRDistOfHIVTagV2)
    
    if AM_FitIncidenceParametersTag in projection:
        if type(projection[AM_FitIncidenceParametersTag]) == dict:
            projection[AM_CSAVRParametersTag].update(projection[AM_FitIncidenceParametersTag])

    # for i in projection[AM_CSAVRParametersTag]:
    #     if i in created_projection[AM_CSAVRParametersTag]:
    #         created_projection[AM_CSAVRParametersTag][i].update(projection[AM_CSAVRParametersTag][i])
    # projection[AM_CSAVRParametersTag] = deepcopy(created_projection[AM_CSAVRParametersTag])
    projection[AM_CSAVRParametersTag] = transferCSAVRParams(projection[AM_CSAVRParametersTag], created_projection[AM_CSAVRParametersTag])

    if AM_FitIncidenceIncScaleParametersTag in projection:
        if type(projection[AM_FitIncidenceIncScaleParametersTag]) == dict:
            projection[AM_CSAVRIncScaleParametersTag].update(projection[AM_FitIncidenceIncScaleParametersTag])

    # for i in projection[AM_CSAVRIncScaleParametersTag]:
    #     if i in created_projection[AM_CSAVRIncScaleParametersTag]:
    #         created_projection[AM_CSAVRIncScaleParametersTag][i].update(projection[AM_CSAVRIncScaleParametersTag][i])
    # projection[AM_CSAVRIncScaleParametersTag] = deepcopy(created_projection[AM_CSAVRIncScaleParametersTag])
    projection[AM_CSAVRIncScaleParametersTag] = transferCSAVRParams(projection[AM_CSAVRIncScaleParametersTag], created_projection[AM_CSAVRIncScaleParametersTag])

    if AM_FitIncidencePopulationValueTag in projection:
        projection[AM_CSAVRPopulationValueTag] = projection[AM_FitIncidencePopulationValueTag]    
        
    if AM_FitIncidenceTypeOfFitTag in projection:
        projection[AM_CSAVRTypeOfFitTag] = projection[AM_FitIncidenceTypeOfFitTag]   

def transferCSAVRParams(params, created_params):
    result = deepcopy(created_params)
    for i in params:
        if i in created_params:
            for j in created_params[i]:
                if j in params[i]:
                    if int(created_params[i][j]['mstID']) == int(params[i][j]['mstID']):
                        created_params[i][j]['value'] = params[i][j]['value']        
                    
    return result




AM_ChildARTCalcTag                          = '<AM_ChildARTCalc_V1>'
AM_EdAdSourceTag                            = '<AM_EdAdSource_V1>'
AM_PercLostFollowupTag                      = '<AM_PercLostFollowup_V1>'
AM_PercLostFollowupChildTag                 = '<AM_PercLostFollowupChild_V1>'

#CSAVR
AM_CSAVRAdjustIRRsTagV1                     = '<AM_CSAVRAdjustIRRs_V1>'
# AM_FitIncidenceFitMethodTag                 = '<AM_FitIncidenceFitMethod_V1>'
AM_CSAVRMetaDataTagV1                       = '<AM_CSAVRMetaData_V1>'
AM_MeanCD4atDiagnosisTagV1                  = '<AM_MeanCD4atDiagnosis_V1>'
AM_MeanCD4atDiagnosisByAgeSexTagV1          = '<AM_MeanCD4atDiagnosisByAgeSex_V1>'
AM_TimeInfToDiagTagV1                       = '<AM_TimeInfToDiag_V1>'
AM_PropOfDiagnosedTagV1                     = '<AM_PropOfDiagnosed_V1>'
AM_PropOfDiagnosedNoARTTagV1                = '<AM_PropOfDiagnosedNoART_V1>'
AM_PropOfDiagnosedByAgeSexCD4TagV1          = '<AM_PropOfDiagnosedByAgeSexCD4_V1>'
AM_CSAVRNumPLHIVTagV1                       = '<AM_CSAVRNumPLHIV_V1>'
AM_CSAVRNumPLHIVByAgeSexCD4TagV1            = '<AM_CSAVRNumPLHIVByAgeSexCD4_V1>'
AM_CSAVRNumDiagnosedTagV1                   = '<AM_CSAVRNumDiagnosed_V1>'
AM_CSAVRNumDiagnosedByAgeSexCD4TagV1        = '<AM_CSAVRNumDiagnosedByAgeSexCD4_V1>'
AM_CSAVRAIDSDeathsTagV1                     = '<AM_CSAVRAIDSDeaths_V1>'
AM_CSAVRAIDSDeathsByAgeSexTagV1             = '<AM_CSAVRAIDSDeathsByAgeSex_V1>'
AM_CSAVRNumNewInfectionsTagV1               = '<AM_CSAVRNumNewInfections_V1>'
AM_FitIncidenceParametersTag                = '<AM_FitIncidenceParameters_V1>'
AM_FitIncidenceIncScaleParametersTag        = '<AM_FitIncidenceIncScaleParameters_V1>'
AM_FitIncidenceUncertaintyParamsTag         = '<AM_FitIncidenceUncertaintyParams_V1>'
AM_CSAVRHIVSexRatioTagV1                    = '<AM_CSAVRHIVSexRatio_V1>'
AM_CSAVRDistOfHIVTagV2                      = '<AM_CSAVRDistOfHIV_V2>'
AM_FitIncidencePopulationValueTag           = '<AM_FitIncidencePopulationValue_V1>'
AM_FitIncidenceTypeOfFitTag                 = '<AM_FitIncidenceTypeOfFit_V1>'

# Shiny90
AM_Shiny90PopTagOld                            = '<AM_Shiny90Pop_V1'
AM_Shiny90NumTestsTagOld                       = '<AM_Shiny90NumTests_V1'
AM_Shiny90NumTested12MTagOld                   = '<AM_Shiny90NumTested12M_V1'
AM_Shiny90TotalTestsTagOld                     = '<AM_Shiny90TotalTests_V1'
AM_Shiny90EverTestedTagOld                     = '<AM_Shiny90EverTested_V1'
AM_Shiny90DiagnosedTagOld                      = '<AM_Shiny90Diagnosed_V1'
AM_Shiny90NumAwareTagOld                       = '<AM_Shiny90NumAware_V1'
AM_Shiny90NumDiagnosesTagOld                   = '<AM_Shiny90NumDiagnoses_V1'
AM_Shiny90NumDiagnosesModTagOld                = '<AM_Shiny90NumDiagnosesMod_V1'
AM_Shiny90NumLateDiagnosesTagOld               = '<AM_Shiny90NumLateDiagnoses_V1'
AM_Shiny90NotDiagHIV1YrTagOld                  = '<AM_Shiny90NotDiagHIV1Yr_V1'
AM_Shiny90IsFittedTagOld                       = '<AM_Shiny90IsFitted_V1'

AM_EPPEpidemicTagV1                           = '<AM_EPPEpidemic_V1>'

AM_PreventionNeedsPrEPTagOld            = '<AM_PreventionNeeds_PrEP_V2>'