import numpy as np

from AvenirCommon.Util import *
from SpectrumCommon.Const.FP import *
from SpectrumCommon.Modvars.FP.FPUtil import * 
from SpectrumCommon.Const.GB  import *

def normalize(L, total):
    sum = 0.0
    norm = [0.0] * len(L)

    for i in range(0, (len(L))):
        sum = sum + L[i]

    if sum == 0.0 and np.isclose(sum, 0.0,  abs_tol = 0.00001):
        multiplier = 0.0
    else:
        multiplier = total / sum

    for i in range(0, (len(L))):
        norm[i] = L[i] * multiplier
    
    return norm

def gbNormalize(L):
    valueGreatThanOne = False
    result = []
    for i in range(0, (len(L)-1)):
        if L[i] > 1.0: #and isclose(sum, 1.0,  abs_tol = 0.00001):
            valueGreatThanOne = True
    
    if valueGreatThanOne:
        result = normalize(L, 100)
    else:
        result = normalize(L, 1)
    return result


def ActiveMethod_Init (db):
    result = [False] * (FP_Max_Methods+1)
    #for m = 0 to FP_Max_Methods do
    #   result[m] = False
   # result[FP_CondomML] = True
   # result[FP_CondomFM] = True

    result[FP_CondomML] = np.any(db['condomMale'])      
    result[FP_fster] =     np.any(db['sterilizationFemale'])       
    result[FP_mster] =  np.any(db['sterilizationMale'])         
    result[FP_vag_barrier] = np.any(db['vaginalBarrierDiaphragm'])   
    result[FP_withdrawal] =   np.any(db['traditionalWithdrawal'])   
    result[FP_PerAbst]= np.any(db['traditionalPeriodicAbstinence'])   
    result[FP_Traditional] =  np.any(db['traditionalOther'])   
    result[FP_inj3month]=   np.any(db['injectableDepo3MoDefault'])   
    result[FP_Jadelle] =  np.any(db['implantJadelle5YrDefault'])      
    result[FP_StandardPill] = np.any(db['pillStandard'])                       
    result[FP_IUDCopper] = np.any(db['IUD_CopperT10Yr'])                
    result[FP_other] =      np.any(db['otherModern'])      

    result = [bool(val) for val in result]         

    return result

def Sources_Init ():
    result = [''] * (FP_max_sources+1)
    result[0] = 'All'
    result[1] = 'Default Source'
    return result

def MethodName_Init():
    result = []
    for m in GBRange(0, FP_Max_Methods):
        result += [GetFP_methodName(m)]
    return result

def MethodMix_Init(db, final_year_idx):
    result = np.zeros((FP_MAX_AGE + 1, FP_Max_Methods + 1, FP_max_need + 1, final_year_idx+1))

    need = FP_all_need

    result[:,FP_CondomML,need] = db['condomMale']              
    result[:,FP_fster,need] = db['sterilizationFemale']             
    result[:,FP_mster,need] = db['sterilizationMale']             
    result[:,FP_vag_barrier,need] = db['vaginalBarrierDiaphragm'] 
    result[:,FP_withdrawal,need] = db['traditionalWithdrawal']    
    result[:,FP_PerAbst,need]= db['traditionalPeriodicAbstinence'] 
    result[:,FP_Traditional,need] = db['traditionalOther']        
    result[:,FP_inj3month,need]= db['injectableDepo3MoDefault']         
    result[:,FP_Jadelle,need] = db['implantJadelle5YrDefault']         
    result[:,FP_StandardPill,need] = db['pillStandard']                       
    result[:,FP_IUDCopper,need] = db['IUD_CopperT10Yr']               
    result[:,FP_other,need] = db['otherModern']                   

    return result


def Effectiveness_Init(eff_db, final_year_idx):
    result = np.zeros((FP_MAX_AGE + 1, FP_Max_Methods + 1, final_year_idx+1))

    for eff in (eff_db):
        #reshape effectiveness values to go by age and year
        eff_years = np.swapaxes(np.array([eff['values']]*(final_year_idx+1)), 0, 1)
        result[:, eff['mstID'], :] += eff_years
                          
    return result

