from numbers import Real
from typing import List, Set, Union, Tuple
import warnings

import pandas as pd

from gbd_mapping import causes, covariates, risk_factors, Cause, ModelableEntity, RiskFactor
from vivarium.framework.artifact import EntityKey
from vivarium_gbd_access import constants as gbd_constants, gbd
from vivarium_gbd_access.utilities import get_draws, query
from vivarium_inputs import globals as vi_globals, utilities as vi_utils, utility_data
from vivarium_inputs.mapping_extension import alternative_risk_factors, AlternativeRiskFactor
from vivarium_inputs.validation.raw import check_metadata

from vivarium_gates_iv_iron.constants.metadata import AGE_GROUP, GBD_2020_ROUND_ID


def get_entity(key: EntityKey) -> ModelableEntity:
    # Map of entity types to their gbd mappings.
    type_map = {
        'cause': causes,
        'covariate': covariates,
        'risk_factor': risk_factors,
        'alternative_risk_factor': alternative_risk_factors
    }
    return type_map[key.type][key.name]


def get_data(key: EntityKey, entity: ModelableEntity, location: str, source: str, gbd_id_type: str,
             age_group_ids: Set[int], gbd_round_id: int, decomp_step: str = 'iterative') -> pd.DataFrame:
    age_group_ids = list(age_group_ids)

    # from interface.get_measure
    # from vivarium_inputs.core.get_data
    location_id = utility_data.get_location_id(location) if isinstance(location, str) else location

    # from vivarium_inputs.core.get_{measure}
    # from vivarium_inputs.extract.extract_data
    check_metadata(entity, key.measure)

    # from vivarium_inputs.extract.extract_{measure}
    # from vivarium_gbd_access.gbd.get_{measure}
    data = get_draws(gbd_id_type=gbd_id_type,
                     gbd_id=entity.gbd_id,
                     source=source,
                     location_id=location_id,
                     sex_id=gbd_constants.SEX.MALE + gbd_constants.SEX.FEMALE,
                     age_group_id=age_group_ids,
                     gbd_round_id=gbd_round_id,
                     decomp_step=decomp_step,
                     status='best')
    return data


def validate_and_reshape_gbd_data(data: pd.DataFrame, entity: ModelableEntity, key: EntityKey,
                                  location: str, gbd_round_id: int, age_group_ids: List[int] = None) -> pd.DataFrame:

    # from vivarium_inputs.core.get_data
    data = vi_utils.reshape(data, value_cols=vi_globals.DRAW_COLUMNS)

    # from interface.get_measure
    data = scrub_gbd_conventions(data, location, age_group_ids)

    estimation_years = get_gbd_estimation_years(gbd_round_id)
    validation_years = pd.DataFrame({'year_start': range(min(estimation_years), max(estimation_years) + 1)})
    validation_years['year_end'] = validation_years['year_start'] + 1

    # validate_for_simulation(data, entity, key.measure, location, years=validation_years,
    #                         age_bins=get_gbd_age_bins(age_group_ids))
    data = vi_utils.split_interval(data, interval_column='age', split_column_prefix='age')
    data = vi_utils.split_interval(data, interval_column='year', split_column_prefix='year')
    data = vi_utils.sort_hierarchical_data(data).droplevel('location')
    return data


def normalize_age_and_years(data: pd.DataFrame, fill_value: Real = None,
                            cols_to_fill: List[str] = vi_globals.DRAW_COLUMNS,
                            gbd_round_id: int = GBD_2020_ROUND_ID,
                            age_group_ids: List[int] = AGE_GROUP.GBD_2020) -> pd.DataFrame:
    data = vi_utils.normalize_sex(data, fill_value, cols_to_fill)

    # vi_inputs.normalize_year(data)
    binned_years = get_gbd_estimation_years(gbd_round_id)
    years = {'annual': list(range(min(binned_years), max(binned_years) + 1)), 'binned': binned_years}

    if 'year_id' not in data:
        # Data doesn't vary by year, so copy for each year.
        df = []
        for year in years['annual']:
            fill_data = data.copy()
            fill_data['year_id'] = year
            df.append(fill_data)
        data = pd.concat(df, ignore_index=True)
    elif set(data.year_id) == set(years['binned']):
        data = vi_utils.interpolate_year(data)
    else:  # set(data.year_id.unique()) == years['annual']
        pass

    # Dump extra data.
    data = data[data.year_id.isin(years['annual'])]

    data = _normalize_age(data, fill_value, cols_to_fill, age_group_ids)
    return data


def _normalize_age(data: pd.DataFrame, fill_value: Real, cols_to_fill: List[str],
                   age_group_ids: List[int] = None) -> pd.DataFrame:
    data_ages = set(data.age_group_id.unique()) if 'age_group_id' in data.columns else set()
    gbd_ages = set(utility_data.get_age_group_ids()) if not age_group_ids else set(age_group_ids)

    if not data_ages:
        # Data does not correspond to individuals, so no age column necessary.
        pass
    elif data_ages == {vi_globals.SPECIAL_AGES['all_ages']}:
        # Data applies to all ages, so copy.
        dfs = []
        for age in gbd_ages:
            missing = data.copy()
            missing.loc[:, 'age_group_id'] = age
            dfs.append(missing)
        data = pd.concat(dfs, ignore_index=True)
    elif data_ages < gbd_ages:
        # Data applies to subset, so fill other ages with fill value.
        key_columns = list(data.columns.difference(cols_to_fill))
        key_columns.remove('age_group_id')
        expected_index = pd.MultiIndex.from_product([data[c].unique() for c in key_columns] + [gbd_ages],
                                                    names=key_columns + ['age_group_id'])

        data = (data.set_index(key_columns + ['age_group_id'])
                .reindex(expected_index, fill_value=fill_value)
                .reset_index())
    else:  # data_ages == gbd_ages
        pass
    return data


