# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from __future__ import absolute_import, division, print_function, unicode_literals

import itertools
import logging
import math
from fractions import Fraction
from functools import reduce
from typing import Any, Callable, Dict, List, Tuple, Union

import numpy as np
import pandas as pd
from balance import adjustment as balance_adjustment, util as balance_util

logger: logging.Logger = logging.getLogger(__package__)


# TODO: Add options for only marginal distributions input
def _run_ipf_numpy(
    original: np.ndarray,
    target_margins: List[np.ndarray],
    convergence_rate: float,
    max_iteration: int,
    rate_tolerance: float,
) -> Tuple[np.ndarray, int, pd.DataFrame]:
    """Run iterative proportional fitting on a NumPy array.

    This reimplements the minimal subset of the :mod:`ipfn` package that is
    required for balance's usage.  The original implementation spends most of
    its time looping in pure Python over every slice of the contingency table,
    which is prohibitively slow for the high-cardinality problems we test
    against.  The logic here mirrors the algorithm used by ``ipfn.ipfn`` but
    applies the adjustments in a vectorised manner, yielding identical
    numerical results with a fraction of the runtime.

    The caller is expected to pass ``target_margins`` that correspond to
    single-axis marginals (which is how :func:`rake` constructs the inputs).
    """

    if original.ndim == 0:
        raise ValueError("`original` must have at least one dimension")

    table = np.asarray(original, dtype=np.float64)
    margins = [np.asarray(margin, dtype=np.float64) for margin in target_margins]

    # Pre-compute shapes and axes that are repeatedly required during the
    # iterative updates.  Each entry in ``axis_shapes`` represents how a
    # one-dimensional scaling factor should be reshaped in order to broadcast
    # along the appropriate axis of ``table``.
    axis_shapes: List[Tuple[int, ...]] = []
    sum_axes: List[Tuple[int, ...]] = []
    for axis in range(table.ndim):
        shape = [1] * table.ndim
        shape[axis] = table.shape[axis]
        axis_shapes.append(tuple(shape))
        sum_axes.append(tuple(i for i in range(table.ndim) if i != axis))

    conv = np.inf
    old_conv = -np.inf
    conv_history: List[float] = []
    iteration = 0

    while (
        iteration <= max_iteration
        and conv > convergence_rate
        and abs(conv - old_conv) > rate_tolerance
    ):
        old_conv = conv

        # Sequentially update the table for each marginal.  Because the
        # marginals correspond to single axes we can compute all scaling
        # factors at once, avoiding the expensive Python loops present in the
        # reference implementation.
        for axis, margin in enumerate(margins):
            current = table.sum(axis=sum_axes[axis])
            factors = np.ones_like(margin, dtype=np.float64)
            np.divide(margin, current, out=factors, where=current != 0)
            table *= factors.reshape(axis_shapes[axis])

        # Measure convergence using the same criterion as ``ipfn.ipfn``.  The
        # implementation there keeps the maximum absolute proportional
        # difference while naturally ignoring NaNs (which arise for 0/0).  We
        # match that behaviour by treating NaNs as zero deviation.
        conv = 0.0
        for axis, margin in enumerate(margins):
            current = table.sum(axis=sum_axes[axis])
            with np.errstate(divide="ignore", invalid="ignore"):
                diff = np.abs(np.divide(current, margin) - 1.0)
            current_conv = float(np.nanmax(diff)) if diff.size else 0.0
            if math.isnan(current_conv):
                current_conv = 0.0
            if current_conv > conv:
                conv = current_conv

        conv_history.append(conv)
        iteration += 1

    converged = int(iteration <= max_iteration)
    iterations_df = pd.DataFrame(
        {"iteration": range(len(conv_history)), "conv": conv_history}
    ).set_index("iteration")

    return table, converged, iterations_df


