import numpy as np
import numpy.linalg as LA
from sympy import re, im, I, E, symbols, sqrt
import argparse



def g_CauchyK_num(S):
    z = symbols('z')
    ret = 0
    N = len(S)
    
    for j in range(N):

        ret += 1/(z + S[j] - I*np.sqrt(1/(2*N)) )
        ret += 1/(z - S[j]- I*np.sqrt(1/(2*N)) )
    
    return ret/(2*N)


def main():
    z = symbols('z')
    p = argparse.ArgumentParser()

    p.add_argument('-s', nargs="+", type=float)
    p.add_argument('-p', type=str)
    
    args = p.parse_args()

    s_list = list(args.s)
    prior = args.p

    N = 4000
    M = 8000

    if prior=='Wigner':
        X = np.triu(np.random.normal(0, 1, (N,N)))
        X = X + np.transpose(X) + np.diag(np.random.normal(loc=0, scale=np.sqrt(2), size=(N)))
        X = X/np.sqrt(N)
    elif prior=='SqWishart':
        X = np.random.randn(N,4*N)
        X = X@np.transpose(X)/N
        X_s, U_x = LA.eigh(X)
        X_s = np.sqrt(X_s)
        X = U_x @ np.diag(X_s) @ np.transpose(U_x)
    
    Y = np.random.randn(N,M)
    Y = Y/np.sqrt(N)
    
    W = np.random.randn(N,M)
    W = W/np.sqrt(N)



    ### Observation
    S = X @ Y + W
    
    U_s, S_s , Vh_s = LA.svd(S)
    
    gS = g_CauchyK_num(S_s)
    
    for s in s_list:
    
        gS_eval = gS.subs(z,s).evalf()
        print('probability of ', s, ':' ,im(gS_eval))

#
if __name__ == "__main__":
    main()
    
