from SpectrumCommon.Const.FP import *
from SpectrumCommon.Const.GB import *
from SpectrumCommon.Modvars.GB.GBDefs import *
from SpectrumCommon.Modvars.FP.FPIndices import *


fp_age_year_shape = lambda d, **kwargs : gb_shape(d,'2D float',
                                                  dim=['fp_age',  gb_year_dim], 
                                                  indices=[fp_age_indices, gb_year_indices], **kwargs)

fp_age_method_need_year_shape = lambda d, **kwargs : gb_shape(d,'4D float',
                                                              dim=['fp_age', 'fp_method', 'fp_need', gb_year_dim], 
                                                              indices=[fp_age_indices, fp_method_indices, fp_need_indices, gb_year_indices], **kwargs)


fp_age_need_year_shape = lambda d, **kwargs : gb_shape(d,'3D float',
                                                       dim=['fp_age', 'fp_need', gb_year_dim], 
                                                       indices=[fp_age_indices, fp_need_indices, gb_year_indices], **kwargs)

fp_need_year_shape = lambda d, **kwargs : gb_shape(d,'2D float',
                                                   dim=['fp_need', gb_year_dim], 
                                                   indices=[fp_need_indices, gb_year_indices], **kwargs)

fp_age_year_shape = lambda d, **kwargs : gb_shape(d,'2D float',
                                                  dim=['fp_age', gb_year_dim], 
                                                  indices=[fp_age_indices, gb_year_indices], **kwargs)

fp_age_shape = lambda d, **kwargs : gb_shape(d,'1D float',
                                             dim=['fp_age'], 
                                             indices=[fp_age_indices], **kwargs)



fp_method_set_shape = lambda d, **kwargs : gb_shape(d,'1D int',
                                                    options=fp_method_indices, **kwargs)

fp_age_method_year_shape = lambda d, **kwargs : gb_shape(d,'3D float',
                                                         dim=['fp_age', 'fp_method',  gb_year_dim], 
                                                         indices=[fp_age_indices,fp_method_indices,gb_year_indices], **kwargs)

fp_method_year_shape = lambda d, **kwargs : gb_shape(d,'2D float',
                                                     dim=['fp_method', gb_year_dim], 
                                                     indices=[fp_method_indices,gb_year_indices], **kwargs)

fp_method_source_year_shape = lambda d, **kwargs : gb_shape(d,'3D float',
                                                              dim=['fp_method', 'fp_source', gb_year_dim], 
                                                              indices=[fp_method_indices, fp_need_indices, gb_year_indices], **kwargs)

fp_method_age_source_year_shape = lambda d, **kwargs : gb_shape(d,'4D float',
                                                              dim=['fp_method', 'fp_age', 'fp_source', gb_year_dim], 
                                                              indices=[fp_method_indices, fp_age_indices, fp_need_indices, gb_year_indices], **kwargs)

