# Contains methods for setting up more complicated IHT ModVars; others can just be set up in ICModVars.py

import numpy as np

import SpectrumCommon.Const.GB.GBConst as gbc

import SpectrumCommon.Const.IC.ICConst as icc
import SpectrumCommon.Const.IC.ICTargetPopulationIDs as ictpi
import SpectrumCommon.Const.IC.ICConstLink as iccl
import SpectrumCommon.Const.IC.ICModVarFields as icmvf

#######################################################################################################
#                                                                                                     #
#                                       Target populations                                            #
#                                                                                                     #
#######################################################################################################


# These targ pop dicts are attached to interventions
def IC_create_interv_targ_pop_rec(num_years: int):

    return {
        icmvf.IC_TARG_POP_MST_ID: ictpi.IC_TP_TOTAL_POPULATION_MST_ID,
        icmvf.IC_TARG_POP_SEX: gbc.GB_Male,
        icmvf.IC_TARG_POP_START_AGE: 0,
        icmvf.IC_TARG_POP_END_AGE: 80,
        icmvf.IC_TARG_POP_DIRECT_ENTRY_VALUES: np.full((num_years), 0.0).tolist(),
    }


# These targ pop dicts define each target population and are not attached / assigned directly to interventions.
def IC_create_targ_pop_def_rec(name: str, string_constant: str, mst_ID: int):

    return {
        icmvf.IC_TARG_POP_DEF_NAME: name,
        icmvf.IC_TARG_POP_DEF_STR_CONST: string_constant,
        icmvf.IC_TARG_POP_DEF_MST_ID: mst_ID,
        icmvf.IC_TARG_POP_DEF_APPLIC_PROG_AREA_IDS: [],
        icmvf.IC_TARG_POP_DEF_REQ_MOD_MST_IDS: [],
    }


#######################################################################################################
#                                                                                                     #
#                                       Interventions                                                 #
#                                                                                                     #
#######################################################################################################


def IC_create_interv_rec(
    mst_ID: str,
    name: str,
    string_const: str,
    active: bool,
    impact: bool,
    prog_area_ID: str,
    subgroup_ID: str,
    num_deliv_channels: int,
    num_years: int,
):

    return {
        icmvf.IC_INTERV_ID: mst_ID,
        icmvf.IC_INTERV_NAME: name,
        icmvf.IC_INTERV_STR_CONST: string_const,
        icmvf.IC_INTERV_ACTIVE: active,
        icmvf.IC_INTERV_CUSTOM: False,
        icmvf.IC_INTERV_IMPACT: impact,
        icmvf.IC_INTERV_PROG_AREA_ID: prog_area_ID,
        icmvf.IC_INTERV_SUBGROUP_ID: subgroup_ID,
        icmvf.IC_INTERV_SPECIAL_CASES: [],
        icmvf.IC_INTERV_STRATIFYING: False,
        icmvf.IC_INTERV_COST_FOR_PHC: False,
        icmvf.IC_INTERV_PIN_EDITOR_VISITED: False,
        icmvf.IC_INTERV_COVERAGE_CHANGED: False,
        icmvf.IC_INTERV_PIN_STR_CONSTANT: "",
        icmvf.IC_INTERV_COVERAGE: np.full((num_years), 0.0).tolist(),
        icmvf.IC_INTERV_TREAT_INPUT_NOT_REQ: [],
        icmvf.IC_INTERV_MAX_CUSTOM_TI_GROUP_MST_ID: [],
        icmvf.IC_INTERV_PIN: np.full((num_years), 0.0).tolist(),
        icmvf.IC_INTERV_PERC_DELIV_THROUGH_CHAN: np.full((num_deliv_channels, icc.IC_NUM_YEAR_ENDPOINTS), 0.0).tolist(),
        icmvf.IC_INTERV_VISIT_COST_PER_CASE: [],
        icmvf.IC_INTERV_TARG_POP_REC: IC_create_interv_targ_pop_rec(num_years),
    }


# This creates an intervention record with only enough info in it to identify and display it in a list.
def IC_create_interv_rec_ID_info_only(mst_ID: str, name: str, string_const: str, prog_area_ID: str, subgroup_ID=str):

    return {
        icmvf.IC_INTERV_ID: mst_ID,
        icmvf.IC_INTERV_NAME: name,
        icmvf.IC_INTERV_STR_CONST: string_const,
        icmvf.IC_INTERV_PROG_AREA_ID: prog_area_ID,
        icmvf.IC_INTERV_SUBGROUP_ID: subgroup_ID,
    }


#######################################################################################################
#                                                                                                     #
#                                  Alternative coverage records                                       #
#                                                                                                     #
#######################################################################################################


