import pandas as pd
import numpy as np
from typing import Tuple, Dict
from .utils.constants import ORIGINAL_PREDICTIONS


def get_adjustments(adjust_table: pd.DataFrame,
                    data: pd.DataFrame,
                    sens_col: str,
                    score_col: str,
                    lam: float = 1.0) -> np.ndarray:
    """
    Given the full adjustment table, this function calculates a partial adjustment table based on the lambda value
    provided by the downstream fairness algorithm.

    Note the sensitive attribute column should be preprocessed to be able to convert to integers. This is so we can
    properly access the columns relating to the group. Moreover, note that the scores are expected to be float values
    that will be converted to the hundreths place to match the grid that was used to generate the adjustment table. This
    creates a row-column entry access point for the adjustment table.

    :param adjust_table: saved table of adjustments from get_adjustment_table_binary
    :param data: the data you want to get adjustments for
    :param sens_col: the string attribute that calls the sensitive attribute column
    :param score_col: the string attribute that calls the score column
    :param lam: lambda value generated by get_lambdas

    :raises ValueError: Lambda values must be between 0 and 1, inclusive

    :return adjusts: an array of adjustments for each row in the associated dataset
    """
    if not (0 <= lam <= 1):
        raise ValueError("Lambda values must be between 0 and 1, inclusive.")

    adjs_col_list = adjust_table.columns[adjust_table.columns.str.contains(
        'adjust')]
    adjs = adjust_table[adjs_col_list].values
    groups = _adjust_group_names(data[sens_col]).values
    scores = data[score_col].values
    adjusts = lam * adjs[(scores * 100).astype(int), groups.astype(int)]
    return adjusts


def get_adjusted_scores(adjust_table: pd.DataFrame,
                        data: pd.DataFrame,
                        sens_col: str,
                        score_col: str,
                        lam: float = 1.0) -> np.ndarray:
    """
    Given the full adjustment table, this function calculates a partial adjustment table based on the lambda value
    provided by the downstream fairness algorithm. Then adjust the scores using the partial adjustments.

    Note the sensitive attribute column should be preprocessed to be able to convert to integers. This is so we can
    properly access the columns relating to the group. Moreover, note that the scores are expected to be float values
    that will be converted to the hundreths place to match the grid that was used to generate the adjustment table. This
    creates a row-column entry access point for the adjustment table.

    :param adjust_table: saved table of adjustments from get_adjustment_table_binary
    :param data: the data you want to get adjustments for
    :param sens_col: the string attribute that calls the sensitive attribute column
    :param score_col: the string attribute that calls the score column
    :param lam: lambda value generated from get_lambdas

    :raises ValueError: Lambda values must be between 0 and 1, inclusive

    :return adjusts: an array of adjustments for each row in the associated dataset
    """
    if not (0 <= lam <= 1):
        raise ValueError("Lambda values must be between 0 and 1, inclusive.")

    adjs_col_list = adjust_table.columns[adjust_table.columns.str.contains(
        'adjust')]
    adjs = adjust_table[adjs_col_list].values
    original = adjust_table[ORIGINAL_PREDICTIONS].values
    groups = _adjust_group_names(data[sens_col]).values
    scores = data[score_col].values
    adjusted_scores = original[(scores * 100).astype(int)] + \
        lam * adjs[(scores * 100).astype(int), groups.astype(int)]
    return adjusted_scores


def get_bias_mitigator(dataset: pd.DataFrame, sens_col: str,
                       score_col: str, label_col: str) -> Tuple[pd.DataFrame, Dict[str, float]]:
    """
    First generate an adjustment table that acts as the full probability adjustment to achieve strong demographic
    parity. Then, based on the adjustments and the training dataset (usually the reference set), finds lambdas that
    minimize disparity on some fairness metric. This outputs the adjustment table and the dictionary of lambdas for
    each fairness metric (true positive rate, false positive rate, and equalized odds).

    :param dataset: a DataFrame consisting of columns that represent labels, scores, and groups.
    :param sens_col: identifier for the group column
    :param score_col: identifier for the score column
    :param label_col: string value to indicate which column has the labels

    :return: A dataframe of full adjustment values for each probability and each group
        and a dictionary of optimal lambdas for each fairness metric
    """
    from .train import get_adjustment_table, get_lambdas

    adjusts = get_adjustment_table(dataset, sens_col, score_col)
    lambdas = get_lambdas(adjusts, dataset, sens_col, score_col, label_col)
    return adjusts, lambdas


def _adjust_group_names(groups: pd.Series) -> pd.Series:
    """
    Helper function to get the group names into a state, where we can appropriately index into arrays using the
    number of groups in a sensitive attribute column.

    :param groups: the column containing the group names

    :return: a modified pandas Series that contains integers corresponding to group names.
    """
    unique_groups = groups.unique()
    # Need to sort the groups because np.unique() sorts the groups
    unique_groups = np.sort(unique_groups)
    group_map = {}
    i = 0
    for group in unique_groups:
        group_map[group] = i
        i += 1
    new_groups = groups.map(group_map)
    return new_groups
