# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from scipy.stats import binom, norm, rankdata
from scipy.special import comb
import pandas as pd
import numpy as np


class DevelopmentCorrelation:
    """
    Mack (1997) test for correlations between subsequent development
    factors. Results should be within confidence interval range
    otherwise too much correlation

    Parameters
    ----------
    triangle: Triangle
        Triangle on which to estimate correlation between subsequent development
        factors.
    p_critical: float (default=0.10)
        Value between 0 and 1 representing the confidence level for the test. A
        value of 0.1 implies 90% confidence.

    Attributes
    ----------
    t_critical: DataFrame
        Boolean value for whether correlation is too high based on ``p_critical``
        confidence level.
    t_expectation: DataFrame
        Values representing the Spearman rank correlation
    t_variance: float
        Variance measure of Spearman rank correlation
    range: tuple
        Range within which ``t_expectation`` must fall for independence assumption
        to be significant.
    """

    def __init__(self, triangle, p_critical=0.5):
        self.p_critical = p_critical
        if triangle.array_backend != "numpy":
            triangle = triangle.set_backend("numpy")
        xp = triangle.get_array_module()

        m1 = triangle.link_ratio
        m1_val = xp.apply_along_axis(rankdata, 2, m1.values) * (m1.values * 0 + 1)
        m2 = triangle[triangle.valuation < triangle.valuation_date].link_ratio
        m2.values = xp.apply_along_axis(rankdata, 2, m2.values) * (m2.values * 0 + 1)
        m1 = m2.copy()
        m1.values = m1_val[..., : m2.shape[2], 1:]
        numerator = ((m1 - m2) ** 2).sum("origin")
        numerator.values = numerator.values[..., :-1]
        numerator.ddims = numerator.ddims[:-1]
        I = triangle.shape[3]
        k = xp.array(range(2, 2 + numerator.shape[3]))
        denominator = ((I - k) ** 3 - I + k)[None, None, None]
        self.t = 1 - 6 * xp.nan_to_num(numerator.values) / denominator
        weight = (I - k - 1)[None, None, None]
        t_expectation = (
            xp.sum(xp.nan_to_num(weight * self.t), axis=3) / xp.sum(weight, axis=3)
        )[..., None]
        idx = triangle.index.set_index(triangle.key_labels).index
        self.t_variance = 2 / ((I - 2) * (I - 3))
        self.t = pd.DataFrame(self.t[..., 0, 0], columns=triangle.vdims, index=idx)
        self.t_expectation = pd.DataFrame(
            t_expectation[..., 0, 0], columns=triangle.vdims, index=idx
        )
        self.range = (
            norm.ppf(0.5 - (1 - p_critical) / 2) * xp.sqrt(self.t_variance),
            norm.ppf(0.5 + (1 - p_critical) / 2) * xp.sqrt(self.t_variance),
        )
        self.t_critical = (self.t_expectation < self.range[0]) | (
            self.t_expectation > self.range[1]
        )


class ValuationCorrelation:
    """
    Mack (1997) test for calendar year effect.A calendar period has impact
    across developments if the probability of the number of small (or large)
    development factors, Z, in that period occurring randomly is less than
    ``p_critical``

    Parameters
    ----------
    triangle: Triangle
        Triangle on which to test whether the calendar effects violate independence
        requirements of the chainladder method.
    p_critical: float (default=0.10)
        Value between 0 and 1 representing the confidence level for the test. 0.1
        implies 90% confidence.
    total: boolean
        Whether to calculate valuation correlation in total across all
        years (True) consistent with Mack 1993 or for each year separately
        (False) consistent with Mack 1997.

    Attributes
    ----------
    z : Triangle or DataFrame
        Z values for each Valuation Period
    z_critical : Triangle or DataFrame
        Boolean value for whether correlation is too high based on ``p_critical``
        confidence level.
    z_expectation : Triangle or DataFrame
        The expected value of Z.
    z_variance : Triangle or DataFrame
        The variance value of Z.
    """

    def __init__(self, triangle, p_critical=0.1, total=True):
        def pZlower(z, n, p=0.5):
            return min(1, 2 * binom.cdf(z, n, p))

        self.p_critical = p_critical
        self.total = total
        triangle = triangle.set_backend("numpy")
        xp = triangle.get_array_module()
        lr = triangle.link_ratio
        m1 = xp.apply_along_axis(rankdata, 2, lr.values) * (lr.values * 0 + 1)
        med = xp.nanmedian(m1, axis=2, keepdims=True)
        m1large = (xp.nan_to_num(m1) > med) + (lr.values * 0)
        m1small = (xp.nan_to_num(m1) < med) + (lr.values * 0)
        m2large = triangle.link_ratio
        m2large.values = m1large
        m2small = triangle.link_ratio
        m2small.values = m1small
        S = xp.nan_to_num(m2small.dev_to_val().sum(axis=2).set_backend('numpy').values)
        L = xp.nan_to_num(m2large.dev_to_val().sum(axis=2).set_backend('numpy').values)
        z = xp.minimum(L, S)
        n = L + S
        m = xp.floor((n - 1) / 2)
        c = comb(n - 1, m)
        EZ = (n / 2) - c * n / (2 ** n)
        VarZ = n * (n - 1) / 4 - c * n * (n - 1) / (2 ** n) + EZ - EZ ** 2
        if not self.total:
            T = []
            for i in range(0, xp.max(m1large.shape[2:]) + 1):
                T.append(
                    [
                        pZlower(i, j, 0.5)
                        for j in range(0, xp.max(m1large.shape[2:]) + 1)
                    ]
                )
            T = np.array(T)
            z_idx, n_idx = z.astype(int), n.astype(int)
            self.probs = T[z_idx, n_idx]
            z_critical = triangle[triangle.valuation > triangle.valuation.min()]
            z_critical = z_critical.dev_to_val().dropna().sum("origin") * 0
            z_critical.values = np.array(self.probs) < p_critical
            z_critical.odims = triangle.odims[0:1]
            self.z_critical = z_critical
            self.z = self.z_critical.copy()
            self.z.values = z
            self.z_expectation = self.z_critical.copy()
            self.z_expectation.values = EZ
            self.z_variance = self.z_critical.copy()
            self.z_variance.values = VarZ
        else:
            ci2 = norm.ppf(0.5 - (1 - p_critical) / 2) * xp.sqrt(xp.sum(VarZ, axis=-1))
            self.range = (xp.sum(VarZ, axis=-1) + ci2, xp.sum(VarZ, axis=-1) - ci2)
            idx = triangle.index.set_index(triangle.key_labels).index
            self.z_critical = pd.DataFrame(
                (
                    (self.range[0] > VarZ.sum(axis=-1))
                    | (VarZ.sum(axis=-1) > self.range[1])
                )[..., 0],
                columns=triangle.vdims,
                index=idx,
            )
            self.z = pd.DataFrame(
                z.sum(axis=-1)[..., 0], columns=triangle.vdims, index=idx
            )
            self.z_expectation = pd.DataFrame(
                EZ.sum(axis=-1)[..., 0], columns=triangle.vdims, index=idx
            )
            self.z_variance = pd.DataFrame(
                VarZ.sum(axis=-1)[..., 0], columns=triangle.vdims, index=idx
            )
