import numpy as np
from numpy import linalg as LA
import sys
from mpi4py import MPI
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()


# COMPRESSED SENSING: LINEAR BREGMAN METHOD
# Translated and adapted into python from tinycs
#
# *tinycs* is a minimal compressed sensing (CS) toolkit designed
# to allow MR imaging scientists to design undersampled
# acquisitions and reconstruct the resulting data with CS without
# needing to be a CS expert.
# 
# The Cartesian reconstruction is based on the split Bregman
# code written by Tom Goldstein, originally available here:
# <http://tag7.web.rice.edu/Split_Bregman.html>


def pdf(k,kw,klo,q):
    
    p                           =   (np.abs(k)/kw)**(-q)
    p[np.where(k==0)]           =   0
    p[np.where(np.abs(k)<=kw)]  =   1
    p[np.where(k<klo)]          =   0
    
    return p

def mask_pdf_1d(n,norm,q,pf):
    
    ks      =   np.arange(0,n) - np.ceil(n/2) - 1
    kmax    =   np.floor(n/2)
    npf     =   np.round(pf*n)
    klo     =   ks[n-npf]

    
    for k in range(int(kmax)):
        P   =   pdf(ks,k+1,klo,q)
        if np.sum(P)>=norm:
            break
    
    P       =   np.fft.fftshift(P)
    
    
    #if np.mod(n,2)!=0:
    #    P   =   np.concatenate(([1],P),axis=None)
    
 

    return P

def mask_pdf_2d(dims,norm,q,pf):
    
    nz      =   dims[1]
    ny      =   dims[0]
    yc      =   round(ny/2)
    zc      =   round(nz/2)
    rmax    =   np.sqrt((ny-yc)**2 + (nz-zc)**2)
    [Z,Y]   =   np.meshgrid(np.arange(0,nz),np.arange(0,ny))
    RR      =   np.sqrt( (Y-yc)**2  + (Z-zc)**2)
    Z       =   np.abs(Z - nz/2 - 0.5)
    Y       =   np.abs(Y - ny/2 - 0.5)
    
    
    for rw in range(1,int(rmax)+1):
        P       =   np.ones([ny,nz])/pf
        C       =   np.logical_and( Z <= rw , Y <= rw)
        W       =   np.logical_or( Z > rw , Y > rw)
        P[W]    =   (RR[W]/rw)**(-q)
        if np.sum(P) >= norm:
            break 
        
    
    
    return [P,C]

def GeneratePattern(dim,R):
    
    # 3D CASE
    if np.size(dim)==3:
        
        nro     =   dim[0]
        npe     =   dim[1]
        nacq    =   round(npe/R)
        q       = 1
        pf      = 1
        P       =   mask_pdf_1d(npe, nacq, q, pf)
        while True:
            M   =   np.random.rand(npe)
            M   =   1*(M<=P)
            if np.sum(M)==nacq:
                break
        # remove partial Fourier plane and compensate sampling density
        M = M!=0
        M = np.tile(M,[nro,1]);
        #M   =   M.T

    # 4D CASE
    if np.size(dim)==4:
        
        nro     =   dim[0]
        npe1    =   dim[1]
        npe2    =   dim[2]
        nacq    =   round(npe1*npe2/R)
        q       =   1
        pf      =   1
        [P,C]   =   mask_pdf_2d([npe1,npe2], nacq, q, pf)
        RR      =   np.random.rand(npe1,npe2)
        M       =   (RR <= P)
        nchosen =   np.sum(M)
        if nchosen > nacq:          # Correct for inexact number chosen
            #outerOn                         =   np.logical_and( M , P!=1 )
            outerOn                         =   np.where((M)*(P!=1))
            numToFlip                       =   nchosen-nacq
            idxs                            =   np.random.permutation(outerOn[0].size)
            idxx                            =   outerOn[0][idxs[0:numToFlip]]
            idxy                            =   outerOn[1][idxs[0:numToFlip]]
            M[idxx,idxy]                    =   False
            
        elif nchosen < nacq:
            outerOff                        =   np.where(~M)
            idxs                            =   np.random.permutation(outerOff[0].size)
            numToFlip                       =   nacq - nchosen
            idxx                            =   outerOff[0][idxs[0:numToFlip]]
            idxy                            =   outerOff[1][idxs[0:numToFlip]]
            M[idxx,idxy]                    =   True
            
            
        
    
        M       =   np.rollaxis(np.tile(np.rollaxis(M,1),[nro,1,1]),2)
        M       =   np.fft.ifftshift(M)
        M       =   M.transpose((1,0,2))
    

    return M