def IC_create_alt_cov_dict(interv_ID: str, num_years: int):

    return {
        icmvf.IC_ALT_COV_INTERV_ID: interv_ID,
        icmvf.IC_ALT_COV_VALUE: np.full((num_years), 0.0).tolist(),
    }


#######################################################################################################
#                                                                                                     #
#                                  Alternative PIN age group set 1 records                            #
#                                                                                                     #
#######################################################################################################


def IC_create_alt_PIN_age_group_set_1_rec(interv_ID: str, num_years: int):

    return {
        icmvf.IC_ALT_PIN_AGE_GROUP_SET_1_INTERV_ID: interv_ID,
        icmvf.IC_ALT_PIN_AGE_GROUP_SET_1_VALUE: np.full((icc.IC_NUM_AGE_GROUPS_SET_1, num_years), 0.0).tolist(),
    }


#######################################################################################################
#                                                                                                     #
#                                  Alternative PIN age group set 2 records                            #
#                                                                                                     #
#######################################################################################################


def IC_create_alt_PIN_age_group_set_2_rec(interv_ID: str, num_years: int):

    return {
        icmvf.IC_ALT_PIN_AGE_GROUP_SET_2_INTERV_ID: interv_ID,
        icmvf.IC_ALT_PIN_AGE_GROUP_SET_2_VALUE: np.full((icc.IC_NUM_AGE_GROUPS_SET_2, num_years), 0.0).tolist(),
    }


#######################################################################################################
#                                                                                                     #
#                                       Drugs and supplies                                            #
#                                                                                                     #
#######################################################################################################


def IC_create_drug_supply_rec(
    ID: str,
    name: str,
    active: bool,
    cost_type_LG_mstID: int,
    drug_supply_type_mst_ID: int,
    requires_cold_storage: bool,
    num_years: int,
    costing_mode: int,
    plan_type: int,
):
    drug_supply_rec = {
        icmvf.IC_DRUG_SUPPLY_ID: ID,
        icmvf.IC_DRUG_SUPPLY_NAME: name,
        icmvf.IC_DRUG_SUPPLY_STR_CONST: "",
        icmvf.IC_DRUG_SUPPLY_COST_TYPE_LG_MST_ID: cost_type_LG_mstID,  # Used by the Medicines, commodities, and supplies Summary result.
        icmvf.IC_DRUG_SUPPLY_ACTIVE: active,
        icmvf.IC_DRUG_SUPPLY_CUSTOM: False,
        icmvf.IC_DRUG_SUPPLY_TYPE_MST_ID: drug_supply_type_mst_ID,
        icmvf.IC_DRUG_SUPPLY_APPLIC_PROG_AREA_LIST: [],
        icmvf.IC_DRUG_SUPPLY_UNIT_COST: np.full((num_years), 0.0).tolist(),
        # Used in Logistics
        icmvf.IC_DRUG_SUPPLY_PERC_WASTAGE: 0,
        icmvf.IC_DRUG_SUPPLY_PERC_SAFETY_STOCK_NEEDED: 0,
        icmvf.IC_DRUG_SUPPLY_SAFETY_STOCK_BASELINE: 0,
        icmvf.IC_DRUG_SUPPLY_PERC_SAFETY_STOCK_REC: np.full((num_years), 0.0).tolist(),
        icmvf.IC_DRUG_SUPPLY_QUANTITY: np.full((num_years), 0.0).tolist(),
        icmvf.IC_DRUG_SUPPLY_UNIT_WEIGHT: 0.0,
        icmvf.IC_DRUG_SUPPLY_UNIT_VOLUME: 0.0,
        icmvf.IC_DRUG_SUPPLY_REQUIRES_COLD_STORAGE: requires_cold_storage,
    }

    if ((costing_mode == gbc.GB_IHT_COSTING_MODE) and (plan_type == iccl.IC_IH_SPECIFIC_PLAN)) or (
        costing_mode == gbc.GB_TB_COSTING_MODE
    ):
        drug_supply_rec[icmvf.IC_DRUG_SUPPLY_LOG_SUPP_CHAIN_COSTS_PERC] = 0
        drug_supply_rec[icmvf.IC_DRUG_SUPPLY_WASTAGE_PERC] = 0

    return drug_supply_rec


#######################################################################################################
#                                                                                                     #
#                                       Treatment input group (TIG) records                           #
#                                                                                                     #
#######################################################################################################


def IC_create_treat_input_group_rec(interv_ID: str, deliv_chan_ID: str, treat_input_group_ID: str, name=str):

    return {
        icmvf.IC_TIG_INTERV_ID: interv_ID,
        icmvf.IC_TIG_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_TIG_ID: treat_input_group_ID,
        icmvf.IC_TIG_NAME: name,
        icmvf.IC_TIG_STR_CONST: "",
    }