def rake(
    sample_df: pd.DataFrame,
    sample_weights: pd.Series,
    target_df: pd.DataFrame,
    target_weights: pd.Series,
    variables: Union[List[str], None] = None,
    transformations: Union[Dict[str, Callable[..., Any]], str] = "default",
    na_action: str = "add_indicator",
    max_iteration: int = 1000,
    convergence_rate: float = 0.0005,
    rate_tolerance: float = 1e-8,
    weight_trimming_mean_ratio: Union[float, int, None] = None,
    weight_trimming_percentile: Union[float, None] = None,
    keep_sum_of_weights: bool = True,
    *args: Any,
    **kwargs: Any,
) -> Dict[str, Any]:
    """
    Perform raking (using the iterative proportional fitting algorithm).
    See: https://en.wikipedia.org/wiki/Iterative_proportional_fitting

    Returns weights normalised to sum of target weights

    Arguments:
    sample_df --- (pandas dataframe) a dataframe representing the sample.
    sample_weights --- (pandas series) design weights for sample.
    target_df ---  (pandas dataframe) a dataframe representing the target.
    target_weights --- (pandas series) design weights for target.
    variables ---  (list of strings) list of variables to include in the model.
                   If None all joint variables of sample_df and target_df are used.
    transformations --- (dict) what transformations to apply to data before fitting the model.
                               Default is "default" (see apply_transformations function).
    na_action --- (string) what to do with NAs. Default is "add_indicator", which adds NaN as a
                  group (called "__NaN__") for each weighting variable (post-transformation);
                  "drop" removes rows with any missing values on any variable from both sample
                  and target.
    max_iteration --- (int) maximum number of iterations for iterative proportional fitting algorithm
    convergence_rate --- (float) convergence criteria; the maximum difference in proportions between
                          sample and target marginal distribution on any covariate in order
                          for algorithm to converge.
    rate_tolerance --- (float) convergence criteria; if convergence rate does not move more
                               than this amount than the algorithm is also considered to
                               have converged.
    weight_trimming_mean_ratio --- (float, int, optional) upper bound for weights expressed as a
                                   multiple of the mean weight. Delegated to
                                   :func:`balance.adjustment.trim_weights`.
    weight_trimming_percentile --- (float, optional) percentile limit(s) for winsorisation.
                                   Delegated to :func:`balance.adjustment.trim_weights`.
    keep_sum_of_weights --- (bool, optional) preserve the sum of weights during trimming before
                            rescaling to the target total. Defaults to True.

    Returns:
    A dictionary including:
    "weight" --- The weights for the sample.
    "model" --- parameters of the model: iterations (dataframe with iteration numbers and
                convergence rate information at all steps), converged (Flag with the output
                status: 0 for failure and 1 for success).

    Examples:
    .. code-block:: python

        import pandas as pd
        from balance.weighting_methods.rake import rake
        sample_df = pd.DataFrame({"x": ["a", "b"]})
        target_df = pd.DataFrame({"x": ["a", "b"]})
        sample_weights = pd.Series([1.0, 1.0])
        target_weights = pd.Series([1.0, 1.0])
        result = rake(sample_df, sample_weights, target_df, target_weights, variables=["x"])
        result["weight"].tolist()
        # [1.0, 1.0]
    """
    assert (
        "weight" not in sample_df.columns.values
    ), "weight shouldn't be a name for covariate in the sample data"
    assert (
        "weight" not in target_df.columns.values
    ), "weight shouldn't be a name for covariate in the target data"

    # TODO: move the input checks into separate funnction for rake, ipw, poststratify
    assert isinstance(sample_df, pd.DataFrame), "sample_df must be a pandas DataFrame"
    assert isinstance(target_df, pd.DataFrame), "target_df must be a pandas DataFrame"
    assert isinstance(
        sample_weights, pd.Series
    ), "sample_weights must be a pandas Series"
    assert isinstance(
        target_weights, pd.Series
    ), "target_weights must be a pandas Series"
    assert sample_df.shape[0] == sample_weights.shape[0], (
        "sample_weights must be the same length as sample_df"
        f"{sample_df.shape[0]}, {sample_weights.shape[0]}"
    )
    assert target_df.shape[0] == target_weights.shape[0], (
        "target_weights must be the same length as target_df"
        f"{target_df.shape[0]}, {target_weights.shape[0]}"
    )

    variables = balance_util.choose_variables(sample_df, target_df, variables=variables)

    logger.debug(f"Join variables for sample and target: {variables}")

    sample_df = sample_df.loc[:, variables]
    target_df = target_df.loc[:, variables]

    assert len(variables) > 1, (
        "Must weight on at least two variables for raking. "
        f"Currently have variables={variables} only"
    )

    sample_df, target_df = balance_adjustment.apply_transformations(
        (sample_df, target_df), transformations
    )

    # TODO: separate into a function that handles NA (for rake, ipw, poststratify)
    if na_action == "drop":
        (sample_df, sample_weights) = balance_util.drop_na_rows(
            sample_df, sample_weights, "sample"
        )
        (target_df, target_weights) = balance_util.drop_na_rows(
            target_df, target_weights, "target"
        )
    elif na_action == "add_indicator":
        from balance.util import _safe_fillna_and_infer

        target_df = _safe_fillna_and_infer(target_df, "__NaN__")
        sample_df = _safe_fillna_and_infer(sample_df, "__NaN__")
    else:
        raise ValueError("`na_action` must be 'add_indicator' or 'drop'")

    # Alphabetize variables to ensure consistency across covariate order
    # (ipfn algorithm is iterative and variable order can matter on the margins)
    alphabetized_variables = list(variables)
    alphabetized_variables.sort()

    logger.debug(
        f"Alphabetized variable order is as follows: {alphabetized_variables}."
    )

    # Cast all data types as string to be explicit about each unique value
    # being its own group and to handle that `fillna()` above creates
    # series of type Object, which won't work for the ipfn script
    categories = []
    for variable in alphabetized_variables:
        target_df[variable] = target_df[variable].astype(str)
        sample_df[variable] = sample_df[variable].astype(str)

        sample_var_set = set(sample_df[variable].unique())
        target_var_set = set(target_df[variable].unique())
        sample_over_set = sample_var_set - target_var_set
        target_over_set = target_var_set - sample_var_set
        if len(sample_over_set):
            raise ValueError(
                "All variable levels in sample must be present in target. "
                f"'{variable}' in target is missing these levels: {sample_over_set}."
            )
        if len(target_over_set):
            logger.warning(
                f"'{variable}' has more levels in target than in sample. "
                f"'{variable}' in sample is missing these levels: {target_over_set}. "
                "These levels are treated as if they do not exist for that variable."
            )
        categories.append(sorted(sample_var_set.intersection(target_var_set)))

    logger.info(
        f"Final covariates and levels that will be used in raking: {dict(zip(alphabetized_variables, categories))}."
    )

    target_df = target_df.assign(weight=target_weights)
    sample_df = sample_df.assign(weight=sample_weights)

    sample_sum_weights = sample_df["weight"].sum()
    target_sum_weights = target_df["weight"].sum()

    # Calculate {# covariates}-dimensional array representation of the sample
    # for the ipfn algorithm

    grouped_sample_series = sample_df.groupby(alphabetized_variables)["weight"].sum()
    index = pd.MultiIndex.from_product(categories, names=alphabetized_variables)
    grouped_sample_full = grouped_sample_series.reindex(index, fill_value=0)
    m_sample = grouped_sample_full.to_numpy().reshape([len(c) for c in categories])
    m_fit_input = m_sample.copy()

    # Calculate target margins for ipfn
    target_margins = []
    for col, cats in zip(alphabetized_variables, categories):
        sums = (
            target_df.groupby(col)["weight"].sum()
            / target_sum_weights
            * sample_sum_weights
        )
        sums = sums.reindex(cats, fill_value=0)
        target_margins.append(sums.values)

    logger.debug(
        "Raking algorithm running following settings: "
        f" convergence_rate: {convergence_rate}; max_iteration: {max_iteration}; rate_tolerance: {rate_tolerance}"
    )

    # returns array with joint distribution of covariates and total weight
    # for that specific set of covariates
    # no longer uses the dataframe version of the ipfn algorithm
    # due to incompatability with latest Python versions
    m_fit, converged, iterations = _run_ipf_numpy(
        m_fit_input,
        target_margins,
        convergence_rate,
        max_iteration,
        rate_tolerance,
    )

    logger.debug(
        f"Raking algorithm terminated with following convergence: {converged}; "
        f"and iteration meta data: {iterations}."
    )

    if not converged:
        logger.warning("Maximum iterations reached, convergence was not achieved")

    combos = list(itertools.product(*categories))
    fit = pd.DataFrame(combos, columns=alphabetized_variables)
    fit["rake_weight"] = m_fit.flatten()

    raked = pd.merge(
        sample_df.reset_index(),
        fit,
        how="left",
        on=alphabetized_variables,
    )

    raked_rescaled = pd.merge(
        raked,
        grouped_sample_series.reset_index().rename(
            columns={"weight": "total_survey_weight"}
        ),
        how="left",
        on=alphabetized_variables,
    ).set_index("index")

    raked_rescaled["rake_weight"] = (
        raked_rescaled["rake_weight"] / raked_rescaled["total_survey_weight"]
    )

    w = balance_adjustment.trim_weights(
        raked_rescaled["rake_weight"],
        target_sum_weights=target_sum_weights,
        weight_trimming_mean_ratio=weight_trimming_mean_ratio,
        weight_trimming_percentile=weight_trimming_percentile,
        keep_sum_of_weights=keep_sum_of_weights,
    ).rename("rake_weight")
    return {
        "weight": w,
        "model": {
            "method": "rake",
            "iterations": iterations,
            "converged": converged,
            "perf": {"prop_dev_explained": np.array([np.nan])},
            # TODO: fix functions that use the perf and remove it from here
        },
    }