def get_norm_factor(MASK,uu):
    UM      =   MASK==1
    return  UM.shape[0]/LA.norm(uu)

def Dxyzt(X):
    
    if np.ndim(X)==3:
        dd0     =   X[:,:,0]
        dd1     =   X[:,:,1]
        DA      =   dd0 - np.vstack((dd0[1::,:],dd0[0,:])) 
        DB      =   dd1 - np.hstack((dd1[:,1::],dd1[:,0:1]))  
        
        return DA + DB

    if np.ndim(X)==4:
        dd0     =   X[:,:,:,0]
        dd1     =   X[:,:,:,1]
        dd2     =   X[:,:,:,2]
        
        DA      =   dd0 - np.vstack((dd0[1::,:,:],dd0[0,:,:][np.newaxis,:,:])) 
        DB      =   dd1 - np.hstack((dd1[:,1::,:],dd1[:,0,:][:,np.newaxis,:]))  
        DC      =   dd2 - np.dstack((dd2[:,:,1::],dd2[:,:,0][:,:,np.newaxis]))
        
        return  DA + DB + DC

def Dxyz(u):
    
    if np.ndim(u)==2:
        
        dx      =   u[:,:]- np.vstack((u[-1,:],u[0:-1,:])) 
        dy      =   u[:,:]- np.hstack((u[:,-1:],u[:,0:-1])) 
        D           =   np.zeros([dx.shape[0],dx.shape[1],2],dtype=complex)
        D[:,:,0]    =   dx
        D[:,:,1]    =   dy
        
        return D

    if np.ndim(u)==3:
        
        dx      =   u[:,:,:]- np.vstack((u[-1,:,:][np.newaxis,:,:],u[0:-1,:,:])) 
        dy      =   u[:,:,:]- np.hstack((u[:,-1,:][:,np.newaxis,:],u[:,0:-1,:])) 
        dz      =   u[:,:,:]- np.dstack((u[:,:,-1][:,:,np.newaxis],u[:,:,0:-1])) 
        
        D           =   np.zeros([dx.shape[0],dx.shape[1],dx.shape[2],3],dtype=complex)
        
        D[:,:,:,0]    =   dx
        D[:,:,:,1]    =   dy
        D[:,:,:,2]    =   dz
        
        return D

def shrink(X,pgam):
    p       =   1
    s       =   np.abs(X)
    tt       =   pgam/(s)**(1-p)
    # t     =   pgam/np.sqrt(s)
    ss      =   s-tt
    ss      =   ss*(ss>0)
    s       =   s + 1*(s<tt)
    ss      =   ss/s
    return ss*X