#######################################################################################################
#                                                                                                     #
#                                       Treatment input (TI) base records                             #
#                                                                                                     #
#######################################################################################################


def IC_create_treat_input_rec(
    interv_ID: str,
    deliv_chan_ID: str,
    treat_input_type_mst_ID: int,
    treat_input_group_ID: str,
    treat_input_ID: str,
    num_years: int,
):

    TI_rec = {
        icmvf.IC_TI_INTERV_ID: interv_ID,
        icmvf.IC_TI_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_TI_GROUP_ID: treat_input_group_ID,
        icmvf.IC_TI_ID: treat_input_ID,
        icmvf.IC_TI_PERC_APPLIED: np.full((num_years), 0.0).tolist(),
        icmvf.IC_TI_NOTE: "",
        icmvf.IC_TI_NOTE_STR_CONST: "",
    }

    if treat_input_type_mst_ID == icc.IC_TI_DRUGS_SUPPLIES_MST_ID:
        TI_rec[icmvf.IC_TI_DRUG_SUPPLY_TIMES_PER_DAY] = 0
        TI_rec[icmvf.IC_TI_DRUG_SUPPLY_DAYS_PER_CASE] = 0
        TI_rec[icmvf.IC_TI_DRUG_SUPPLY_NUMBER] = 0

    elif treat_input_type_mst_ID == icc.IC_TI_OUTPAT_VISITS_INPAT_DAYS_MST_ID:
        TI_rec[icmvf.IC_TI_OUTPAT_INPAT_UNITS_PER_CASE] = 0.0

    elif treat_input_type_mst_ID == icc.IC_TI_MED_PERSONNEL_TIME_MST_ID:
        TI_rec[icmvf.IC_TI_MED_PERSONNEL_CUSTOM_STAFF_TYPE] = False
        TI_rec[icmvf.IC_TI_MED_PERSONNEL_SKILLED_STAFF_COST] = False
        TI_rec[icmvf.IC_TI_MED_PERSONNEL_NUM_OF_DAYS] = 0
        TI_rec[icmvf.IC_TI_MED_PERSONNEL_MINUTES] = 0.0

    return TI_rec


#######################################################################################################
#                                                                                                     #
#                                       Cost per case records                                         #
#                                                                                                     #
#######################################################################################################


def IC_create_cost_per_case_rec(interv_ID: str, deliv_chan_ID: str, value=list):

    return {
        icmvf.IC_F_INTERV_ID: interv_ID,
        icmvf.IC_F_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_F_VALUE: value,
    }


#######################################################################################################
#                                                                                                     #
#                             Minutes per average case (default staff) records                        #
#                                                                                                     #
#######################################################################################################


def IC_create_min_per_avg_case_static_staff_rec(
    interv_ID: str, deliv_chan_ID: str, static_staff_mst_ID: int, value=list
):

    return {
        icmvf.IC_F_INTERV_ID: interv_ID,
        icmvf.IC_F_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_F_STATIC_STAFF_MST_ID: static_staff_mst_ID,
        icmvf.IC_F_VALUE: value,
    }


#######################################################################################################
#                                                                                                     #
#                             Minutes per average case (custom staff) records                         #
#                                                                                                     #
#######################################################################################################


def IC_create_min_per_avg_case_custom_staff_rec(
    interv_ID: str, deliv_chan_ID: str, custom_staff_mst_ID: int, value=list
):

    return {
        icmvf.IC_F_INTERV_ID: interv_ID,
        icmvf.IC_F_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_F_DYNAMIC_STAFF_MST_ID: custom_staff_mst_ID,
        icmvf.IC_F_VALUE: value,
    }


#######################################################################################################
#                                                                                                     #
#                             Inpatient days per average case records                                 #
#                                                                                                     #
#######################################################################################################


def IC_create_inpatient_days_outpatient_visits_per_avg_case_rec(
    interv_ID: str, deliv_chan_ID: str, treat_input_ID: int, value: list
):

    return {
        icmvf.IC_F_INTERV_ID: interv_ID,
        icmvf.IC_F_DELIV_CHAN_ID: deliv_chan_ID,
        icmvf.IC_F_TREAT_INPUT_ID: treat_input_ID,
        icmvf.IC_F_VALUE: value,
    }


#######################################################################################################
#                                                                                                     #
#                                            Results                                                  #
#                                                                                                     #
#######################################################################################################