def get_gbd_estimation_years(gbd_round_id: int) -> List[int]:
    """Gets the estimation years for a particular gbd round."""
    from db_queries import get_demographics
    warnings.filterwarnings("default", module="db_queries")

    return get_demographics(gbd_constants.CONN_DEFS.EPI, gbd_round_id=gbd_round_id)['year_id']


def scrub_gbd_conventions(data: pd.DataFrame, location: str, age_group_ids: List[int] = None) -> pd.DataFrame:
    data = vi_utils.scrub_location(data, location)
    data = vi_utils.scrub_sex(data)
    data = _scrub_age(data, age_group_ids)
    data = vi_utils.scrub_year(data)
    data = vi_utils.scrub_affected_entity(data)
    return data


def process_exposure(data: pd.DataFrame, key: str, entity: Union[RiskFactor, AlternativeRiskFactor],
                     location: str, gbd_round_id: int, age_group_ids: List[int] = None) -> pd.DataFrame:
    data['rei_id'] = entity.gbd_id

    # from vivarium_inputs.extract.extract_exposure
    allowable_measures = [vi_globals.MEASURES['Proportion'], vi_globals.MEASURES['Continuous'],
                          vi_globals.MEASURES['Prevalence']]
    proper_measure_id = set(data.measure_id).intersection(allowable_measures)
    if len(proper_measure_id) != 1:
        raise vi_globals.DataAbnormalError(f'Exposure data have {len(proper_measure_id)} measure id(s). '
                                           f'Data should have exactly one id out of {allowable_measures} '
                                           f'but came back with {proper_measure_id}.')
    data = data[data.measure_id == proper_measure_id.pop()]

    # from vivarium_inputs.core.get_exposure
    data = data.drop('modelable_entity_id', 'columns')

    if entity.name in vi_globals.EXTRA_RESIDUAL_CATEGORY:
        # noinspection PyUnusedLocal
        cat = vi_globals.EXTRA_RESIDUAL_CATEGORY[entity.name]
        data = data.drop(labels=data.query('parameter == @cat').index)
        data[vi_globals.DRAW_COLUMNS] = data[vi_globals.DRAW_COLUMNS].clip(lower=vi_globals.MINIMUM_EXPOSURE_VALUE)

    if entity.distribution in ['dichotomous', 'ordered_polytomous', 'unordered_polytomous']:
        tmrel_cat = utility_data.get_tmrel_category(entity)
        exposed = data[data.parameter != tmrel_cat]
        unexposed = data[data.parameter == tmrel_cat]
        #  FIXME: We fill 1 as exposure of tmrel category, which is not correct.
        data = pd.concat([normalize_age_and_years(exposed, fill_value=0, gbd_round_id=gbd_round_id),
                          normalize_age_and_years(unexposed, fill_value=1, gbd_round_id=gbd_round_id)],
                         ignore_index=True)

        # normalize so all categories sum to 1
        cols = list(set(data.columns).difference(vi_globals.DRAW_COLUMNS + ['parameter']))
        data = data.set_index(cols + ['parameter'])
        sums = (
            data.groupby(cols)[vi_globals.DRAW_COLUMNS].sum()
                .reindex(index=data.index)
        )
        data = data.divide(sums).reset_index()
    else:
        data = vi_utils.normalize(data, fill_value=0)

    data = data.filter(vi_globals.DEMOGRAPHIC_COLUMNS + vi_globals.DRAW_COLUMNS + ['parameter'])
    data = validate_and_reshape_gbd_data(data, entity, key, location, gbd_round_id, age_group_ids)
    return data


def _scrub_age(data: pd.DataFrame, age_group_ids: List[int] = None) -> pd.DataFrame:
    if 'age_group_id' in data.index.names:
        age_bins = get_gbd_age_bins(age_group_ids).set_index('age_group_id')
        id_levels = data.index.levels[data.index.names.index('age_group_id')]
        interval_levels = [pd.Interval(age_bins.age_start[age_id], age_bins.age_end[age_id], closed='left')
                           for age_id in id_levels]
        data.index = data.index.rename(names='age', level='age_group_id').set_levels(levels=interval_levels, level='age')
    return data


def get_gbd_age_bins(age_group_ids: List[int] = None) -> pd.DataFrame:
    # If no age group ids are specified, use the standard GBD 2019 age bins
    if not age_group_ids:
        age_group_ids = gbd.get_age_group_id()
    # from gbd.get_age_bins()
    q = f"""
                SELECT age_group_id,
                       age_group_years_start,
                       age_group_years_end,
                       age_group_name
                FROM age_group
                WHERE age_group_id IN ({','.join([str(a) for a in age_group_ids])})
                """
    raw_age_bins = query(q, 'shared')

    # from utility_data.get_age_bins()
    age_bins = (
        raw_age_bins[['age_group_id', 'age_group_name', 'age_group_years_start', 'age_group_years_end']]
        .rename(columns={'age_group_years_start': 'age_start', 'age_group_years_end': 'age_end'})
    )

    # set age start for birth prevalence age bin to -1 to avoid validation issues
    age_bins.loc[age_bins['age_end'] == 0.0, 'age_start'] = -1.0
    return age_bins