def CSMETHOD(ITOT,R):
    
    ''' Compressed Function.

    Args:
        ITOT:     	a numpy matrix with the full sampled (3D or 4D) dynamical data
        R:        	the acceleration factor
    '''
    # Method parameters
    ninner          =   5
    nbreg           =   10
    lmbda           =   4
    mu              =   20
    gam             =   1
    
    
    if np.ndim(ITOT)==3:
        [row,col,numt2]     =   ITOT.shape
    elif np.ndim(ITOT)==4:
        [row,col,dep,numt2] =   ITOT.shape
    else:
        raise Exception('Dynamical data is requested')
    
    
    MASK            =   GeneratePattern(ITOT.shape,R)
    CS1             =   np.zeros(ITOT.shape,dtype=complex)


    nit				=	0
    nit_tot			=	(numt2-1)/20
    
    
    if np.ndim(ITOT)==3:
        
        for t in range(numt2):
        
            if rank==0:
                print('{3D COMPRESSED SENSING} t = ',t)
        
            Kdata               =   np.fft.fft2(ITOT[:,:,t])*MASK
        
            data_ndims          =   Kdata.ndim
            mask                =   Kdata!=0         # not perfect, but good enough
            # normalize the data so that standard parameter values work
            norm_factor         =   get_norm_factor(mask, Kdata)
            Kdata               =   Kdata*norm_factor
            # Reserve memory for the auxillary variables
            Kdata0              =   Kdata
            img                 =   np.zeros([row,col],dtype=complex)
            X                   =   np.zeros([row,col, data_ndims])
            B                   =   np.zeros([row,col, data_ndims])
            # Build Kernels
            scale               =   np.sqrt(row*col)
            murf                =   np.fft.ifft2(mu*mask*Kdata)*scale
            uker                =   np.zeros([row,col])
            uker[0,0]           =   4
            uker[0,1]           =   -1  ; uker[1,0]     =   -1
            uker[-1,0]          =   -1  ; uker[0,-1]    =   -1

            uker                =   1/(mu*mask + lmbda*np.fft.fftn(uker) + gam)

            #  Do the reconstruction
            for outer in range(nbreg):
                for inner in range(ninner):
                    # update u
                    rhs      =   murf + lmbda*Dxyzt(X-B) + gam*img
                    img      =   np.fft.ifft2(np.fft.fft2(rhs)*uker)
                    # update x and y
                    A        =   Dxyz(img) + B
                    X        =   shrink(A, 1/lmbda)
                    # update bregman parameters
                    B       =   A - X
                Kdata   = Kdata + Kdata0 - mask*np.fft.fftn(img)/scale
                murf    = np.fft.ifftn(mu*mask*Kdata)*scale
    
            # undo the normalization so that results are scaled properly
            img = img / norm_factor / scale
    
    
            CS1[:,:,t]  =   img
    
    if np.ndim(ITOT)==4:
    
        for t in range(numt2):
            
            
            if rank==0:
                print('[4D CS] R = {re}  t = {te}/{tef}'.format(re=R,te=t,tef=numt2))
                #if np.mod(t,nit_tot)<1:
                #    sys.stdout.write('\r')
                #    # Progress bar
                #    if numt2==3:
                #        sys.stdout.write("{4d-CS} [%-6s] %d%%" % ('=='*nit, 100*t/numt2))
                #    else:
                #        sys.stdout.write("{4d-CS} [%-40s] %d%%" % ('=='*nit, 100*t/numt2))
                #    sys.stdout.flush()
                #    nit		=	nit +1

		

            Kdata_0             =   np.fft.fftn(ITOT[:,:,:,t])
            Kdata               =   Kdata_0*MASK
            data_ndims          =   Kdata.ndim
            mask                =   Kdata!=0         # not perfect, but good enough
            # normalize the data so that standard parameter values work
            norm_factor         =   get_norm_factor(mask, Kdata)
            Kdata               =   Kdata*norm_factor
            # Reserve memory for the auxillary variables
            Kdata0              =   Kdata
            img                 =   np.zeros([row,col,dep],dtype=complex)
            X                   =   np.zeros([row,col,dep, data_ndims])
            B                   =   np.zeros([row,col,dep, data_ndims])
            # Build Kernels
            scale               =   np.sqrt(row*col*dep)
            murf                =   np.fft.ifftn(mu*mask*Kdata)*scale
            uker                =   np.zeros([row,col,dep])
            uker[0,0,0]         =   8
            uker[1,0,0]         =   -1      ;   uker[0,1,0]     =   -1  ;   uker[0,0,1]  =  -1
            uker[-1,0,0]        =   -1      ;   uker[0,-1,0]    =   -1  ;   uker[0,0,-1] =  -1
            uker                =   1/(mu*mask + lmbda*np.fft.fftn(uker) + gam)

            #  Do the reconstruction
            for outer in range(nbreg):
                for inner in range(ninner):
                    # update u
                    rhs      =   murf + lmbda*Dxyzt(X-B) + gam*img
                    img      =   np.fft.ifft2(np.fft.fft2(rhs)*uker)
                    # update x and y
                    A        =   Dxyz(img) + B
                    X        =   shrink(A, 1/lmbda)
                    # update bregman parameters
                    B       =   A - X
                Kdata   = Kdata + Kdata0 - mask*np.fft.fftn(img)/scale
                murf    = np.fft.ifftn(mu*mask*Kdata)*scale
    
            # undo the normalization so that results are scaled properly
            img = img / norm_factor / scale
    
    
            CS1[:,:,:,t]  =   img
    
    
    
    
    return CS1