def _lcm(a: int, b: int) -> int:
    """
    Calculates the least common multiple (LCM) of two integers.

    The least common multiple (LCM) of two or more numbers is the smallest positive integer that is divisible by each of the given numbers.
    In other words, it is the smallest multiple that the numbers have in common.
    The LCM is useful when you need to find a common denominator for fractions or synchronize repeating events with different intervals.

    For example, let's find the LCM of 4 and 6:

    The multiples of 4 are: 4, 8, 12, 16, 20, ...
    The multiples of 6 are: 6, 12, 18, 24, 30, ...
    The smallest multiple that both numbers have in common is 12, so the LCM of 4 and 6 is 12.

    The calculation is based on the property that the product of two numbers is equal to the product of their LCM and GCD: a * b = LCM(a, b) * GCD(a, b).
    (proof: https://math.stackexchange.com/a/589299/1406)

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The least common multiple of the two integers.

    Implementation note:
        Uses the mathematical relationship between LCM and GCD to efficiently compute the result.
        The absolute value ensures the LCM is always positive, even for negative inputs.
    """
    # NOTE: this function uses math.gcd which calculates the greatest common divisor (GCD) of two integers.
    # The greatest common divisor (GCD) of two or more integers is the largest positive integer that divides each of the given integers without leaving a remainder.
    # In other words, it is the largest common factor of the given integers.
    # For example, the GCD of 12 and 18 is 6, since 6 is the largest integer that divides both 12 and 18 without leaving a remainder.
    # Similarly, the GCD of 24, 36, and 48 is 12, since 12 is the largest integer that divides all three of these numbers without leaving a remainder.
    return abs(a * b) // math.gcd(a, b)


