
import numpy as np
from os import environ
from copy import deepcopy

from AvenirCommon.Database import GB_get_db_json, GB_file_exists
from SpectrumCommon.Const.GB import GB_BaseYear 
from SpectrumCommon.Const.FP import *
from SpectrumCommon.Modvars.FP.FPModvarDef import *


from SpectrumCommon.Const.DP import *


#@autoreload
def init_modvars(countryID, first_year, final_year,  extra):
    #Variables needed for multiple modvars
    #final_year = 2030
    # if countryID in FP_COUNTRY_PROXY:
    #     countryID = FP_COUNTRY_PROXY[countryID]
    
    #base_year = GB_BaseYear
    #first_year_idx = first_year - base_year
    final_year_idx = final_year - first_year 

    # db_first_yr_idx = first_year-DP_DefaultSurveyYear
    # db_final_yr_idx = final_year-DP_DefaultSurveyYear
    db_first_yr_idx = first_year-1990
    db_final_yr_idx = final_year-1990
    
    exploreLiST = extra['exploreLiST']
    RapidProj = extra['RapidProj']
    # LiSTProj = extra['LiSTProj']
    
    mode = FP_MARUninterpolated if exploreLiST else FP_MAR

    if RapidProj:
        mode = FP_AllUninterpolated if exploreLiST else FP_All

    # if LiSTProj:
    #     mode = FP_MARUninterpolated if exploreLiST else FP_MAR
    
    dir = FP_DB_Dir[mode] + '/' + FP_DB_SubDir[mode]
    version = FP_DB_Version[mode]
    
    #Load databases
    db = GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'famplan', f'{dir}/{countryID}_{version}.JSON') 
    eff_db = GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'famplan', 'FPEffectiveness_V1.JSON') 
    
    #extend database years
    for key in db:
        if type(db[key]) is list:
            for age in range(FP_All_Ages, FP_MAX_AGE+1) :
                if final_year>FP_FinalDBYear:
                    #flatline years   
                    for yr in range(FP_FinalDBYear, final_year):  
                        last_item = db[key][age][-1]
                        db[key][age].append(last_item)  
                elif final_year<FP_FinalDBYear:
                    #truncate years
                    for yr in range(final_year, FP_FinalDBYear):  
                       db[key][age].pop(-1)
                    
                if first_year<FP_FirstDBYear:
                    #flatline years   
                    for yr in range(first_year, FP_FirstDBYear):  
                        first_item = db[key][age][0]
                        db[key][age].insert(0, first_item)
                        pass
                
                
                #cut off historical years, leaving values only between first and final year
                if db_first_yr_idx>0:
                    db[key][age] = db[key][age][db_first_yr_idx:db_final_yr_idx+1]
            db[key] = np.array(db[key])

    

    #replace baseline tfr with tfr from demproj file
    DP_TFR = FP_GetDPTFR(countryID, first_year, final_year)
    if DP_TFR:
        db['TFR'][0] = DP_TFR
    year_zeroes = np.zeros((final_year_idx+1))
    age_year_zeroes = np.zeros((FP_MAX_AGE+1, final_year_idx+1))
    age_zeroes = np.zeros((FP_MAX_AGE+1))
    method_src_year_zeroes = np.zeros((FP_Max_Methods+1, FP_max_sources+1, final_year_idx+1))
    age_need_year_zeroes =  np.zeros((FP_MAX_AGE+1, FP_max_need+1, final_year_idx+1))
    
    method_age_src_year_zeroes = np.zeros((FP_Max_Methods+1, FP_MAX_AGE+1, FP_max_sources+1, final_year_idx+1))

    method_src_year_hundred = np.full((FP_Max_Methods+1, FP_max_sources+1, final_year_idx+1), 100)

    # abortion_percent = np.full((FP_MAX_AGE+1, final_year_idx+1), 32)
    abortion_percent = np.zeros((FP_MAX_AGE+1, final_year_idx+1))
    # abortion_percent[FP_All_Ages] += db['abortionsPerUnwantedPregnancy']
    abortion_percent[FP_All_Ages] += db['PropUnintendedPregEndInAbort']*100

    sterility =  db['childlessWomen45to49Sterility']

    in_union = db['inUnion']
    WRA = Fp_GetDPWRAFlatlined (countryID, first_year, final_year_idx, extra)
    MWRA = FP_ResetMWRA(final_year_idx, in_union, WRA, [FP_All_Ages])

    method_mix = MethodMix_Init(db, final_year_idx)
    effectiveness = Effectiveness_Init(eff_db, final_year_idx)
    avg_effectiveness =  FP_AverageEffectiveness(final_year_idx, method_mix, effectiveness)

    effectiveness_meta = deepcopy(effectiveness)

    #modvar creation list
    return {
        # db,
        #General
        FP_ActiveMethodsTag                : ActiveMethod_Init(db),
        FP_AgeGroupOptionTag               : FP_Single_Age_Group,
        FP_GoalTag                         : FP_Prevalence,
        FP_AbortionOptionTag               : FP_calculate_abortions,
        FP_NeedOptionTag                   : False,
        FP_SourcesTag                      : Sources_Init(),
        FP_MethodNameTag                   : MethodName_Init(),
        FP_NumSourcesTag                   : 1,
        FP_NumNeedsTag                     : FP_max_need,
        FP_NumCustomMethodsTag             : 0,
        FP_PregnancyRateTag                : 0.42,
        FP_PercentPregWithMiscarriageTag   : 0.13,
        FP_PropUnintendedPregEndInAbortTag : db['PropUnintendedPregEndInAbort'],
        FP_PropUnsafeAbortTag              : db['propUnsafeAbort'],
        FP_MortalityRateSafeAbortTag       : 2,
        FP_MortalityRateUnsafeAbortTag     : db['MortalityRateUnsafeAbort'],
        FP_PropPregWantedLaterTag          : 0,
        FP_PropPregUnwantedTag             : 0,
        FP_MiscarriagePercentTag           : FP_MiscarriagePercent,
        FP_CalcMethOptionTag               : FP_TotFec,
        FP_MethodMixTag                    : method_mix,
        FP_MethodPermTermTag               : FP_PermanentMethods,
        FP_MethodLongTermTag               : FP_LongTermtMethods,
        FP_MethodshortTermTag              : FP_ShortTermMethods,
        FP_EffectivenessTag                : effectiveness,
        FP_EffectivenessMetaTag            : effectiveness_meta, #Is this needed? Yes, it is needed for RAPID
        FP_UnitsPerCYPTag                  : UnitsPerCYP_Init(final_year_idx),
        FP_AverageAgeTag                   : AverageAge_Init(countryID, final_year_idx),
        FP_DurationOfUseTag                : DurationOfUse_Init(final_year_idx),
        FP_SourceMixTag                    : method_src_year_hundred,
        FP_CostPerUATag                    : method_src_year_zeroes,
        FP_FeesTag                         : method_src_year_zeroes,
        FP_GrossCostTag                    : method_src_year_zeroes,
        FP_RevenueTag                      : method_src_year_zeroes,
        FP_NetCostTag                      : method_src_year_zeroes,
        FP_CommoditiesTag                  : method_src_year_zeroes,
        FP_PercentInUnionTag               : in_union,
        FP_PPITag                          : db['PPI_MedianDuration'],
        FP_TARTag                          : db['TAR'],
        FP_AbortionPercentTag              : abortion_percent,
        FP_SterilityTag                    : sterility,
        FP_MarriageIndexTag                : age_year_zeroes, 
        FP_PPIIndexTag                     : age_year_zeroes, 
        FP_AbortionIndexTag                : age_year_zeroes,
        FP_SterilityIndexTag               : age_year_zeroes,
        FP_PrevalenceIndexTag              : age_year_zeroes,
        FP_TFRTag                          : db['TFR'],
        FP_TFRMetaTag                      : {'default': db['TFR'].tolist()}, #is this needed? Yes, it is needed for RAPID
        FP_WRATag                          : WRA,
        FP_MWRATag                         : MWRA,
        FP_MortalityRateTag                : age_year_zeroes,
        FP_BirthsTag                       : age_year_zeroes,
        FP_BirthsAvertedTag                : age_year_zeroes,
        FP_AbortionsTag                    : age_year_zeroes,
        FP_AverageEffectivenessTag         : avg_effectiveness,
        FP_PrevalenceTag                   : prevalence_Init(db['CPR'], final_year_idx), 
        FP_MPrevalenceTag                  : age_need_year_zeroes,
        FP_AddMdrnUsersPrevTag             : age_need_year_zeroes,
        FP_UnmetNeedTag                    : unmetNeed_Init(db['unmetNeed'], final_year_idx),
        FP_UnmetNeedReductionTag           : np.zeros((FP_max_need+1, final_year_idx+1)),
        FP_WantedTFRTag                    : age_zeroes,
        FP_TFRdifferenceTag                : year_zeroes,
        FP_TestUnmetNeedFromPrevTag        : True,
        FP_CalcCostRegressionTag           : False,
        FP_ASFRInput_TFRGoalTag            : FP_TFR_MT_TFRGOAL,
        FP_ChildSurvInputsTag              : ChildSurvInputs_Init(db),
        FP_PACInputsTag                    : PACInputs_Init(db['MMR'], final_year_idx),
        FP_LegalAbortNeedTreatTag          : 0,
        FP_PACOutputsTag                   : np.zeros((FP_CostPerMaternDeathAvert+1, final_year_idx+1)),  
        FP_LAMinsusceptibleTag             : LAMinsusceptible_Init(),
        FP_EPSurveyYearTFRTag              : 0,
        FP_TFTag                           : db['TF'],
        FP_PRTag                           : age_zeroes,
        FP_DiscontinuationRateTag          : np.zeros((FP_Max_Methods+1, final_year_idx+1)),  
        FP_UsersTag                        : method_age_src_year_zeroes,
        FP_AcceptorsTag                    : method_age_src_year_zeroes,
        FP_GrowthRateTag                   : method_age_src_year_zeroes,
        FP_PregnanciesTag                  : np.zeros((FP_MAX_AGE+1, FP_mistimed+1, final_year_idx+1)),  
        FP_ChildSurvOutputsTag             : np.zeros((FP_FertAdjU5MR+1, final_year_idx+1)), 
        FP_PregAvertedTag                  : year_zeroes,
        FP_MaternalDeathsAvertedTag        : year_zeroes,
        FP_UnsafeAbortionsAvertedTag       : year_zeroes,
        FP_RelativeRisksTag                : relativeRisks_Init(db, final_year_idx),
        FP_SourceTag                       : FPSource_Init(db),
        FP_DataSourceTag                   : db['mainDataSources'],
        FP_AvgHouseholdTag                 : 0,#db['householdSize'],
        FP_GoalAnnualChangeTag             : [{'useAnnualChange' :False, 'annualChange' :0} for g in GBRange(0, FP_Pregnancy)],
        FP_RRAgeBirthOrderResultTag        : np.zeros((FP_RRBirthOrderMax + 1, FP_RRAgeOrderMax + 1, final_year_idx+1)),
        FP_RRBirthIntervalResultTag        : np.zeros((FP_RRBirthIntervalMax + 1, final_year_idx+1))
    }




