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



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('-p', type=str)
    
    args = p.parse_args()

    prior = args.p

    N = 4000
    M = 8000
    
    c = 3

    if prior=='Gaussian':
        Y = np.random.randn(N,M)
        Y = Y/np.sqrt(N)
        
    elif prior=='Uniform':
        G = np.random.randn(N,M)
        G = G/np.sqrt(N)
        U_y, t, Vh_y = LA.svd(G)
        sv_y = 2 * np.random.rand(N) + 1

        S_y =  np.hstack((np.diag(sv_y),np.zeros((N,M-N))))

        Y =  U_y @ S_y @ Vh_y
    
    X = np.triu(np.random.normal(0, 1, (N,N)),1)
    X = X + np.transpose(X) + np.diag(np.random.normal(loc=0, scale=np.sqrt(2), size=(N)))
    X = X/np.sqrt(N)
    X = X + c*np.eye(N)
    
    W = np.random.randn(N,M)
    W = W/np.sqrt(N)



    ### Observation
    S = X @ Y + W
    
    ### SVD
    U_s, S_s , Vh_s = LA.svd(S)

    gS = g_CauchyK_num(S_s)


    
    with open('G_'+prior+'.txt','wb') as f:
        pickle.dump(gS,f)
#
if __name__ == "__main__":
    main()
    
