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

            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 sense algorith with SENSE... in contruction!.

    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 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'):
    ''' Simulation of a typical magnetization. A x-dependent plane is added into the
    reference phase.
    '''
    # 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_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