"""
Functions for calculating the statistical significant differences
between two dependent or independent correlation coefficients.
The Fisher and Steiger method is adopted from the R package
http://personality-project.org/r/html/paired.r.html
and is described in detail in the book 'Statistical Methods for Psychology'
The Zou method is adopted from
http://seriousstats.wordpress.com/2012/02/05/comparing-correlations/
Credit goes to the authors of above mentioned packages!
Author: Philipp Singer (www.philippsinger.info)

------------------------------------------------------------
README.md:
CorrelationStats
This Python script enables you to compute statistical significance tests on
both dependent and independent correlation coefficients. For each case two methods to
choose from are available.

For details, please refer to: http://www.philippsinger.info/?p=347
#copied from on 4/24/2015 from
https://github.com/psinger/CorrelationStats/blob/master/corrstats.py
"""

__author__ = "psinger"
from math import atanh
from typing import Optional, Tuple

import numpy as np
from scipy.stats import norm, t


def rz_ci(r: float, n: int, conf_level: float = 0.95) -> np.ndarray:
    zr_se = pow(1 / (n - 3), 0.5)
    moe = norm.ppf(1 - (1 - conf_level) / float(2)) * zr_se
    zu = atanh(r) + moe
    zl = atanh(r) - moe
    return np.tanh((zl, zu))


def rho_rxy_rxz(rxy: float, rxz: float, ryz: float) -> float:
    num = (ryz - 1 / 2.0 * rxy * rxz) * (
        1 - pow(rxy, 2) - pow(rxz, 2) - pow(ryz, 2)
    ) + pow(ryz, 3)
    den = (1 - pow(rxy, 2)) * (1 - pow(rxz, 2))
    return num / float(den)


def dependent_corr(
    xy: float,
    xz: float,
    yz: float,
    n: int,
    twotailed: bool = True,
    conf_level: float = 0.95,
    method: str = "steiger",
) -> Tuple[float, float]:
    """
    Calculates the statistic significance between two dependent correlation coefficients
    :param xy: correlation coefficient between x and y
    :param xz: correlation coefficient between x and z
    :param yz: correlation coefficient between y and z
    :param n: number of elements in x, y and z
    :param twotailed: whether to calculate a one or two tailed test, only works for 'steiger' method
    :param conf_level: confidence level, only works for 'zou' method
    :param method: defines the method uses, 'steiger' or 'zou'
    @return: t and p-val
    """
    if method != "steiger" and "zou":
        raise Exception("Wrong method!")

    if method == "steiger":
        d = xy - xz
        determin = 1 - xy * xy - xz * xz - yz * yz + 2 * xy * xz * yz
        av = (xy + xz) / 2
        cube = (1 - yz) * (1 - yz) * (1 - yz)

        t2 = d * np.sqrt(
            (n - 1) * (1 + yz) / ((2 * (n - 1) / (n - 3)) * determin + av * av * cube)
        )
        p = 1 - t.cdf(abs(t2), n - 2)

        if twotailed:
            p *= 2

        output1 = t2
        output2 = p
    elif method == "zou":
        L1 = rz_ci(xy, n, conf_level=conf_level)[0]
        U1 = rz_ci(xy, n, conf_level=conf_level)[1]
        L2 = rz_ci(xz, n, conf_level=conf_level)[0]
        U2 = rz_ci(xz, n, conf_level=conf_level)[1]
        rho_r12_r13 = rho_rxy_rxz(xy, xz, yz)
        lower = (
            xy
            - xz
            - pow(
                (
                    pow((xy - L1), 2)
                    + pow((U2 - xz), 2)
                    - 2 * rho_r12_r13 * (xy - L1) * (U2 - xz)
                ),
                0.5,
            )
        )
        upper = (
            xy
            - xz
            + pow(
                (
                    pow((U1 - xy), 2)
                    + pow((xz - L2), 2)
                    - 2 * rho_r12_r13 * (U1 - xy) * (xz - L2)
                ),
                0.5,
            )
        )
        output1 = lower
        output2 = upper

    return (output1, output2)


def independent_corr(
    xy: float,
    ab: float,
    n: int,
    n2: Optional[int] = None,
    twotailed: bool = True,
    conf_level: float = 0.95,
    method: str = "fisher",
) -> Tuple[float, float]:
    """
    Calculates the statistic significance between two independent correlation coefficients
    :param xy: correlation coefficient between x and y
    :param ab: correlation coefficient between a and b
    :param n: number of elements in xy
    :param n2: number of elements in ab (if distinct from n)
    :param twotailed: whether to calculate a one or two tailed test, only works for 'fisher' method
    :param conf_level: confidence level, only works for 'zou' method
    :param method: defines the method uses, 'fisher' or 'zou'
    return z and p-val
    """
    if n2 is None:
        n2 = n

    if not method == "fisher" and "zou":
        raise Exception("Wrong method!")

    if method == "fisher":
        xy_z = 0.5 * np.log((1 + xy) / (1 - xy))
        xz_z = 0.5 * np.log((1 + ab) / (1 - ab))

        se_diff_r = np.sqrt(1 / (n - 3) + 1 / (n2 - 3))
        diff = xy_z - xz_z
        z = abs(diff / se_diff_r)
        p = 1 - norm.cdf(z)
        if twotailed:
            p *= 2
        output1 = z
        output2 = p
    elif method == "zou":
        L1 = rz_ci(xy, n, conf_level=conf_level)[0]
        U1 = rz_ci(xy, n, conf_level=conf_level)[1]
        L2 = rz_ci(ab, n2, conf_level=conf_level)[0]
        U2 = rz_ci(ab, n2, conf_level=conf_level)[1]
        lower = xy - ab - pow((pow((xy - L1), 2) + pow((U2 - ab), 2)), 0.5)
        upper = xy - ab + pow((pow((U1 - xy), 2) + pow((ab - L2), 2)), 0.5)
        output1 = lower
        output2 = upper

    return output1, output2