def UnitsPerCYP_Init(final_year_idx):
    result = np.zeros((FP_Max_Methods + 1, final_year_idx+1))
    result[FP_all_methods]  += 105.0
    result[FP_CondomFM]     += 120.0
    result[FP_CondomML]     += 120.0
    result[FP_inj3month]    += 4.0
    result[FP_inj2month]    += 6.0
    result[FP_inj1month]    += 13.0
    result[FP_Uninject]     += 4.0
    result[FP_Implanon]     += 2.5
    result[FP_Jadelle]      += 3.8
    result[FP_Sino_Implant] += 3.2
    result[FP_StandardPill] += 15.0
    result[FP_Progestin]    += 15.0
    result[FP_Peri_coital]  += 15.0
    result[FP_VFT]          += 120.0
    result[FP_LAM]          += 0.25
    result[FP_SDM]          += 1.5

    return result
                
def AverageAge_Init(countryID, final_year_idx):
    result = np.zeros((FP_Max_Methods+1, final_year_idx+1))

    if countryID in ['IND', 'NPL', 'BGD']: #india, nepal, bangladesh
        result[FP_mster] += 26.0 
        result[FP_fster] += 26.0 
    else:
        result[FP_mster] += 35.0 
        result[FP_fster] += 35.0 

    return result    
      
def DurationOfUse_Init(final_year_idx):
    result = np.zeros((FP_Max_Methods+1, final_year_idx+1))

    result[FP_IUDCopper]    += 4.6
    result[FP_IUDLngIus]    += 3.3
    result[FP_Implanon]     += 2.5
    result[FP_Jadelle]      += 3.8
    result[FP_Sino_Implant] += 3.2
    result[FP_vag_barrier]  += 1.0

    return result


def prevalence_Init(cpr, final_year_idx):
    result = np.zeros((FP_MAX_AGE+1, FP_max_need+1, final_year_idx+1))
    result[:, FP_all_need] = cpr
    return result
 
 
def unmetNeed_Init(unmet_need, final_year_idx):
    result = np.zeros((FP_MAX_AGE+1, FP_max_need+1, final_year_idx+1))
    result[:, FP_all_need] = unmet_need
    return result 
    
   

def ChildSurvInputs_Init(db):
    result = np.zeros((FP_RelU5MR+1))
    result[FP_RelRiskyBirth] = -0.55
    result[FP_RelIMR]        = 2.31
    result[FP_RelU5MR]       = 3.54
    result[FP_PerBirthRisk]  = 0#db['anyFertilityRisk']
    return result


def PACInputs_Init(MMR, final_year_idx):
    result = np.zeros((FP_CostCounServCase+1, final_year_idx+1))

    # # PAC Default values 
    result[FP_MortalityUntreatVsTreat] += 3.0
    result[FP_MMRMaternDeath][0] = MMR 
    return result

def LAMinsusceptible_Init():
    result = np.zeros((FP_LAM4_5+1))
    result[FP_LAM0_1] = 0.96
    result[FP_LAM2_3] = 0.906
    result[FP_LAM4_5] = 0.815
    return result

    

def defaultRelativeRisk(db):
    risks = [db['MA_L18FirstBirth'],
             db['MA_L18_BO2_3'],
             db['MA_L18_BO_G3'],
             db['MA18_34BO2_3'],
             db['MA18_34BO_G3'],
             db['MA_G35BO2_3'],
             db['MA_G35BO_G3'],
             db['MA18_34FirstBirth'],
             db['MA_G35FirstBirth']]
             
    #Database values should be already normalized
    #gbNormalize(risks)
    total = sum(risks)
    result =  [r/total*100 for r in risks]
    return result

def defaultBirthInterval(db):
    birth_interval = [db['firstBirthInterval'],
                      db['L18Months'],
                      db['BO18_24Months'],
                      db['BO_24MonthsOrGreater']]

    #Database values should be already normalized
    #gbNormalize(birth_interval)#maybe normalize risk in blob storage
    total = sum(birth_interval)
    result =  [r/total*100 for r in birth_interval]
    return result

