import typing

import numpy as np
import pytest

import rapidstats as rs

np.random.seed(208)

N_ROWS = 1_000
CHOICES = ["A", "B", "C", "D", "E", "F", "G", "H"]
WEIGHTS = [0.0001, 0.2, 0.3, 0.4, 0.5, 6.0, 0.7, 0.8]
_sw = sum(WEIGHTS)
WEIGHTS = [w / _sw for w in WEIGHTS]

REFERENCE = np.random.rand(N_ROWS)
CURRENT = np.random.rand(int(N_ROWS * 0.8))

REFERENCE_CAT = np.random.choice(CHOICES, size=N_ROWS, p=WEIGHTS)
CURRENT_CAT = np.random.choice(CHOICES, size=int(N_ROWS * 0.8), p=WEIGHTS)


@pytest.mark.parametrize("ref", [REFERENCE, REFERENCE_CAT])
def test_identical(ref):
    assert rs.drift.psi(ref, ref.copy()) == pytest.approx(0.0)


def reference_psi(expected, actual, buckettype="bins", buckets=10, axis=0):
    """Calculate the PSI (population stability index) across all variables

    Args:
       expected: numpy matrix of original values
       actual: numpy matrix of new values
       buckettype: type of strategy for creating buckets, bins splits into even splits, quantiles splits into quantile buckets
       buckets: number of quantiles to use in bucketing variables
       axis: axis by which variables are defined, 0 for vertical, 1 for horizontal

    Returns:
       psi_values: ndarray of psi values for each variable

    Author:
       Matthew Burke
       github.com/mwburke
       mwburke.github.io.com
    """

    def psi(expected_array, actual_array, buckets):
        """Calculate the PSI for a single variable

        Args:
           expected_array: numpy array of original values
           actual_array: numpy array of new values, same size as expected
           buckets: number of percentile ranges to bucket the values into

        Returns:
           psi_value: calculated PSI value
        """

        def scale_range(input, min, max):
            input += -(np.min(input))
            input /= np.max(input) / (max - min)
            input += min
            return input

        breakpoints = np.arange(0, buckets + 1) / (buckets) * 100

        if buckettype == "bins":
            breakpoints = scale_range(
                breakpoints, np.min(expected_array), np.max(expected_array)
            )
        elif buckettype == "quantiles":
            breakpoints = np.stack(
                [np.percentile(expected_array, b) for b in breakpoints]
            )

        expected_fractions = np.histogram(expected_array, breakpoints)[0] / len(
            expected_array
        )
        actual_fractions = np.histogram(actual_array, breakpoints)[0] / len(
            actual_array
        )

        def sub_psi(e_perc, a_perc):
            """Calculate the actual PSI value from comparing the values.
            Update the actual value to a very small number if equal to zero
            """
            if a_perc == 0:
                a_perc = 0.0001
            if e_perc == 0:
                e_perc = 0.0001

            value = (e_perc - a_perc) * np.log(e_perc / a_perc)
            return value

        psi_value = sum(
            sub_psi(expected_fractions[i], actual_fractions[i])
            for i in range(0, len(expected_fractions))
        )

        return psi_value

    if len(expected.shape) == 1:
        psi_values = np.empty(len(expected.shape))
    else:
        psi_values = np.empty(expected.shape[1 - axis])

    for i in range(0, len(psi_values)):
        if len(psi_values) == 1:
            psi_values = psi(expected, actual, buckets)
        elif axis == 0:
            psi_values[i] = psi(expected[:, i], actual[:, i], buckets)
        elif axis == 1:
            psi_values[i] = psi(expected[i, :], actual[i, :], buckets)

    return psi_values


@pytest.mark.parametrize("tup", [(REFERENCE, CURRENT)])
@pytest.mark.parametrize("bin_count", [5, 10, 20])
def test_psi(tup, bin_count):
    ref, cur = tup

    x = reference_psi(ref, cur, buckettype="bins", buckets=bin_count)
    y = rs.drift.psi(ref, cur, bin_count=bin_count)

    assert pytest.approx(x) == y


@pytest.mark.parametrize("tup", [(REFERENCE, CURRENT), (REFERENCE_CAT, CURRENT_CAT)])
@pytest.mark.parametrize("bins", [None, [0.1, 0.5, 0.9]])
@pytest.mark.parametrize(
    "bin_count", list(typing.get_args(rs.drift.BinMethod)) + [1, 5, 10]
)
@pytest.mark.parametrize("include_nulls", [True, False])
@pytest.mark.parametrize("epsilon", [None, 0.01])
def test_runs(tup, bins, bin_count, include_nulls, epsilon):
    ref, cur = tup

    rs.drift.psi(
        ref,
        cur,
        bins=bins,
        bin_count=bin_count,
        include_nulls=include_nulls,
        epsilon=epsilon,
    )