def CSMETHOD_SENSE(ITOT,R,R_SENSE): 
    ''' Compressed Function.

    Args:
        ITOT:     	a numpy matrix with the full sampled (3D or 4D) dynamical data
        R:        	the acceleration factor
    '''
    # Method parameters
    ninner          =   5
    nbreg           =   10
    lmbda           =   4
    mu              =   20
    gam             =   1

    [row,col,dep,numt2] =   ITOT.shape
    MASK                =   {}
    ITOTCS              =   {}
    MASK[0]             =   GeneratePattern([row,int(np.ceil(col/2)),dep,numt2],R)
    MASK[1]             =   GeneratePattern([row,int(np.ceil(col/2)),dep,numt2],R)
    SenseMAP                        =   {}               
    [SenseMAP[0],SenseMAP[1]]       =   Sensitivity_Map([row,col,dep])

    col                 =   int(np.ceil(col/2))
    ITOTCS[0]           =   np.zeros([row,col,dep,numt2],dtype=complex)
    ITOTCS[1]           =   np.zeros([row,col,dep,numt2],dtype=complex)
        
    for rs in range(R_SENSE):
        for t in range(numt2):
            if rank==0:
                print('[4D CS] R = {re}  t = {te}/{tef}'.format(re=R,te=t,tef=numt2))

            Kdata_0             =   np.fft.fftn(ITOT[:,:,:,t])
            Kdata_0             =   Kdata_0*SenseMAP[rs]
            Kdata_0             =   Kdata_0[:,0::R_SENSE,:]
            
            Kdata               =   Kdata_0*MASK[rs]
            data_ndims          =   Kdata.ndim
            mask                =   Kdata!=0         # not perfect, but good enough
            # normalize the data so that standard parameter values work
            norm_factor         =   get_norm_factor(mask, Kdata)
            Kdata               =   Kdata*norm_factor
            # Reserve memory for the auxillary variables
            Kdata0              =   Kdata
            img                 =   np.zeros([row,col,dep],dtype=complex)
            X                   =   np.zeros([row,col,dep, data_ndims])
            B                   =   np.zeros([row,col,dep, data_ndims])
            # Build Kernels
            scale               =   np.sqrt(row*col*dep)
            murf                =   np.fft.ifftn(mu*mask*Kdata)*scale
            uker                =   np.zeros([row,col,dep])
            uker[0,0,0]         =   8
            uker[1,0,0]         =   -1      ;   uker[0,1,0]     =   -1  ;   uker[0,0,1]  =  -1
            uker[-1,0,0]        =   -1      ;   uker[0,-1,0]    =   -1  ;   uker[0,0,-1] =  -1
            uker                =   1/(mu*mask + lmbda*np.fft.fftn(uker) + gam)

            #  Do the reconstruction
            for outer in range(nbreg):
                for inner in range(ninner):
                    # update u
                    rhs      =   murf + lmbda*Dxyzt(X-B) + gam*img
                    img      =   np.fft.ifft2(np.fft.fft2(rhs)*uker)
                    # update x and y
                    A        =   Dxyz(img) + B
                    X        =   shrink(A, 1/lmbda)
                    # update bregman parameters
                    B       =   A - X
                Kdata   = Kdata + Kdata0 - mask*np.fft.fftn(img)/scale
                murf    = np.fft.ifftn(mu*mask*Kdata)*scale

            # undo the normalization so that results are scaled properly
            img = img / norm_factor / scale

            ITOTCS[rs][:,:,:,t]  =   img

    
    return [ITOTCS[0],ITOTCS[1]]
    