def defaultRelRiskCoef(ageBirthAdv, birthIntAdv, birthOrderAdv):
    #set advance constants, ask john of bob for values
      #First Birth
    ageBirthAdv[FP_RRIntrcept][FP_RRFirstBirth][FP_RRLess18] =  7.265480
    ageBirthAdv[FP_RRC1][FP_RRFirstBirth][FP_RRLess18]       = -0.091041
    ageBirthAdv[FP_RRC2][FP_RRFirstBirth][FP_RRLess18]       =  0.001016

    ageBirthAdv[FP_RRIntrcept][FP_RRFirstBirth][FP_RR18To34] =  8.798430
    ageBirthAdv[FP_RRC1][FP_RRFirstBirth][FP_RR18To34]       =  0.303922
    ageBirthAdv[FP_RRC2][FP_RRFirstBirth][FP_RR18To34]       = -0.000318

    ageBirthAdv[FP_RRIntrcept][FP_RRFirstBirth][FP_RRMore34] =  0.125492
    ageBirthAdv[FP_RRC1][FP_RRFirstBirth][FP_RRMore34]       = -0.007625
    ageBirthAdv[FP_RRC2][FP_RRFirstBirth][FP_RRMore34]       =  0.000222

      #Second to third birth
    ageBirthAdv[FP_RRIntrcept][FP_RRSecToThirdBirth][FP_RRLess18] =  2.161788
    ageBirthAdv[FP_RRC1][FP_RRSecToThirdBirth][FP_RRLess18]       = -0.047509
    ageBirthAdv[FP_RRC2][FP_RRSecToThirdBirth][FP_RRLess18]       =  0.000438

    ageBirthAdv[FP_RRIntrcept][FP_RRSecToThirdBirth][FP_RR18To34] = 26.088305
    ageBirthAdv[FP_RRC1][FP_RRSecToThirdBirth][FP_RR18To34]       =  0.274474
    ageBirthAdv[FP_RRC2][FP_RRSecToThirdBirth][FP_RR18To34]       = -0.001101

    ageBirthAdv[FP_RRIntrcept][FP_RRSecToThirdBirth][FP_RRMore34] =  0.696853
    ageBirthAdv[FP_RRC1][FP_RRSecToThirdBirth][FP_RRMore34]       = -0.022429
    ageBirthAdv[FP_RRC2][FP_RRSecToThirdBirth][FP_RRMore34]       =  0.000796

      #Four or more births
    ageBirthAdv[FP_RRIntrcept][FP_RRMoreFourBirth][FP_RRLess18] =  0.047730
    ageBirthAdv[FP_RRC1][FP_RRMoreFourBirth][FP_RRLess18]       = -0.001463
    ageBirthAdv[FP_RRC2][FP_RRMoreFourBirth][FP_RRLess18]       =  0.000012

    ageBirthAdv[FP_RRIntrcept][FP_RRMoreFourBirth][FP_RR18To34] = 39.282154
    ageBirthAdv[FP_RRC1][FP_RRMoreFourBirth][FP_RR18To34]       = -0.344732
    ageBirthAdv[FP_RRC2][FP_RRMoreFourBirth][FP_RR18To34]       = -0.000165

    ageBirthAdv[FP_RRIntrcept][FP_RRMoreFourBirth][FP_RRMore34] = 15.168927
    ageBirthAdv[FP_RRC1][FP_RRMoreFourBirth][FP_RRMore34]       = -0.061253
    ageBirthAdv[FP_RRC2][FP_RRMoreFourBirth][FP_RRMore34]       = -0.000866


    birthOrderAdv[FP_RRIntrcept][FP_RRBOGreater3] = -18.90215906
    birthOrderAdv[FP_RRC1][FP_RRBOGreater3]       =  19.41538934
    birthOrderAdv[FP_RRC2][FP_RRBOGreater3]       = -1.32560387

    birthOrderAdv[FP_RRIntrcept][FP_RRMAGreat35] = -1.940569547
    birthOrderAdv[FP_RRC1][FP_RRMAGreat35]       =  5.700909531
    birthOrderAdv[FP_RRC2][FP_RRMAGreat35]       = -0.471392363


    birthIntAdv[FP_RRIntrcept][FP_RRFirstBirth] = 16.554243
    birthIntAdv[FP_RRC1][FP_RRFirstBirth]       =  0.202914
    birthIntAdv[FP_RRC2][FP_RRFirstBirth]       =  0.000887

    birthIntAdv[FP_RRIntrcept][FP_RRLess18Months] = 1.854915
    birthIntAdv[FP_RRC1][FP_RRLess18Months]       = -0.011982
    birthIntAdv[FP_RRC2][FP_RRLess18Months]       = -0.000019

    birthIntAdv[FP_RRIntrcept][FP_RR18to23Months] = 5.648637
    birthIntAdv[FP_RRC1][FP_RR18to23Months]       = 0.029261
    birthIntAdv[FP_RRC2][FP_RR18to23Months]       = -0.000311

    birthIntAdv[FP_RRIntrcept][FP_RRMore24Months] = 75.942205
    birthIntAdv[FP_RRC1][FP_RRMore24Months]       = -0.220192
    birthIntAdv[FP_RRC2][FP_RRMore24Months]       = -0.000557


