
import numpy as np #for all numerical computations
import pandas as pd #similar to dataframe in R
import scipy as sp #scientific computations (includes stats)
import matplotlib.pyplot as plt #plotting 
import scipy.io as sio
import cvxpy as cp # convex optimization
import random
import sys
import time
import gc
import math

from sklearn.utils.extmath import randomized_svd

def randomized_SPCA(data,s):
    
    A=np.matmul(np.transpose(data),data)
    Z_F2= np.linalg.norm(data, 'fro')**2 # Calculate the Frebenius norm
    p_tilde= np.linalg.norm(data,axis=0)**2/Z_F2
    p=np.minimum(np.array(p_tilde)*(s),np.ones(data.shape[1],dtype=int)) #min(sp,1)
        
    idx=np.array([],dtype="int32")

    while idx.shape[0]==0:

        sample=np.random.binomial(1,p,len(p))

        S= np.diag(np.multiply(sample,np.reciprocal(np.sqrt(p)))) # sampling and rescaling matrix

        # top right singular vector of ZS
        
        U, Sig, VT =randomized_svd(np.matmul(data,S), n_components=1, random_state=10)
        y_rspca=np.matmul(S,VT[0]) # VT[0]: top right singular vector of ZS

        # SVD normalization
        dummy=np.where(sample!=0)[0]
        idx=np.append(idx,dummy)

    A_red=A[np.ix_(idx, idx)].copy()
    
    UA, SigA, VTA =randomized_svd(A_red, n_components=1, random_state=10)

    y_rspca_norm = np.zeros(len(p))
    y_rspca_norm[idx]=VTA[0].copy()  ### padding other elements with zeros

    f_rspca=np.sum(np.square(np.matmul(data,y_rspca_norm)))/np.linalg.norm(data, 2)**2


    print("rspca done.")
    return y_rspca_norm, f_rspca