def CSMETHOD_peaksystole(ITOT,R,tstar):
    
    ''' Compressed Function.

    Args:
        ITOT:     	a numpy matrix with the full sampled (3D or 4D) dynamical data
        R:        	the acceleration factor
        tstar:      the time when the flux in the inlet it's maximum
    '''
    # Method parameters
    ninner          =   5
    nbreg           =   10
    lmbda           =   4
    mu              =   20
    gam             =   1
    
    

    [row,col,dep,numt2] =   ITOT.shape

    MASK            =   GeneratePattern(ITOT.shape,R)
    CS1             =   np.zeros([row,col,dep],dtype=complex)  
    
    for t in range(tstar,tstar+1):
        Kdata               =   np.fft.fftn(ITOT[:,:,:,t])*MASK
        data_ndims          =   Kdata.ndim
        mask                =   Kdata!=0         # not perfect, but good enough
        # normalize the data so that standard parameter values work
        norm_factor         =   get_norm_factor(mask, Kdata)
        Kdata               =   Kdata*norm_factor
        # Reserve memory for the auxillary variables
        Kdata0              =   Kdata
        img                 =   np.zeros([row,col,dep],dtype=complex)
        X                   =   np.zeros([row,col,dep, data_ndims])
        B                   =   np.zeros([row,col,dep, data_ndims])
        # Build Kernels
        scale               =   np.sqrt(row*col*dep)
        murf                =   np.fft.ifftn(mu*mask*Kdata)*scale
        uker                =   np.zeros([row,col,dep])
        uker[0,0,0]         =   8
        uker[1,0,0]         =   -1      ;   uker[0,1,0]     =   -1  ;   uker[0,0,1]  =  -1
        uker[-1,0,0]        =   -1      ;   uker[0,-1,0]    =   -1  ;   uker[0,0,-1] =  -1
        uker                =   1/(mu*mask + lmbda*np.fft.fftn(uker) + gam)

        #  Do the reconstruction
        for outer in range(nbreg):
            for inner in range(ninner):
                # update u
                rhs      =   murf + lmbda*Dxyzt(X-B) + gam*img
                img      =   np.fft.ifft2(np.fft.fft2(rhs)*uker)
                # update x and y
                A        =   Dxyz(img) + B
                X        =   shrink(A, 1/lmbda)
                # update bregman parameters
                B       =   A - X
            Kdata   = Kdata + Kdata0 - mask*np.fft.fftn(img)/scale
            murf    = np.fft.ifftn(mu*mask*Kdata)*scale

        # undo the normalization so that results are scaled properly
        img = img / norm_factor / scale
        CS1[:,:,:]  =   img
    
    
    return CS1

def phase_contrast(M1,M0,VENC,scantype='0G'):
    param       =   1
    if scantype=='-G+G':
        param   =   0.5
    return VENC*param*(np.angle(M1) - np.angle(M0))/np.pi
 