# Get WRA from DP and set in FP *)
def Fp_GetDPWRAFlatlined (countryID, firstYear, final_year_idx, extra):
    subnatCode = GBGetKeyValue(extra, 'subNatCode', 0)
    # UPDCountryName = formatCountryFName(countryID, DPUPDVersion, subnatCode)\
    # if GB_file_exists(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'demproj', 'UPD/' + UPDCountryName):
        # UPDdb = GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'demproj', 'UPD/' + UPDCountryName) 
    pop1CountryName = formatCountryFName(countryID, DPInitialConditionsDBVersion)#, subnatCode)
    if GB_file_exists(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'demproj', DPInitialConditionsDBSubDir + pop1CountryName):
        Pop1DefaultDB = np.array(GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'demproj', DPInitialConditionsDBSubDir + pop1CountryName)['pop1'])
        
        first_year_pop = Pop1DefaultDB[firstYear - GB_BaseYear, :, :, :, :]
        female_pop = first_year_pop[1, :, :, :]

        WRA = np.zeros((FP_MAX_AGE+1, final_year_idx+1))
            # +3 is because FP at age 15-20 starts at the 4 id for DP, offset +3*)
        WRA[FP_A15_19] += np.sum(female_pop[15:20])
        WRA[FP_A20_24] += np.sum(female_pop[20:25])
        WRA[FP_A25_29] += np.sum(female_pop[25:30])
        WRA[FP_A30_34] += np.sum(female_pop[30:35])
        WRA[FP_A35_39] += np.sum(female_pop[35:40])
        WRA[FP_A40_44] += np.sum(female_pop[40:45])
        WRA[FP_A45_49] += np.sum(female_pop[45:50])
        
        WRA[0] = np.sum(WRA[1:], axis=0)

        return WRA
    