def _proportional_array_from_dict(
    input_dict: Dict[str, float], max_length: int = 10000
) -> List[str]:
    """
    Generates a proportional array based on the input dictionary.

    Args:
        input_dict (Dict[str, float]): A dictionary where keys are strings and values are their proportions (float).
        max_length (int): check if the length of the output exceeds the max_length. If it does, it will be scaled down to that length. Default is 10k.

    Returns:
        A list of strings where each key is repeated according to its proportion.

    Examples:
    .. code-block:: python

            _proportional_array_from_dict({"a":0.2, "b":0.8})
                # ['a', 'b', 'b', 'b', 'b']
            _proportional_array_from_dict({"a":0.5, "b":0.5})
                # ['a', 'b']
            _proportional_array_from_dict({"a":1/3, "b":1/3, "c": 1/3})
                # ['a', 'b', 'c']
            _proportional_array_from_dict({"a": 3/8, "b": 5/8})
                # ['a', 'a', 'a', 'b', 'b', 'b', 'b', 'b']
            _proportional_array_from_dict({"a":3/5, "b":1/5, "c": 2/10})
                # ['a', 'a', 'a', 'b', 'c']
    """
    # Filter out items with zero value
    filtered_dict = {k: v for k, v in input_dict.items() if v > 0}

    # Normalize the values if they don't sum to 1
    sum_values = sum(filtered_dict.values())
    if sum_values != 1:
        filtered_dict = {k: v / sum_values for k, v in filtered_dict.items()}

    # Convert the proportions to fractions
    fraction_dict = {
        k: Fraction(v).limit_denominator() for k, v in filtered_dict.items()
    }

    # Find the least common denominator (LCD) of the fractions
    lcd = 1
    for fraction in fraction_dict.values():
        lcd *= fraction.denominator // math.gcd(lcd, fraction.denominator)

    # Calculate the scaling factor based on max_length
    scaling_factor = min(1, max_length / lcd)

    result = []
    for key, fraction in fraction_dict.items():
        # Calculate the count for each key based on its proportion and scaling_factor
        k_count = round((fraction * lcd).numerator * scaling_factor)
        # Extend the result array with the key repeated according to its count
        result.extend([key] * k_count)

    return result


def _find_lcm_of_array_lengths(arrays: Dict[str, List[str]]) -> int:
    """
    Finds the least common multiple (LCM) of the lengths of arrays in the input dictionary.

    Args:
        arrays: A dictionary where keys are strings and values are lists of strings.

    Returns:
        The LCM of the lengths of the arrays in the input dictionary.

    Example:
    .. code-block:: python

            arrays = {
                        "v1": ["a", "b", "b", "c"],
                        "v2": ["aa", "bb"]
                    }
            _find_lcm_of_array_lengths(arrays)
                # 4

            arrays = {
                        "v1": ["a", "b", "b", "c"],
                        "v2": ["aa", "bb"],
                        "v3": ["a1", "a2", "a3"]
                    }
            _find_lcm_of_array_lengths(arrays)
                # 12
    """
    array_lengths = [len(arr) for arr in arrays.values()]
    lcm_length = reduce(lambda x, y: _lcm(x, y), array_lengths)
    return lcm_length


