"""Module containing functions that standardize the format of GBD outputs."""
from typing import Mapping, Iterable
from itertools import chain

import pandas as pd
import numpy as np
from gbd_mapping import (causes, risk_factors, etiologies, coverage_gaps,
                         Cause, Risk, Sequela, CoverageGap, Etiology, ModelableEntity)
from gbd_mapping.id import scalar, UNKNOWN

from vivarium_inputs.mapping_extension import HealthcareEntity, HealthTechnology

try:
    from vivarium_gbd_access import gbd
except ModuleNotFoundError:
    class GbdDummy:
        def __getattr__(self, item):
            raise ModuleNotFoundError("Required package vivarium_gbd_access not found.")
    gbd = GbdDummy()

GBD_ROUND_ID_MAP = {3: 'GBD_2015', 4: 'GBD_2016'}


class DataError(Exception):
    """Base exception for errors in data loading."""
    pass


class InvalidQueryError(DataError):
    """Exception raised when the user makes an invalid request for data (e.g. exposures for a sequela)."""
    pass


class UnhandledDataError(DataError):
    """Exception raised when we receive data from the databases that we don't know how to handle."""
    pass


class DataMissingError(DataError):
    """Exception raised when data has unhandled missing entries."""
    pass


class DuplicateDataError(DataError):
    """Exception raised when data has duplication in the index."""
    pass


def standardize_dimensions(data: pd.DataFrame, dimensions: pd.MultiIndex,
                           fill_na_value: Mapping[str, float]=None) -> pd.DataFrame:
    """Take input data and make it dense over the specified dimensions.

    The behavior of the function depends on which dimension is being considered and the nature of the sparsity:

    Cases:
    1) dimension missing - nothing needed
    2) fully dense in the expected range - nothing needed
    3) some sparsity:
       a) in the year dimension: interpolate
       b) in age/sex dimension and dense for a contiguous subset of the expected range and missing elsewhere: fill
    4) all others - panic

    Parameters
    ----------
    data :
        The data to standardize
    dimensions :
        A multi index whose individual index columns must be some subset of {'age', 'sex', 'year'}
        which represent dimensions of the data; and whose levels provide the expected extent of
        the data in those dimensions.
    fill_na_value :
        A mapping between

    Returns
    -------
        A dataframe with either full extent or no extent in all of the given dimensions.

    Raises
    ------
    UnhandledDataError :
        If the data is somehow malformed in a manner we don't deal with.
    """
    dimensions = dimensions.to_frame().reset_index(drop=True)

    dimensions = pd.concat([dimensions.assign(measure=measure) for measure in data.measure.unique()], ignore_index=True)
    draw_columns = [c for c in data.columns if 'draw_' in c]

    assert set(dimensions.columns) <= {'age_group_id', 'sex', 'year', 'measure', 'location'}
    assert set(data.columns) <= set(dimensions.columns.tolist() + draw_columns)
    assert set(data.measure.unique()) <= fill_na_value.keys()

    # Case 1: Remove any dimensions not present in the data.
    applicable_dimensions = dimensions[[c for c in dimensions.columns if c in data]].copy()
    # Case 4: Check preconditions and panic if it's an unhandled case.
    verify_well_formed(data, applicable_dimensions)

    # Case 3a: Interpolate over year.
    if 'year' in applicable_dimensions.columns:
        data = interpolate_years(data, applicable_dimensions)

    # Case 3b: Fill in missing data
    dimension_columns, extents = zip(*applicable_dimensions.items())
    expected_index = pd.MultiIndex.from_arrays(extents, names=dimension_columns)
    data = data.set_index(list(dimension_columns))

    missing = expected_index.difference(data.index)
    if not missing.empty:
        to_add = pd.DataFrame(columns=draw_columns, index=missing)
        to_add = to_add.reset_index().set_index(list(set(dimension_columns) - {'measure'})).sort_index()
        data = data.reset_index().set_index(list(set(dimension_columns) - {'measure'})).sort_index()
        for measure, fill in fill_na_value.items():
            to_add.loc[to_add['measure'] == measure, draw_columns] = fill
        data = data.append(to_add)

    return data.reset_index()


