import numpy as np
import numpy.linalg as LA
from sympy import re, im, I, Symbol, sqrt

### returns numerical Cauchy tranform of symmetrized distribution of the input set of singular values
def g_CauchyK_num(S):
    z = Symbol('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 = Symbol('z')
    
    N = 2000
    M = 3000
    prior = "Wigner" ## or Wishart
    
    ### Constructing X
    
    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 == "Wishart":
        X = np.random.normal(0, 1, (N,2*N))
        X = X/np.sqrt(N)
        X = X @ np.transpose(X)
    
    X_e, U_x = LA.eigh(X)
    X = np.diag(X_e)
    

    Y = np.random.randn(N,M)
    Y = Y/np.sqrt(N)
    
    W = np.random.randn(N,M)
    W = W/np.sqrt(N)

    S = X @ Y + W
    U_s, S_s , Vh_s = LA.svd(S)
    
    gS = g_CauchyK_num(S_s)
    
    A11=np.zeros((N,N));
    A12=S;
    A21=np.transpose(S);
    A22=np.zeros((M,M));

    S_C = np.bmat([[A11, A12], [A21, A22]])
    
    entry1 = np.zeros(N,dtype=complex)
    entry2 = np.zeros(N,dtype=complex)
    entry3 = np.zeros(N,dtype=complex)
    entry4 = np.zeros(N,dtype=complex)
    entry5 = np.zeros(N,dtype=complex)
    
    entry1_th = np.zeros(N,dtype=complex)
    entry2_th = np.zeros(N,dtype=complex)
    entry3_th = np.zeros(N)
    entry4_th = np.zeros(N)
    entry5_th = np.zeros(N)
    
    
    for i in range(N):
        zz = S_s[i] - 0.01581J
        
        G = zz*np.eye(N+M) - S_C
        Res = LA.inv(G)

        entry1[i] = Res[0,0]
        entry2[i] = Res[-1,-1]
        
        entry3[i] = Res[-2,-1]
        entry4[i] = Res[0,-1]
        entry5[i] = Res[0,1]
        
        
        gS_eval = gS.subs(z,S_s[i]).evalf()
        
        zeta1 = gS_eval + (1/2)*(1/zz) ### alpha = 2/3
        p2 = (2/3)*gS_eval + (1/3)*(1/zz)
        zeta3 = zeta1
        
        
        entry1_th[i] = (  (1/zeta3) * 1/( zz/zeta3 - 1 - X_e[0]**2) ).evalf()
        entry2_th[i] = p2
        entry3_th[i] = 0
        entry4_th[i] = 0
        entry5_th[i] = 0
                
        
    
    filename = 'X_'+prior+'_entry1.npy'
    np.save( filename, entry1)
    filename = 'X_'+prior+'_entry1_th.npy'
    np.save( filename, entry1_th)
    
    filename = 'X_'+prior+'_entry2.npy'
    np.save( filename, entry2)
    filename = 'X_'+prior+'_entry2_th.npy'
    np.save( filename, entry2_th)
    
    filename = 'X_'+prior+'_entry3.npy'
    np.save( filename, entry3)
    filename = 'X_'+prior+'_entry3_th.npy'
    np.save( filename, entry3_th)
    
    filename = 'X_'+prior+'_entry4.npy'
    np.save( filename, entry4)
    filename = 'X_'+prior+'_entry4_th.npy'
    np.save( filename, entry4_th)
    
    filename = 'X_'+prior+'_entry5.npy'
    np.save( filename, entry5)
    filename = 'X_'+prior+'_entry5_th.npy'
    np.save( filename, entry5_th)
    
    filename = 'X_'+prior+'_Observation_SVs.npy'
    np.save( filename, S_s)
    
    
if __name__ == "__main__":
    main()
    