FP_modvar_shapes = {
        # db,
        #General
        FP_ActiveMethodsTag                : gb_shape("Active methods selected in the configuration", 
                                                      '1D bool',
                                                      dim=['fp_method'],
                                                      range='True, False',
                                                      indices=[fp_method_indices], 
                                                      note="Atleast one method needs to be active"),
        FP_AgeGroupOptionTag               : gb_int_shape('Selected age group option', 
                                                           options={'Single age group': FP_Single_Age_Group, 
                                                                    'Five year age groups':FP_Five_Year_Age_Groups},
                                                            has_default=False) ,
        FP_GoalTag                         : gb_int_shape('Selected Famplan goal', 
                                                           options={'Unmet need goal'          : FP_Unmet_need,
                                                                    'Desired TFR goal'         : FP_Desired_TFR,
                                                                    'Prevalence goal(default)' : FP_Prevalence,
                                                                    'TFR goal'                 : FP_TFR,
                                                                    'Modern CPR goal'          : FP_mCPR ,
                                                                    'Additional modern goal'   : FP_AddMdrnMeth,
                                                                    'Pregnancy goal'           : FP_Pregnancy },
                                                            has_default=False),   


        FP_AbortionOptionTag               : gb_int_shape('Abortion option flag', 
                                                           options={'Calculate abortions':FP_calculate_abortions, 
                                                                    'Specify TAR':FP_specify_TAR},
                                                           has_default=False,
                                                           defaultDataTags=['fp_easydata']) ,
        FP_NeedOptionTag                   : gb_bool_shape('Need option',has_default=False),

        FP_SourceMixTag                    : fp_method_source_year_shape("Source mix", 
                                                                         range='0-100', 
                                                                         note="Values must sum to 100 accross the method dimension"),
        FP_MethodMixTag                    : fp_age_method_need_year_shape("Method mix", 
                                                                           range='0-100', 
                                                                           note="Values must sum to 100 accross the method dimension",
                                                                           defaultDataTags=['fp_easydata']),
        FP_NumSourcesTag                   : gb_int_shape('Number of sources'),
        FP_NumNeedsTag                     : gb_int_shape('Number of needs'),
        FP_NumCustomMethodsTag             : gb_int_shape('Number of custom created contraceptive methods'),
        FP_PregnancyRateTag                : gb_float_shape('Pregnancy Rate', range='0-1.0'),
        FP_PercentPregWithMiscarriageTag   : gb_float_shape('Percent pregnant with miscarriage', range='0-1.0'),
        FP_PropUnintendedPregEndInAbortTag : gb_float_shape('Proportion of unintended pregnancies ending in abortion', defaultDataTags=['fp_easydata']),
        FP_PropUnsafeAbortTag              : gb_float_shape('Proportion of unsafe abortions', defaultDataTags=['fp_easydata']),
        FP_MortalityRateSafeAbortTag       : gb_float_shape('Mortality rate of safe abortions'),
        FP_MortalityRateUnsafeAbortTag     : gb_float_shape('Mortality rate of unsafe abortions', defaultDataTags=['fp_easydata']),
        FP_PropPregWantedLaterTag          : gb_float_shape('Proportion of pregnancies wanted later'),
        FP_PropPregUnwantedTag             : gb_float_shape('Proportion of unwanted pregnancies'),
        FP_MiscarriagePercentTag           : gb_int_shape('Miscarriage percent', range='0-100'),
        FP_CalcMethOptionTag               : gb_int_shape('Calculation method option', options={                                                                      
                                                                            'Use total fecundity and proximate determinants' : FP_TotFec,
                                                                            'Use pregnancy rate approach' : FP_PregRate 
                                                                        }), #,
        FP_MethodPermTermTag               : fp_method_set_shape('Permanent contraceptive methods'),
        FP_MethodLongTermTag               : fp_method_set_shape('Long term contraceptive methods'),
        FP_MethodshortTermTag              : fp_method_set_shape('Short term contraceptive methods'),
        FP_EffectivenessTag                : fp_age_method_year_shape('Units per CYP', defaultDataTags=['fp_effectiveness_db']),
        FP_UnitsPerCYPTag                  : fp_method_year_shape('Units per CYP'),
        FP_AverageAgeTag                   : fp_method_year_shape('Average age'),
        FP_DurationOfUseTag                : fp_method_year_shape('Duration of use'),
        FP_SourceMixTag                    : fp_method_source_year_shape("Source mix", 
                                                                         range='0-100', 
                                                                         note="Values must sum to 100 accross the method dimension"),
        FP_CostPerUATag                    : fp_method_source_year_shape("Cost per UA"),
        FP_FeesTag                         : fp_method_source_year_shape("Fees"),
        FP_GrossCostTag                    : fp_method_source_year_shape("Gross cost"),
        FP_RevenueTag                      : fp_method_source_year_shape("Revenue"),
        FP_NetCostTag                      : fp_method_source_year_shape("Net cost"),
        FP_CommoditiesTag                  : fp_method_source_year_shape("Commodities"),
        FP_PercentInUnionTag               : fp_age_year_shape('Percent in union', defaultDataTags=['fp_easydata']),
        FP_PPITag                          : fp_age_year_shape('PPI', defaultDataTags=['fp_easydata']),
        FP_TARTag                          : fp_age_year_shape('Total abortion rate (TAR)'),
        FP_AbortionPercentTag              : fp_age_year_shape('TAbortion percent'),
        FP_SterilityTag                    : fp_age_year_shape('Sterility'),
        FP_MarriageIndexTag                : fp_age_year_shape('Marriage index'), 
        FP_PPIIndexTag                     : fp_age_year_shape('PPI index'), 
        FP_AbortionIndexTag                : fp_age_year_shape('Abortion index'),
        FP_SterilityIndexTag               : fp_age_year_shape('Sterility index'),
        FP_PrevalenceIndexTag              : fp_age_year_shape('Prevalence index'),
        FP_TFRTag                          : fp_age_year_shape('Total fertility rate (TFR)', defaultDataTags=['fp_easydata']),
        FP_WRATag                          : fp_age_year_shape('Women of reproductive age (WRA)', defaultDataTags=['fp_easydata']),
        FP_MWRATag                         : fp_age_year_shape('Married women of reproductive age (MWRA)', defaultDataTags=['fp_easydata']),
        FP_MortalityRateTag                : fp_age_year_shape('Mortality rate'),
        FP_BirthsTag                       : fp_age_year_shape('Births'),
        FP_BirthsAvertedTag                : fp_age_year_shape('Births averted'),
        FP_AbortionsTag                    : fp_age_year_shape('Abortions'),
        FP_AverageEffectivenessTag         : fp_age_need_year_shape('Average effectiveness', defaultDataTags=['fp_effectiveness_db']),
        FP_PrevalenceTag                   : fp_age_need_year_shape('Contraception prevalence rate', defaultDataTags=['fp_easydata']), 
        FP_MPrevalenceTag                  : fp_age_need_year_shape('Modern prevalence'),
        FP_AddMdrnUsersPrevTag             : fp_age_need_year_shape('Additional modern users prevalence'),
        FP_UnmetNeedTag                    : fp_age_need_year_shape('Unmet need', defaultDataTags=['fp_easydata']),
        FP_UnmetNeedReductionTag           : fp_need_year_shape('Unmet need reduction'), 
        FP_WantedTFRTag                    : fp_age_shape('Wanted TFR'),
        FP_TFRdifferenceTag                : gb_year_shape('TFR differences'),
        FP_TestUnmetNeedFromPrevTag        : gb_bool_shape('Test unmet need from prev', has_default=False),
        FP_CalcCostRegressionTag           : gb_bool_shape('Calculate cost regression', has_default=False),
        FP_ASFRInput_TFRGoalTag            : gb_int_shape('ASFR input', 
                                                          options={                                                                      
                                                              'TFR MT TFR Goal' : FP_TFR_MT_TFRGOAL,
                                                              'ASFR TFR Goal' : FP_ASFR_TFRGOAL 
                                                          },
                                                          has_default=False),
        FP_ChildSurvInputsTag              : gb_shape('Child survival inputs','1D float',
                                                  dim=['fp_child_survival'], 
                                                  indices=[fp_child_survival_indices],
                                                  defaultDataTags=['fp_easydata']),

        FP_PACInputsTag                    : gb_shape('Pac inputs','1D float',
                                                  dim=['fp_pac_inputs', gb_year_dim], 
                                                  indices=[fp_pac_inputs_indices, gb_year_indices]),
        FP_LegalAbortNeedTreatTag          : gb_float_shape('Legal abortion needing treatment'),
        FP_PACOutputsTag                   : gb_shape('Pac outputs','1D float',
                                                  dim=['fp_pac_outputs', gb_year_dim], 
                                                  indices=[fp_pac_outputs_indices, gb_year_indices]),
        FP_LAMinsusceptibleTag             : gb_shape('LAM insusceptible','1D float',
                                                  dim=['fp_lam'], 
                                                  indices=[fp_lam_indices]),
        FP_EPSurveyYearTFRTag              : gb_float_shape('EP Survey Year TFR'),
        FP_TFTag                           : fp_age_year_shape('Total fecundity', defaultDataTags=['fp_easydata']),
        FP_PRTag                           : fp_age_shape('Pregnancy rate'),
        FP_DiscontinuationRateTag          : fp_method_year_shape('Discontinuation rate'),
        FP_UsersTag                        : fp_method_age_source_year_shape('Users'),
        FP_AcceptorsTag                    : fp_method_age_source_year_shape('Acceptors'),
        FP_GrowthRateTag                   : fp_method_age_source_year_shape('Growth rates'),
        FP_PregnanciesTag                  : gb_shape('Pac outputs','1D float',
                                                  dim=['fp_age', 'fp_pregnancies', gb_year_dim], 
                                                  indices=[fp_age_indices, fp_pregnancies_indices, gb_year_indices]),
        
        FP_ChildSurvOutputsTag             : gb_shape('Child survival outputs','2D float',
                                                  dim=['fp_child_survival', gb_year_dim], 
                                                  indices=[fp_child_survival_indices, gb_year_indices]),
        FP_PregAvertedTag                  : gb_year_shape('Pregnancies averted'),
        FP_MaternalDeathsAvertedTag        : gb_year_shape('Maternal deaths averted'),
        FP_UnsafeAbortionsAvertedTag       : gb_year_shape('Unsafe abortions averted averted'),
        FP_SourceTag                       : gb_shape('Source of Famplan data', 'dict', innerShape={
                                                'main' : gb_str_shape('Source of Famplan data'),
                                                'goal' : gb_str_shape('Source of Famplan data')
                                                },
                                                defaultDataTags=['fp_easydata']),
        FP_DataSourceTag                   : gb_str_shape('Source of Famplan data', defaultDataTags=['fp_easydata']),
        FP_AvgHouseholdTag                 : gb_int_shape('Average household size'), 
        FP_GoalAnnualChangeTag             : gb_shape('Annual change goal', 'dictList', innerShape={
                'useAnnualChange' : gb_bool_shape('Use annual change flag', has_default=False), 
                'annualChange' : gb_float_shape('Annual change value')} 
        ),
        FP_RelativeRisksTag                : gb_shape('Relative risk', 'dict', innerShape={
            'ageBirthOrder' : gb_shape("age and birth order", 
                                       "2D float", 
                                       dim=['fp_rr_birth_order', 'fp_rr_age_order'],
                                       indices=[fp_rr_birth_order_indices, fp_rr_age_order_indices]), 

            'birthInterval' : gb_shape("birth interval",
                                       "1D float",
                                       dim=['fp_rr_birth_interval'],
                                       indices=[fp_rr_birth_interval_indices]), 
            
            'ageBirthAdv'   : gb_shape("advance age and birth order",
                                       "3D float", 
                                       dim=['fp_rr_advance', 'fp_birth_order', 'fp_age_order'],
                                       indices=[fp_rr_advance_indices, fp_rr_birth_order_indices, fp_rr_age_order_indices]), 
            'birthIntAdv'   : gb_shape("advance birth interval",
                                       "2D float",
                                       dim=['fp_rr_advance', 'fp_birth_interval'],
                                       indices=[fp_rr_advance_indices, fp_rr_birth_interval_indices]), 
            'birthOrderAdv' : gb_shape("advance birth order",
                                       "2D float",
                                       dim=['fp_rr_advance', 'fp_rr_mothers_age'],
                                       indices=[fp_rr_advance_indices, fp_rr_mothers_age_indices]), 
            
        }, defaultDataTags=['fp_easydata']),

        
        FP_RRAgeBirthOrderResultTag : gb_shape("age and birth order result",
                                             "3D float", 
                                             dim=['fp_rr_birth_order', 'fp_rr_age_order', gb_year_dim],
                                             indices=[fp_rr_birth_order_indices, fp_rr_age_order_indices, gb_year_indices]), 
        FP_RRBirthIntervalResultTag : gb_shape("birth interval result",
                                             "2D float",
                                             dim=['fp_rr_birth_interval', gb_year_dim],
                                             indices=[fp_rr_birth_interval_indices, gb_year_indices]), 

    }