def relativeRisks_Init(db, final_year_idx):
    #result = [[[0]*(final_year_idx+1)]*FP_RRAgeOrderMax]*FP_RRBirthOrderMax

    ageBirthOrder = np.zeros((FP_RRBirthOrderMax + 1, FP_RRAgeOrderMax + 1))
    birthInterval = np.zeros(FP_RRBirthIntervalMax + 1)
    ageBirthAdv   = np.zeros((FP_RRC2 + 1, FP_RRBirthOrderMax + 1, FP_RRAgeOrderMax + 1))
    birthIntAdv   = np.zeros((FP_RRC2 + 1, FP_RRExtraBirthIntervals + 1))
    birthOrderAdv = np.zeros(( FP_RRC2 + 1, FP_RRMA18To34 + 1))
   

    default_risks = defaultRelativeRisk(db)
    ageBirthOrder[FP_RRFirstBirth][FP_RRLess18]       = default_risks[0] 
    ageBirthOrder[FP_RRSecToThirdBirth][FP_RRLess18]  = default_risks[1] 
    ageBirthOrder[FP_RRMoreFourBirth][FP_RRLess18]    = default_risks[2] 
    ageBirthOrder[FP_RRSecToThirdBirth][FP_RR18To34]  = default_risks[3] 
    ageBirthOrder[FP_RRMoreFourBirth][FP_RR18To34]    = default_risks[4] 
    ageBirthOrder[FP_RRSecToThirdBirth][FP_RRMore34]  = default_risks[5] 
    ageBirthOrder[FP_RRMoreFourBirth][FP_RRMore34]    = default_risks[6]
    ageBirthOrder[FP_RRFirstBirth][FP_RR18To34]       = default_risks[7]
    ageBirthOrder[FP_RRFirstBirth][FP_RRMore34]       = default_risks[8]

    default_birth_interval = defaultBirthInterval(db)
    birthInterval[FP_RRFirstBirth]    = default_birth_interval[0] 
    birthInterval[FP_RRLess18Months]  = default_birth_interval[1]
    birthInterval[FP_RR18to23Months]  = default_birth_interval[2] 
    birthInterval[FP_RRMore24Months]  = default_birth_interval[3]

    defaultRelRiskCoef(ageBirthAdv, birthIntAdv, birthOrderAdv)

    # # flatline result values with baseline values
    # for t in GBRange(1, final_year_idx):
    #     for i in GBRange(1, FP_RRBirthOrderMax):
    #         for ma in GBRange(1, FP_RRAgeOrderMax):
    #             ageBirthOrderResult[i][ma][t] = ageBirthOrder[i][ma]

    #     for i in GBRange(1, FP_RRBirthIntervalMax):
    #         birthIntervalResult[i][t] = birthInterval[i]
    return {
        'ageBirthOrder' : ageBirthOrder.tolist(),      
        'birthInterval' : birthInterval.tolist(), 
        
        'ageBirthAdv'   : ageBirthAdv.tolist(), 
        'birthIntAdv'   : birthIntAdv.tolist(), 
        'birthOrderAdv' : birthOrderAdv.tolist(),
        
        }


def FPSource_Init(db):
    group = db['Group']
    country_name = db['name']
    main_data_sources = db['mainDataSources']
    result = main_data_sources

    # if (group == 'A') or (group == 'AB'):
    #     result = f'Survey data are from: {main_data_sources} {country_name}'
    # elif (group == 'B'):
    #     result = f'Contraceptive prevalence is from: {main_data_sources} {country_name} \n Other proximate determinants data are based upon averages calculated for countries with similar levels of fertility.'
    # elif (group == 'C'):
    #     result = 'All contraceptive information and other proximate determinants data are based upon averages calculated for countries with similar levels of fertility.'
    # else:
    #     result = main_data_sources

    # result += '\n\nAbortions (i.e., terminations per unwanted pregnancy) is based on subregional estimations by Bearak et al 2020.'

    return {
        'main':result,
        'goal':result
    } 