def verify_well_formed(data: pd.DataFrame, dimensions: pd.DataFrame):
    for dimension in dimensions.columns:
        if dimension == 'sex':
            continue
        existing = data[dimension].sort_values().drop_duplicates()
        expected = dimensions[dimension].sort_values().drop_duplicates()

        contiguous_overlap = ((expected >= existing.min()) & (expected <= existing.max()))
        if dimension == 'age_group_id' and not set(existing) == set(expected[contiguous_overlap]):
            raise UnhandledDataError(f'The data is malformed in the {dimension} dimension.')

        if dimension == 'year' and (existing.min() > expected.min() or existing.max() < expected.max()):
            raise UnhandledDataError("Cannot interpolate years without data on both edges")


def interpolate_years(data: pd.DataFrame, dimensions: pd.DataFrame) -> pd.DataFrame:

    existing_extent = data['year'].sort_values().drop_duplicates()
    expected_extent = dimensions['year'].sort_values().drop_duplicates()

    out = []
    for year in expected_extent:
        if year in existing_extent.values:
            out.append(data[data['year'] == year])
        else:
            from core_maths.interpolate import interpolate
            index_columns = list(dimensions.columns)
            value_columns = list(data.columns.difference(index_columns))

            previous_year = int(existing_extent[existing_extent < year].iloc[-1])
            next_year = int(existing_extent[existing_extent > year].iloc[0])
            start = data[data['year'] == previous_year].sort_values(index_columns).reset_index(drop=True)
            end = data[data['year'] == next_year].sort_values(index_columns).reset_index(drop=True)

            interpolated = interpolate(start, end, index_columns, 'year', value_columns, previous_year, next_year)
            interpolated = interpolated[interpolated['year'] == year]
            out.append(interpolated)

    return pd.concat(out, ignore_index=True)


def select_draw_data(data, draw, column_name, src_column=None):
    if column_name:
        if src_column is not None:
            if isinstance(src_column, str):
                column_map = {src_column.format(draw=draw): column_name}
            else:
                column_map = {src.format(draw=draw): dest for src, dest in zip(src_column, column_name)}
        else:
            column_map = {'draw_{draw}'.format(draw=draw): column_name}

        # if 'measure' is in columns, then keep it, else do
        # not keep it (need measure for the relative risk estimations)
        if 'parameter' in data.columns:
            keep_columns = ['year_id', 'age', 'sex_id', 'parameter'] + list(column_map.keys())
        else:
            keep_columns = ['year_id', 'age', 'sex_id'] + list(column_map.keys())

        data = data[keep_columns]
        data = data.rename(columns=column_map)

        return normalize_for_simulation(data)
    return data


def normalize_for_simulation(df):
    """
    Parameters
    ----------
    df : DataFrame
        dataframe to change

    Returns
    -------
    Returns a df with column year_id changed to year, and year_start and year_end
    created as bin ends around year_id with year_start set to year_id;
    sex_id changed to sex, and sex values changed from 1 and 2 to Male and Female

    Notes
    -----
    Used by -- load_data_from_cache

    Assumptions -- None

    Questions -- None

    Unit test in place? -- Yes
    """
    if "sex_id" in df:
        if set(df["sex_id"]) == {3}:
            df_m = df.copy()
            df_f = df.copy()
            df_m['sex'] = 'Male'
            df_f['sex'] = 'Female'
            df = pd.concat([df_m, df_f], ignore_index=True)
        else:
            df["sex"] = df.sex_id.map({1: "Male", 2: "Female", 3: "Both"}).astype(
                pd.api.types.CategoricalDtype(categories=["Male", "Female", "Both"], ordered=False))

        df = df.drop("sex_id", axis=1)

    if "year_id" in df:
        # FIXME: use central comp interpolation tools
        if 2006 in df.year_id.unique() and 2007 not in df.year_id.unique():
            df = df.loc[(df.year_id != 2006)]

        df = df.rename(columns={"year_id": "year"})
        idx = df.index

        mapping = df[['year']].drop_duplicates().sort_values(by="year")
        mapping['year_start'] = mapping['year']
        mapping['year_end'] = mapping['year'].shift(-1).fillna(mapping.year.max()+1)

        df = df.set_index("year", drop=False)
        mapping = mapping.set_index("year", drop=False)

        df[["year_start", "year_end"]] = mapping[["year_start", "year_end"]]

        df = df.set_index(idx)

    return df


