"""  Created on 03/06/2022::
------------- simul_commonprob.py -------------

**Authors**: L. Mingarelli
"""

import numpy as np
from scipy.stats import norm, mvn
from tqdm import tqdm

def simul_commonprob(margprob, corr=0, method="integrate", n1=10**5, n2=10, pbar=False):
    """
    Compute common probabilities of binary random variates
    generated by thresholding normal variates at 0.

    The output of this function is used by rmvbin. For all combinations of marginprob[i], marginprob[j]
    and corr[k], the probability that both components of a normal random variable with mean qnorm(marginprob[i,j])
    and correlation corr[k] are larger than zero is computed.
    The probabilities are either computed by numerical integration of the multivariate normal density,
    or by Monte Carlo simulation.
    For normal usage of rmvbin it is not necessary to use this function, one simulation result is provided
    as variable SimulVals in this package and loaded by default.

    Args:
        margprob: vector of marginal probabilities.
        corr: vector of correlation values for normal distribution.
        method: either "integrate" or "monte carlo".
        n1: number of normal variates if method is "monte carlo".
        n2: number of repetitions if method is "monte carlo".
        pbar (bool): Whether to display progress bar

    Returns:
        an array of dimension (len(margprob), len(margprob), len(corr)).
    """
    lm = len(margprob)
    lr = len(corr)

    z = np.zeros((lm, lm, lr))

    method = ["integrate", "monte carlo"].index(method)
    if method == -1:
        raise ValueError("invalid method")

    for k in range(lr):
        sigma = np.array([[1, corr[k]], [corr[k], 1]])

        if pbar: range_prog = tqdm(range(lm), desc=f'{corr[k]}')
        else:    range_prog = range(lm)
        for m in range_prog:
            q1 = norm.ppf(margprob[m])
            for n in range(m, lm):
                corr[k] = round(corr[k], 12)
                margprob[m] = round(margprob[m], 12)
                margprob[n] = round(margprob[n], 12)
                if not pbar: print(corr[k], margprob[m], margprob[n], ": ", end="")
                q2 = norm.ppf(margprob[n])

                if corr[k] == -1:
                    z[m, n, k] = max(margprob[m] + margprob[n] - 1, 0)
                elif corr[k] == 0:
                    z[m, n, k] = margprob[m] * margprob[n]
                elif corr[k] == 1:
                    z[m, n, k] = min(margprob[m], margprob[n])
                elif margprob[m] * margprob[n] == 0:
                    z[m, n, k] = 0
                elif margprob[m] == 1:
                    z[m, n, k] = margprob[n]
                elif margprob[n] == 1:
                    z[m, n, k] = margprob[m]
                elif method == 0: # Integrate
                    a, _ = mvn.mvnun([0, 0], [np.inf, np.inf], [q1, q2], sigma)
                    if not np.isfinite(a):
                        z[m, n, k] = np.nan
                    else:
                        z[m, n, k] = a
                else:               # Monte Carlo
                    x1 = np.random.multivariate_normal([q1, q2], sigma, size=(n2, n1))
                    x2 = np.mean((x1[:, :, 0] > 0) & (x1[:, :, 1] > 0), axis=1)
                    z[m, n, k] = np.mean(x2)
                if not pbar: print("done")
                z[n, m, k] = z[m, n, k]

    D = {}
    for j in range(lm):
        for k in range(j, lm):
            mj = round(margprob[j], 10)
            mk = round(margprob[k], 10)
            D[(mj, mk)] = D[(mj, mk)] = np.vstack((corr, z[j, k, :]))


    return D