def GenerateMagnetization(Sq,VENC,noise,scantype='0G'):
    # MRI PARAMETERS
    gamma = 267.513e6  # rad/Tesla/sec   Gyromagnetic ratio for H nuclei
    B0 = 1.5  # Tesla           Magnetic Field Strenght
    TE = 5e-3  # Echo-time
    PHASE0 = np.zeros(Sq.shape)
    PHASE1 = np.zeros(Sq.shape)
    RHO0 = np.zeros(Sq.shape, dtype=complex)
    RHO1 = np.zeros(Sq.shape, dtype=complex)

    if np.ndim(Sq)==3:
        [row,col,numt2] = Sq.shape
        [X,Y]               		=   np.meshgrid(np.linspace(0,col,col),np.linspace(0,row,row))
        for k in range(numt2):
	        if noise:
	    	    Drho                    =   np.random.normal(0,0.2,[row,col])
	    	    Drho2                   =   np.random.normal(0,0.2,[row,col])
	        else:
		        Drho                    =   np.zeros([row,col])
		        Drho2                   =   np.zeros([row,col])
        
	        varPHASE0                   =   np.random.randint(-10,11,size=(row,col))*np.pi/180*(np.abs(Sq[:,:,k])<0.001) #Hugo's observation
	        modulus                     =   0.5 + 0.5*(np.abs(Sq[:,:,k])>0.001)
        
	        if scantype=='0G':
		        PHASE0[:,:,k]               =   (gamma*B0*TE+0.01*X)*(np.abs(Sq[:,:,k])>0.001) + 10*varPHASE0
		        PHASE1[:,:,k]               =   (gamma*B0*TE+0.01*X)*(np.abs(Sq[:,:,k])>0.001) + 10*varPHASE0 + np.pi*Sq[:,:,k]/VENC
        
	        if scantype=='-G+G':
		        PHASE0[:,:,k]               =   gamma*B0*TE*np.ones([row,col]) + 10*varPHASE0 - np.pi*Sq[:,:,k]/VENC
		        PHASE1[:,:,k]               =   gamma*B0*TE*np.ones([row,col]) + 10*varPHASE0 + np.pi*Sq[:,:,k]/VENC
	    
	        RHO0[:,:,k]                 =   modulus*np.cos(PHASE0[:,:,k]) + Drho + 1j*modulus*np.sin(PHASE0[:,:,k]) + 1j*Drho2
	        RHO1[:,:,k]                 =   modulus*np.cos(PHASE1[:,:,k]) + Drho + 1j*modulus*np.sin(PHASE1[:,:,k]) + 1j*Drho2

    
    if np.ndim(Sq)==4:
        [row, col, dep, numt2] = Sq.shape
        [X, Y, Z] = np.meshgrid(np.linspace(0, col, col), np.linspace(
            0, row, row), np.linspace(0, dep, dep))

        for k in range(numt2):

	        if noise:
		        Drho = np.random.normal(0, 0.2, [row, col, dep])
		        Drho2 = np.random.normal(0, 0.2, [row, col, dep])
	        else:
		        Drho = np.zeros([row, col, dep])
		        Drho2 = np.zeros([row, col, dep])

	        varPHASE0 = np.random.randint(-10, 11, size=(row, col, dep)) * \
                    np.pi/180*(np.abs(Sq[:, :, :, k]) < 0.001)
	        modulus = 0.5 + 0.5*(np.abs(Sq[:, :, :, k]) > 0.001)

	        if scantype == '0G':
		        PHASE0[:, :, :, k] = (gamma*B0*TE+0.01*X) * \
                            (np.abs(Sq[:, :, :, k]) > 0.001) + 10*varPHASE0
		        PHASE1[:, :, :, k] = (gamma*B0*TE+0.01*X)*(np.abs(Sq[:, :, :, k])
                                                     > 0.001) + 10*varPHASE0 + np.pi*Sq[:, :, :, k]/VENC

	        if scantype == '-G+G':
		        PHASE0[:, :, :, k] = gamma*B0*TE * \
		            np.ones([row, col, dep]) + varPHASE0 - np.pi*Sq[:, :, :, k]/VENC
		        PHASE1[:, :, :, k] = gamma*B0*TE * \
		            np.ones([row, col, dep]) + varPHASE0 + np.pi*Sq[:, :, :, k]/VENC

	        RHO0[:, :, :, k] = modulus*np.cos(PHASE0[:, :, :, k]) + \
                    Drho + 1j*modulus*np.sin(PHASE0[:, :, :, k]) + 1j*Drho2
	        RHO1[:, :, :, k] = modulus*np.cos(PHASE1[:, :, :, k]) + \
                    Drho + 1j*modulus*np.sin(PHASE1[:, :, :, k]) + 1j*Drho2



    return [RHO0,RHO1]

def undersampling(Sqx,Sqy,Sqz,options,savepath):
	
    R           =   options['cs']['R']
   
    for r in R:
		
        if rank==0:
            print('Using Acceleration Factor R = ' + str(r))
            print('Component x of M0')
        
        [M0,M1]	=	GenerateMagnetization(Sqx,options['cs']['VENC'],options['cs']['noise'])
        

        print('\n Component x of M0')
        M0_cs  	=   CSMETHOD(M0,r)
        print('\n Component x of M1')
        M1_cs  	=   CSMETHOD(M1,r)

        Sqx_cs	=	phase_contrast(M1_cs,M0_cs,options['cs']['VENC'])
        del M0,M1
        del M0_cs, M1_cs
        
        [M0,M1]	=	GenerateMagnetization(Sqy,options['cs']['VENC'],options['cs']['noise'])
        

        print('\n Component y of M0')
        M0_cs  	=   CSMETHOD(M0,r)
        print('\n Component y of M1')
        M1_cs  	=   CSMETHOD(M1,r)
        

        Sqy_cs	=	phase_contrast(M1_cs,M0_cs,options['cs']['VENC'])
        
        del M0,M1
        del M0_cs, M1_cs
        
        [M0,M1]	=	GenerateMagnetization(Sqz,options['cs']['VENC'],options['cs']['noise'])
   
        if rank==0:
            print('\n Component z of M0')
        M0_cs  	=   CSMETHOD(M0,r)
        if rank==0:
            print('\n Component z of M1')
        M1_cs	    =   CSMETHOD(M1,r)
        if rank==0:
            print(' ')
		
        Sqz_cs	=	phase_contrast(M1_cs,M0_cs,options['cs']['VENC'])	
		
		
        if rank==0:
            print('saving the sequences in ' + savepath)
            seqname	=	options['cs']['name'] +'_R' + str(r) + '.npz'
            print('sequence name: '  + seqname)
            np.savez_compressed( savepath + seqname, x=Sqx_cs, y=Sqy_cs,z=Sqz_cs)
    
        del Sqx_cs,Sqy_cs,Sqz_cs
	