def get_age_group_midpoint_from_age_group_id(df):
    """Creates an "age" column from the "age_group_id" column

    Parameters
    ----------
    df: df for which you want an age column that has an age_group_id column

    Returns
    -------
    df with an age column

    Notes
    -----
    Assumptions -- We assume that using a midpoint of age 82.5 for the 80+ year old age group is
    ok for the purposes of Vivarium. Everett proposed that we could get the life expectancy at age 80
    for each location and use that as the midpoint for the 80+ group, but Abie suggested that we
    keep the midpoint as 82.5 for now. GBD populations have data for each age group up until the
    age 95+ age group, at which point I'm assuming we can use 97.5 as the midpoint.
    """
    if df.empty:
        df['age'] = 0
        return df

    df = df.copy()
    idx = df.index
    mapping = gbd.get_age_bins()
    mapping = mapping.set_index('age_group_id')
    mapping['age'] = mapping[['age_group_years_start', 'age_group_years_end']].mean(axis=1)

    df = df.set_index('age_group_id')
    df[['age', 'age_group_start', 'age_group_end']] = mapping[['age', 'age_group_years_start', 'age_group_years_end']]

    # Assumption: We're using 82.5 as the midpoint for the age 80+ age group. May want to change in the future.
    df.loc[df.age == 102.5, 'age'] = 82.5

    df = df.set_index(idx)

    return df


def get_id_for_measure(entity: ModelableEntity, measure: str) -> int:
    """Selects the appropriate gbd id type for each entity and measure pair.

    Parameters
    ----------
    entity :
        A data containers from the `gbd_mapping` package.
    measure :
        A GBD measures requested for the provided entity.

    Returns
    -------
    The appropriate GBD id for use with central comp tools for the
    provided entity and measure.

    Raises
    ------
    InvalidQueryError
        If the entities passed are inconsistent with the requested measures.
    """
    measure_types = {
        'death': (Cause, 'gbd_id'),
        'prevalence': ((Cause, Sequela), 'gbd_id'),
        'incidence': ((Cause, Sequela), 'gbd_id'),
        'exposure': ((Risk, CoverageGap), 'gbd_id'),
        'exposure_standard_deviation': ((Risk, CoverageGap), 'gbd_id'),
        'relative_risk': ((Risk, CoverageGap), 'gbd_id'),
        'population_attributable_fraction': ((Risk, Etiology, CoverageGap), 'gbd_id'),
        'cause_specific_mortality': ((Cause,), 'gbd_id'),
        'excess_mortality': ((Cause,), 'gbd_id'),
        'annual_visits': (HealthcareEntity, 'utilization'),
        'disability_weight': (Sequela, 'gbd_id'),
        'remission': (Cause, 'dismod_id'),
        'cost': ((HealthcareEntity, HealthTechnology), 'cost'),
    }

    valid_types, id_attr = measure_types.get(measure, (None, None))

    if not valid_types or not id_attr:
        raise InvalidQueryError(f"You've requested an invalid measure: {measure}")
    elif not isinstance(entity, valid_types) or entity[id_attr] is UNKNOWN:
        raise InvalidQueryError(f"Entity {entity.name} has no data for measure '{measure}'")
    else:
        measure_id = entity[id_attr]

    return measure_id


def get_additional_id_columns(data, entity):
    id_column_map = {
        'cause': 'cause_id',
        'sequela': 'sequela_id',
        'covariate': 'covariate_id',
        'risk_factor': 'rei_id',
        'etiology': 'rei_id',
        'coverage_gap': 'coverage_gap',
        'healthcare_entity': 'healthcare_entity',
        'health_technology': 'health_technology',
    }
    out = {id_column_map[entity.kind]}
    if entity.kind == 'coverage_gap' and 'relative_risk' in data['measure'].unique():
        out.add('rei_id')
    # TODO: Why?
    out |= set(data.columns) & set(id_column_map.values())
    return out