def IC_init_interv_costed_deliv_chan_year_arr(num_intervs: int, num_costed_deliv_chans: int, num_years: int, value=0.0):
    return np.full((num_intervs, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_deliv_chan_all_year_arr(num_intervs: int, num_deliv_chans_all: int, num_years: int, value=0.0):
    return np.full((num_intervs, num_deliv_chans_all, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_drug_supply_deliv_chan_all_year_arr(
    num_drugs_supplies: int, num_deliv_chans_all: int, num_years: int, value=0.0
):
    return np.full((num_drugs_supplies, num_deliv_chans_all, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_staff_costed_deliv_chan_year_arr(
    num_intervs: int, num_staff_types: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full((num_intervs, num_staff_types, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_static_staff_costed_deliv_chan_year_arr(
    num_intervs: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full(
        (num_intervs, iccl.IC_HW_NUM_STATIC_STAFF_TYPES, num_costed_deliv_chans, num_years),
        value,
        gbc.GB_CALC_FLOAT_TYPE,
    )


def IC_init_interv_year_arr(num_intervs: int, num_years: int, value=0.0):
    return np.full((num_intervs, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_arr(num_intervs: int, value=0.0):
    return np.full((num_intervs), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_empty_arr(value=0):
    return np.empty(value)


def IC_init_year_arr(num_years: int, value=0.0):
    return np.full((num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_costed_deliv_chan_year_arr(num_costed_deliv_chans: int, num_years: int, value=0.0):
    return np.full((num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_static_staff_costed_deliv_chan_year_arr(num_costed_deliv_chans: int, num_years: int, value=0.0):
    return np.full(
        (iccl.IC_HW_NUM_STATIC_STAFF_TYPES, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE
    )


def IC_init_static_staff_interv_year_arr(num_intervs: int, num_years: int, value=0.0):
    return np.full((iccl.IC_HW_NUM_STATIC_STAFF_TYPES, num_intervs, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_static_staff_prog_area_year_arr(num_prog_areas: int, num_years: int, value=0.0):
    return np.full((iccl.IC_HW_NUM_STATIC_STAFF_TYPES, num_prog_areas, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_static_staff_prog_area_deliv_chan_year_arr(
    num_prog_areas: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full(
        (iccl.IC_HW_NUM_STATIC_STAFF_TYPES, num_prog_areas, num_costed_deliv_chans, num_years),
        value,
        gbc.GB_CALC_FLOAT_TYPE,
    )


def IC_init_dynamic_staff_costed_deliv_chan_year_arr(
    num_dynamic_staff_types: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full((num_dynamic_staff_types, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_dynamic_staff_interv_year_arr(num_dynamic_staff_types, num_intervs: int, num_years: int, value=0.0):
    return np.full((num_dynamic_staff_types, num_intervs, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_dynamic_staff_prog_area_year_arr(
    num_dynamic_staff_types: int, num_prog_areas: int, num_years: int, value=0.0
):
    return np.full((num_dynamic_staff_types, num_prog_areas, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_dynamic_staff_prog_area_deliv_chan_year_arr(
    num_dynamic_staff_types: int, num_prog_areas: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full(
        (num_dynamic_staff_types, num_prog_areas, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE
    )


def IC_init_drug_supply_year_arr(num_drugs_supplies: int, num_years: int, value=0.0):
    return np.full((num_drugs_supplies, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_drug_supply_interv_costed_deliv_chan_year_arr(
    num_drugs_supplies: int, num_intervs: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full((num_drugs_supplies, num_intervs, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_drug_supply_costed_deliv_chan_year_arr(
    num_drugs_supplies: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full((num_drugs_supplies, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_drug_supply_prog_area_year_arr(num_drugs_supplies: int, num_prog_areas: int, num_years: int, value=0.0):
    return np.full((num_drugs_supplies, num_prog_areas, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_visit_costed_deliv_chan_year_arr(
    num_intervs: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full(
        (num_intervs, icc.IC_NUM_VISIT_TYPES, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE
    )


def IC_init_visit_costed_deliv_chan_arr(num_costed_deliv_chans: int, value=0.0):
    return np.full((icc.IC_NUM_VISIT_TYPES, num_costed_deliv_chans), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_costing_sect_costed_deliv_chan_year_arr(
    num_intervs: int, num_costing_sect: int, num_costed_deliv_chans: int, num_years: int, value=0.0
):
    return np.full((num_intervs, num_costing_sect, num_costed_deliv_chans, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_interv_costing_sect_year_arr(num_intervs: int, num_costing_sect: int, num_years: int, value=0.0):
    return np.full((num_intervs, num_costing_sect, num_years), value, gbc.GB_CALC_FLOAT_TYPE)


def IC_init_costing_sect_year_arr(num_costing_sect: int, num_years: int, value=0.0):
    return np.full((num_costing_sect, num_years), value, gbc.GB_CALC_FLOAT_TYPE)