def _realize_dicts_of_proportions(
    dict_of_dicts: Dict[str, Dict[str, float]],
) -> Dict[str, List[str]]:
    """
    Generates proportional arrays of equal length for each input dictionary.

    This can be used to get an input dict of proportions of values, and produce a dict with arrays that realizes these proportions.
    It can be used as input to the Sample object so it could be used for running raking.

    Args:
        dict_of_dicts: A dictionary of dictionaries, where each key is a string and
                   each value is a dictionary with keys as strings and values as their proportions (float).

    Returns:
        A dictionary with the same keys as the input and equal length arrays as values.

    Examples:
    .. code-block:: python

            dict_of_dicts = {
                "v1": {"a": 0.2, "b": 0.6, "c": 0.2},
                "v2": {"aa": 0.5, "bb": 0.5}
            }

            realize_dicts_of_proportions(dict_of_dicts)
            # {'v1': ['a', 'b', 'b', 'b', 'c', 'a', 'b', 'b', 'b', 'c'], 'v2': ['aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb']}


            dict_of_dicts = {
                "v1": {"a": 0.2, "b": 0.6, "c": 0.2},
                "v2": {"aa": 0.5, "bb": 0.5},
                "v3": {"A": 0.2, "B": 0.8},
            }
            realize_dicts_of_proportions(dict_of_dicts)
                # {'v1': ['a', 'b', 'b', 'b', 'c', 'a', 'b', 'b', 'b', 'c'],
                #  'v2': ['aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb'],
                #  'v3': ['A', 'B', 'B', 'B', 'B', 'A', 'B', 'B', 'B', 'B']}
            # The above example could have been made shorter. But that's a limitation of the function.

            dict_of_dicts = {
                "v1": {"a": 0.2, "b": 0.6, "c": 0.2},
                "v2": {"aa": 0.5, "bb": 0.5},
                "v3": {"A": 0.2, "B": 0.8},
                "v4": {"A": 0.1, "B": 0.9},
            }
            realize_dicts_of_proportions(dict_of_dicts)
                # {'v1': ['a', 'b', 'b', 'b', 'c', 'a', 'b', 'b', 'b', 'c'],
                #  'v2': ['aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb', 'aa', 'bb'],
                #  'v3': ['A', 'B', 'B', 'B', 'B', 'A', 'B', 'B', 'B', 'B'],
                #  'v4': ['A', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B']}
    """
    # Generate proportional arrays for each dictionary
    arrays = {k: _proportional_array_from_dict(v) for k, v in dict_of_dicts.items()}

    # Find the LCM over the lengths of all the arrays
    lcm_length = _find_lcm_of_array_lengths(arrays)

    # Extend each array to have the same LCM length while maintaining proportions
    result = {}
    for k, arr in arrays.items():
        factor = lcm_length // len(arr)
        extended_arr = arr * factor
        result[k] = extended_arr

    return result


def prepare_marginal_dist_for_raking(
    dict_of_dicts: Dict[str, Dict[str, float]],
) -> pd.DataFrame:
    """
    Realizes a nested dictionary of proportions into a DataFrame.

    Args:
        dict_of_dicts: A nested dictionary where the outer keys are column names
                   and the inner dictionaries have keys as category labels
                   and values as their proportions (float).

    Returns:
        A DataFrame with columns specified by the outer keys of the input dictionary
        and rows containing the category labels according to their proportions.
        An additional "id" column is added with integer values as row identifiers.

    Examples:
    .. code-block:: python

        from balance.weighting_methods.rake import prepare_marginal_dist_for_raking
        df = prepare_marginal_dist_for_raking(
            {"A": {"a": 0.5, "b": 0.5}, "B": {"x": 0.2, "y": 0.8}}
        )
        df.columns.tolist()
        # ['A', 'B', 'id']
    """
    target_dict_from_marginals = _realize_dicts_of_proportions(dict_of_dicts)
    target_df_from_marginals = pd.DataFrame.from_dict(target_dict_from_marginals)
    # Add an id column:
    target_df_from_marginals["id"] = range(target_df_from_marginals.shape[0])

    return target_df_from_marginals