def validate_data(data: pd.DataFrame, key_columns: Iterable[str]=None):
    """Validates that no data is missing and that the provided key columns make a valid (unique) index.

    Parameters
    ----------
    data:
        The data table to be validated.
    key_columns:
        An iterable of the column names used to uniquely identify a row in the data table.

    Raises
    ------
    DataMissingError
        If the data contains any null (NaN or NaT) values.
    DuplicatedDataError
        If the provided key columns are insufficient to uniquely identify a record in the data table.
    """

    #  check draw_cols only since we may have nan in coverage_gap data from aux_data
    draw_cols = [f'draw_{i}' for i in range(1000)]
    if np.any(data[draw_cols].isnull()):
        raise DataMissingError()

    if key_columns and np.any(data.duplicated(key_columns)):
        raise DuplicateDataError()


def get_age_group_ids(restrictions, measure='yld'):
    assert measure in ['yld', 'yll']

    age_restriction_map = {scalar(0.0): [2, None],
                           scalar(0.01): [3, 2],
                           scalar(0.10): [4, 3],
                           scalar(1.0): [5, 4],
                           scalar(5.0): [6, 5],
                           scalar(10.0): [7, 6],
                           scalar(15.0): [8, 7],
                           scalar(20.0): [9, 8],
                           scalar(30.0): [11, 10],
                           scalar(40.0): [13, 12],
                           scalar(45.0): [14, 13],
                           scalar(50.0): [15, 14],
                           scalar(55.0): [16, 15],
                           scalar(65.0): [18, 17],
                           scalar(95.0): [235, 32], }

    age_start, age_end = restrictions[f'{measure}_age_start'], restrictions[f'{measure}_age_end']
    min_age_group = age_restriction_map[age_start][0]
    max_age_group = age_restriction_map[age_end][1]
    return list(range(min_age_group, max_age_group + 1))


def get_measure_id(measure):
    name_measure_map = {'death': 1,
                        'DALY': 2,
                        'YLD': 3,
                        'YLL': 4,
                        'prevalence': 5,
                        'incidence': 6,
                        'remission': 7,
                        'excess_mortality': 9,
                        'proportion': 18,
                        'continuous': 19, }
    return name_measure_map[measure]


def standardize_all_age_groups(data: pd.DataFrame):
    if set(data.age_group_id) != {22}:
        return data
    whole_age_groups = gbd.get_age_group_id()
    missing_age_groups = set(whole_age_groups).difference(set(data.age_group_id))
    df = []
    for i in missing_age_groups:
        missing = data.copy()
        missing['age_group_id'] = i
        df.append(missing)

    return pd.concat(df, ignore_index=True)


def standardize_data(data: pd.DataFrame, fill_value: int) -> pd.DataFrame:
    # age_groups that we expect to exist for each risk
    whole_age_groups = gbd.get_age_group_id()
    sex_id = data.sex_id.unique()
    year_id = data.year_id.unique()
    location_id = data.location_id.unique()

    index_cols = ['year_id', 'location_id', 'sex_id', 'age_group_id']
    draw_cols = [c for c in data.columns if 'draw_' in c]

    other_cols = {c: data[c].unique() for c in data.dropna(axis=1).columns if c not in index_cols and c not in draw_cols}
    index_cols += [*other_cols.keys()]
    data = data.set_index(index_cols)

    # expected indexes to be in the data
    expected = pd.MultiIndex.from_product([year_id, location_id, sex_id, whole_age_groups] + [*other_cols.values()],
                                          names=(['year_id', 'location_id', 'sex_id', 'age_group_id'] + [
                                              *other_cols.keys()]))

    new_data = data.copy()
    missing = expected.difference(data.index)

    # assign dtype=float to prevent the artifact error with mixed dtypes
    to_add = pd.DataFrame({column: fill_value for column in draw_cols}, index=missing, dtype=float)

    new_data = new_data.append(to_add).sort_index()

    return new_data.reset_index()


def filter_to_most_detailed(data):
    for column, entity_list in [('cause_id', causes),
                                ('rei_id', chain(etiologies, risk_factors)),
                                ('coverage_gap_id', coverage_gaps)]:
        most_detailed = {e.gbd_id for e in entity_list if e is not None}
        if column in data:
            data = data[data[column].isin(most_detailed)]
    return data
