import numpy as np
import numpy.linalg as LA
from sympy import re, im, I, Symbol, sqrt
from sympy.functions.elementary.miscellaneous import cbrt
import argparse

### 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')
    
    p = argparse.ArgumentParser()
    p.add_argument('-p', type=str)
    args = p.parse_args()
    
    prior = args.p #Gaussian or Uniform
    N = 2000
    M = 4000
    a = 1/2
    
    ### Constructing Y
    
    if prior == "Gaussian":
        Y = np.random.randn(N,M)
        Y = Y/np.sqrt(N)
        U_y, S_y, Vh_Y = LA.svd(Y)
        Y =  np.hstack((np.diag(S_y),np.zeros((N,M-N))))

    elif prior == "Uniform":
        S_y = 2 * np.random.rand(N) + 1
        Y =  np.hstack((np.diag(S_y),np.zeros((N,M-N))))
    

    
    c = 3
    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)

    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)
    entry6 = np.zeros(N,dtype=complex)
    
    entry1_th = np.zeros(N,dtype=complex)
    entry2_th = np.zeros(N)
    entry3_th = np.zeros(N,dtype=complex)
    entry4_th = np.zeros(N)
    entry5_th = np.zeros(N,dtype=complex)
    entry6_th = np.zeros(N)
    
    
    for i in range(N):
        
        if i%100 == 0:
            with open('start.txt', 'a') as f:
                f.write(prior +' : ' + str(i)+'\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[0,1]
        
        entry3[i] = Res[N,N]
        entry4[i] = Res[N,N+1]
        
        entry5[i] = Res[0,N]
        entry6[i] = Res[0,N+1]
        
        
        gS_eval = gS.subs(z,zz).evalf()
        
        q1 = gS_eval
        q2 = a*gS_eval + (1-a)/zz
        
        q4 = -3 *c +   ( 3**(2/3) * (a * (-4+c**2) * zz - 2*a* (gS_eval**2)* zz+ 2* gS_eval *(-1+a+a* zz**2)) )/ \
        ( a*zz * cbrt( (9 * c * gS_eval * (-1+a * (1-gS_eval* zz+zz**2)))/(a * zz) \
                     + (1/3)*sqrt( (729 * (c**2) * (gS_eval**2) * (-1+a * (1-gS_eval* zz+zz**2))**2)/( (a**2) * zz**2) \
                               +(-3 * c**2+6 * (2+gS_eval**2-(gS_eval* (-1+a+a* zz**2))/(a * zz)))**3  ) )) \
+ cbrt( (27 * c * gS_eval * (-1+a * (1-gS_eval* zz+zz**2)))/(a * zz) \
                     + sqrt( (729 * (c**2) * (gS_eval**2) * (-1+a * (1-gS_eval* zz+zz**2))**2)/( (a**2) * zz**2) \
                               +(-3 * c**2+6 * (2+gS_eval**2-(gS_eval* (-1+a+a* zz**2))/(a * zz)))**3  ) )
        q4 = q4/6
        
        q3 = ( (zz - (1/a)*q2)/(2*q4 +c) )*q4
        
        
        
        beta1 = (1/a)*q2 + q3
        beta2 = q1
        beta3 = q1
        beta4 = q4 + c
        
        Z1 = (zz - beta1)*(zz-beta2)
        Z2 = beta4**2 + beta3*(zz-beta1)
        
        
        entry1_th[i] = ( (zz - beta2 - beta3 * S_y[0]**2)/(Z1 - Z2 * S_y[0]**2) ).evalf()
        entry2_th[i] = 0
        
        entry3_th[i] =  ( (zz - beta1 )/(Z1 - Z2 * S_y[0]**2) ).evalf()
        entry4_th[i] = 0
        
        entry5_th[i] = ( (beta4 * S_y[0] )/(Z1 - Z2 * S_y[0]**2) ).evalf()
        entry6_th[i] = 0
                
        
    
    filename = 'Y_'+prior+'_entry1.npy'
    np.save( filename, entry1)
    filename = 'Y_'+prior+'_entry1_th.npy'
    np.save( filename, entry1_th)
    
    filename = 'Y_'+prior+'_entry2.npy'
    np.save( filename, entry2)
    filename = 'Y_'+prior+'_entry2_th.npy'
    np.save( filename, entry2_th)
    
    filename = 'Y_'+prior+'_entry3.npy'
    np.save( filename, entry3)
    filename = 'Y_'+prior+'_entry3_th.npy'
    np.save( filename, entry3_th)
    
    filename = 'Y_'+prior+'_entry4.npy'
    np.save( filename, entry4)
    filename = 'Y_'+prior+'_entry4_th.npy'
    np.save( filename, entry4_th)
    
    filename = 'Y_'+prior+'_entry5.npy'
    np.save( filename, entry5)
    filename = 'Y_'+prior+'_entry5_th.npy'
    np.save( filename, entry5_th)
    
    filename = 'Y_'+prior+'_entry6.npy'
    np.save( filename, entry6)
    filename = 'Y_'+prior+'_entry6_th.npy'
    np.save( filename, entry6_th)
    
    filename = 'Y_'+prior+'_Observation_SVs.npy'
    np.save( filename, S_y)
    
    
if __name__ == "__main__":
    main()
    
