from ...forall import *
#######################################################################################################################
# Специальные дискретные случайные величины
#######################################################################################################################
def SDRV_1(na,nb,pa,pb):
    """Банк выдал кредиты двум группам заемщиков: `na` заемщиков в первой группе и `nb` – во второй.<br> 
    Известно, что заемщики из первой группы возвращают кредит с вероятностью `pa`,<br>
    а заемщики из второй группы – с вероятностью `pb`.<br>
    Пусть X – суммарное количество возвращенных кредитов для обеих групп.<br>
    Предполагая независимость заемщиков, найдите:<br>
    - стандартное отклонение X; 
    - асимметрию X. 

    Args:
        na (numerical): Количество заемщиков в 1-ой группе
        nb (numerical): Количество заемщиков в 2-ой группе
        pa (numerical): Вероятность возвращения кредита в 1-ой группе
        pb (numerical): Вероятность возвращения кредита в 2-ой группе

    ## Prints
        `answer` каждое значение по очереди.<br>C запятой вместо точки и сокращением до соответствующего количества знаков после запятой

    Returns:
        `answer` (tuple): Соответствующие величины
    """
    qa = 1-pa
    qb = 1-pb

    sigmaX = (na*pa*qa + nb*pb*qb)**0.5
    

    AsX = (na*pa*qa*(qa-pa) + nb*pb*qb*(qb-pb))/(sigmaX**3)
    
    answer = (sigmaX,AsX)
    
    print('std_X = ' + rrstr(sigmaX,1))
    print('As_X  = ' + rrstr(AsX,3))
    return answer
#######################################################################################################################
def SDRV_2(sigmaX,sigmaY,sigmaZ,s):
    """Независимые пуассоновские случайные величины X,Y,Z
    имеют следующие стандартные отклонения:
    - σX=`sigmaX`
    - σY=`sigmaY`
    - σZ=`sigmaZ`<br>
    Пусть S=X+Y+Z<br>
    Найдите:<br>
    - вероятность P(S=`s`)
    - наиболее вероятное значение суммы S
    - стандартное отклонение σS
    - асимметрию As(S)
    - эксцесс Ex(S)
    Args:
        sigmaX (numerical): Стандартное отклонение X
        sigmaY (numerical): Стандартное отклонение Y
        sigmaZ (numerical): Стандартное отклонение Z
        s (numerical): Значение S, вероятность которого нужно найти

    ## Prints
       `answer` каждое значение по очереди.<br>C запятой вместо точки и сокращением до соответствующего количества знаков после запятой

    Returns:
        `answer` (tuple): Соответствующие величины
    """
    import scipy.stats

    S =  scipy.stats.poisson(sigmaX**2 + sigmaY**2 + sigmaZ**2)
    
    AsS = (sigmaX**2 + sigmaY**2 + sigmaZ**2)/((S.var()**0.5)**3)
    
    numerator = (sigmaX**2 + sigmaY**2 + sigmaZ**2) + 3*(sigmaX**4 + sigmaY**4 + sigmaZ**4) + 6*(sigmaX**2 * sigmaY**2 + sigmaY**2 * sigmaZ**2 + sigmaZ**2 * sigmaX**2)
    denominator = (S.var()**0.5)**4
    ExS = (numerator/denominator) - 3
    
    answer = (S.pmf(s),S.median(),S.std(),AsS,ExS)
    
    print(f'P(S = {s}) = ' + rrstr(answer[0],3))
    print('Наиболее вероятное S = ' + str(answer[1]))
    print('Cтандартное отклонение σS = ' + rrstr(answer[2],3))
    print('асимметрия As(S) = ' + rrstr(answer[3],3))
    print(rrstr(answer[4],3),'эксцесс Ex(S)')

    return answer
#######################################################################################################################
def SDRV_3(N,n,G):
    """Монеты в количестве `N`<br>
    штук подбрасываются до тех пор, пока `n` раз<br>
    не выпадет `G` гербов.<br>
    Пусть X – число бросков до первого появления `G` гербов,<br>
    а Y – число бросков до последнего появления `G`гербов<br>
    (Y = общее число бросков). <br>
    Найдите:<br>
    - математическое ожидание X
    - стандартное отклонение X
    - коэффициент корреляции между X и Y
    - математическое ожидание XY

    Args:
        N (numerical): Общее количество монет
        n (numerical): Пока столько раз не выпадет `G` гербов
        G (numerical): Количество гербов

    ## Prints
       `answer` каждое значение по очереди.<br>C запятой вместо точки и сокращением до соответствующего количества знаков после запятой

    Returns:
        `answer` (tuple): Соответствующие величины
    """
    import scipy.stats
    import math
    def C(n,k):
        return math.comb(n,k)

    M = scipy.stats.geom(1/2 ** N * C(N,G))
    
    answer = (M.mean(),(M.var())**0.5,M.var()/(M.std() * (n*M.var())**0.5),M.var()+n*M.mean()**2)


    print(rrstr(answer[0],3),'Математическое ожидание X')
    print(rrstr(answer[1],3),'Стандартное отклонение X')
    print(rrstr(answer[2],3),'Коэффициент корреляции X и Y')
    print(rrstr(answer[3],1),'Математическое ожидание XY')

    return answer
#######################################################################################################################
def SDRV_4(N,r,b,n):
    """Корзина содержит `N` шаров,
    среди которых `r` – красных и `b` – синих.
    Из корзины, случайным образом, без возвращения извлекаются `n` шаров.
    Пусть X и Y обозначают количество красных и синих шаров среди извлеченных, соответственно. Найдите ковариацию Cov(X,Y).

    Args:
        N (numerical): Всего шаров в корзине
        r (numerical): Красных шаров в корзине
        b (numerical): Синих шаров в корзине
        n (numerical): Извлекается шаров без возвращения
    
    ## Prints
       `answer` C запятой вместо точки и сокращением до соответствующего количества знаков после запятой

    Returns:
        `answer` (tuple): Соответствующая величина
    """
    import math
    import scipy.stats

    Ex = r/N * n
    Ey = b/N * n
    Z = N - r - b

    def C(n,k):
        return math.comb(n,k)


    dXY = dict()
    for i in range(r+1):
        for j in range(b+1):
            for k in range(Z+1):
                if (i+j+k)==n:
                    p = C(r,i)*C(b,j)*C(Z,k)/C(N,n)
                    dXY.setdefault(i*j,0)
                    dXY[i*j]+=p

    val = [i for i in dXY.keys()]
    chn = [i for i in dXY.values()]
    
    Exy = scipy.stats.rv_discrete(values=(val,chn)).mean()
    Cov = Exy - Ex*Ey
    
    answer = tuple([Cov])
    
    print(rrstr(answer[0],2))
    
    return answer
#######################################################################################################################
SDRV = [SDRV_1,SDRV_2,SDRV_3,SDRV_4]