def FP_GetDPTFR(countryID, first_year, final_year):
    
    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:
        return None
    
    dataFirstYear = 1970
    dataFinalYear = DP_UPDFinalYear
#  
    DP_TFR = []

    for year in GBRange(first_year, min(final_year, dataFinalYear)):
        DP_TFR.append(UPDdb['tfr'][str(year)])

    for year in GBRange(dataFinalYear + 1, final_year):
        DP_TFR.append(UPDdb['tfr'][str(dataFinalYear)])
    
    return DP_TFR



def FP_ResetMWRA(final_year_idx, in_union, WRA, ages):
    
    MWRA = np.zeros((FP_MAX_AGE+1, final_year_idx+1))
    for a in ages:
        MWRA[a] = WRA[a]*in_union[a]/100
    
    if len(ages)>1:
        MWRA[FP_All_Ages] = np.sum(MWRA[1:], axis=0)

    return MWRA
            

# Calculate average effectiveness
def FP_AverageEffectiveness(final_year_idx, method_mix, effectiveness):
    avg_effectiveness = np.zeros((FP_MAX_AGE+1, FP_max_need+1, final_year_idx+1))
    
    mix_effect = np.sum(method_mix[:,:,FP_all_need,:]*effectiveness, axis=1)/np.sum(method_mix[:,:,FP_all_need,:], axis=1)

    avg_effectiveness[:, FP_all_need, :] = mix_effect
    return avg_effectiveness



    # FP := GetFPData(proj);
    # for a := FP_All_Ages to FP_MAX_AGE do
    # begin
    #     mixSum := 0;
    #     mixEffect := 0;
    #     for m := 1 to FP_Max_Methods do
    #     begin
    #     if FP.GetmethodActive(m) then //and (m <> FP_LAM)
    #     begin
    #         mixSum := mixSum + FP.GetMethodMix(a,m,FP_All_Need,t);
    #         mixEffect := mixEffect + FP.GetMethodMix(a,m,FP_All_Need,t)
    #                                 *FP.GetEffectiveness(a,m,t);
    #     end;
    #     end;
    #     if mixSum<>0 then
    #     FP.SetAverageEffectiveness(a,FP_All_Need, t, mixEffect/mixSum);