def undersampling_peakpv(Sqx,Sqy,Sqz,options,R):
	
    Sqx_cs		=	{}
    Sqy_cs		=	{}
    Sqz_cs		=	{}
    [Mx0,Mx1]	=	GenerateMagnetization(Sqx,options['cs']['VENC'],options['cs']['noise'],scantype='0G')
    [My0,My1]	=	GenerateMagnetization(Sqy,options['cs']['VENC'],options['cs']['noise'],scantype='0G')
    [Mz0,Mz1]	=	GenerateMagnetization(Sqz,options['cs']['VENC'],options['cs']['noise'],scantype='0G')
    

    Mx0_cs  	=   CSMETHOD(Mx0,R)
    Mx1_cs  	=   CSMETHOD(Mx1,R)
    My0_cs  	=   CSMETHOD(My0,R)
    My1_cs  	=   CSMETHOD(My1,R)
    Mz0_cs  	=   CSMETHOD(Mz0,R)
    Mz1_cs	    =   CSMETHOD(Mz1,R)

    Sqx_cs	    =	phase_contrast(Mx1_cs,Mx0_cs,options['cs']['VENC'],scantype='0G')	
    Sqy_cs	    =	phase_contrast(My1_cs,My0_cs,options['cs']['VENC'],scantype='0G')	
    Sqz_cs	    =	phase_contrast(Mz1_cs,Mz0_cs,options['cs']['VENC'],scantype='0G')	
    
    
    
    return [Sqx_cs,Sqy_cs,Sqz_cs]
	
def undersampling_short(Mx,My,Mz,options):
	
    R           =   options['cs']['R']
    savepath    =   options['cs']['savepath']


    R_SENSE     =   1
    if 'R_SENSE' in options['cs']:
        R_SENSE     =   options['cs']['R_SENSE'][0]

    for r in R:	
        if rank==0:
            print('Using Acceleration Factor R = ' + str(r))
    

        if R_SENSE==2:    
            [MxS0_cs,MxS1_cs]  	=   CSMETHOD_SENSE(Mx,r,2)
            [MyS0_cs,MyS1_cs]  	=   CSMETHOD_SENSE(My,r,2)
            [MzS0_cs,MzS1_cs]  	=   CSMETHOD_SENSE(Mz,r,2)
            if rank==0:
                print('saving the sequences in ' + savepath)
                seqname_s0	    =	options['cs']['name'] +'S0_R' + str(r) + '.npz'
                seqname_s1	    =	options['cs']['name'] +'S1_R' + str(r) + '.npz'
                print('sequence name: '  + seqname_s0)
                np.savez_compressed( savepath + seqname_s0, x=MxS0_cs, y=MyS0_cs,z=MzS0_cs)
                print('sequence name: '  + seqname_s1)
                np.savez_compressed( savepath + seqname_s1, x=MxS1_cs, y=MyS1_cs,z=MzS1_cs)
            del MxS0_cs, MyS0_cs, MzS0_cs
            del MxS1_cs, MyS1_cs, MzS1_cs
        elif R_SENSE==1:
            Mx_cs  	=   CSMETHOD(Mx,r)
            My_cs  	=   CSMETHOD(My,r)
            Mz_cs  	=   CSMETHOD(Mz,r)
            if rank==0:
                print('saving the sequences in ' + savepath)
                seqname	    =	options['cs']['name'] +'_R' + str(r) + '.npz'
                print('sequence name: '  + seqname)
                np.savez_compressed( savepath + seqname, x=Mx_cs, y=My_cs,z=Mz_cs)
            del Mx_cs,My_cs,Mz_cs
        else:
            raise Exception('Only implemented for 2-fold SENSE!!')






# THE END