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()

# SENSE:    Simulation of SENSitive Encoding algorithm proposed by K. Pruessmann, et. al. in:
#           "SENSE: Sensitivity Enconding for Fast MRI" Mag. Res. in Medicine 42. (1999)
#           written by Jeremias Garay (j.e.garay.labra@rug.nl)

def Sensitivity_Map(shape):

    [Nx,Ny,Nz]  =   shape
    [X,Y,Z]     =   np.meshgrid(np.linspace(0,Ny,Ny),np.linspace(0,Nx,Nx),np.linspace(0,Nz,Nz))
    Xsense1     =   (X/(Nx*2)-1)**2
    Xsense2     =   ((Nx-X)/(Nx*2)-1)**2
    S_MAPS      =   [np.fft.fftshift(Xsense1),np.fft.fftshift(Xsense2)]

    return  S_MAPS

def SENSE_recon(S1,M1,S2,M2):
    
    [Nx,Ny,Nz,Nt]       =   M1.shape
    M                   =   np.zeros([Nx,int(2*Ny),Nz,Nt],dtype=complex)
    sm1                 =   np.fft.fftshift(S1)[:,:,0]
    sm2                 =   np.fft.fftshift(S2)[:,:,0]
    
    for j in range(Ny):
        for k in range(Nx):
            l1 = M1[k,j,:,:];    a1 = sm1[k,j];   a2 = sm1[k,j+Ny]  
            l2 = M2[k,j,:,:];    b1 = sm2[k,j];   b2 = sm2[k,j+Ny]
            B                   =   (l1*b1 - l2*a1)/(a2*b1 - b2*a1)
            A                   =   (l1*b2 - l2*a2)/(a1*b2 - a2*b1)
            M[k,j,:,:]          =   A
            M[k,j+Ny,:,:]       =   B
        

    return M

def SENSE_recon2(S1,M1,S2,M2):
    # With matrices as in the original paper!

    [Nx,Ny,Nz,Nt]       =   M1.shape
    M                   =   np.zeros([Nx,int(2*Ny),Nz,Nt],dtype=complex)
    sm1                 =   np.fft.fftshift(S1)[:,:,0]
    sm2                 =   np.fft.fftshift(S2)[:,:,0]
    sigma2              =   0.049**2
    sigma2              =   1
    Psi                 =   np.diagflat(np.array([sigma2,sigma2])) # Error matrix Psi
    Psi_inv             =   np.linalg.inv(Psi)
    
    for j in range(Ny):
        for k in range(Nx):
            l1 = M1[k,j,:,:];    a1 = sm1[k,j];   a2 = sm1[k,j+Ny]  
            l2 = M2[k,j,:,:];    b1 = sm2[k,j];   b2 = sm2[k,j+Ny]
            S                   =   np.array([[a1,a2],[b1,b2]])
            U                   =   np.linalg.inv((np.transpose(S)*Psi_inv*S))*np.transpose(S)*Psi_inv 
            a                   =   np.array([l1,l2])
            a_resized           =   np.resize(a,(2,Nz*Nt))   
            v_resized           =   np.dot(U,a_resized)
            v                   =   np.resize(v_resized,(2,Nz,Nt))
            M[k,j,:,:]          =   v[0,:,:]
            M[k,j+Ny,:,:]       =   v[1,:,:]
        

    return M

def SENSE_METHOD(Seq,R): 
    '''
    Args:
        ITOT:     	a numpy matrix with the full sampled (3D or 4D) dynamical data
        R:        	the acceleration factor
    '''

    [row,col,dep,numt2]         =   Seq.shape
    Seq_red                     =   {}
    SenseMAP                    =   {}         
    [SenseMAP[0],SenseMAP[1]]   =   Sensitivity_Map([row,col,dep])

    col2                        =   int(np.ceil(col/2))
   
    for rs in range(R):
        Seq_red[rs]                   =   np.zeros([row,col2,dep,numt2],dtype=complex)
        for t in range(numt2):
            Kdata_0                 =   np.fft.fftn(Seq[:,:,:,t])
            Kdata_0                 =   Kdata_0*SenseMAP[rs]
            Kdata_0                 =   Kdata_0[:,0::R,:]
            Seq_red[rs][:,:,:,t]    =   np.fft.ifftn(Kdata_0)

    Seq_recon                   =   SENSE_recon2(SenseMAP[0],Seq_red[0],SenseMAP[1],Seq_red[1])
        
    return Seq_recon
     
def undersampling(Mx,My,Mz,options):
	
    R           =   options['SENSE']['R']

    for r in R:	
        if rank==0:
            print('Using Acceleration Factor R = ' + str(r))
            print('applying into x component')
        Mx_s  	=   SENSE_METHOD(Mx,r)
        if rank==0:
            print('applying into y component')
        My_s  	=   SENSE_METHOD(My,r)
        if rank==0:
            print('applying into z component')
        Mz_s  	=   SENSE_METHOD(Mz,r)
    
    return [Mx_s,My_s,Mz_s]