"""
CPU solvers
======================================
"""

import scipy
import scipy.sparse.linalg
import numpy
from ..src._helper import helper

def cDiff(x, d_indx):
        """
        Compute image gradient, which needs the results of indxmap_diff(Nd)
        :param x: The image array
        :param d_indx: The index of the shifted image
        :type x: numpy.float array, matrix size = Nd
        :type d_indx: int32
        :returns: x_diff: Image gradient determined by d_indx
        :rtype: x_diff: numpy.complex64      
        """
        x_diff=numpy.asarray(x.copy(),order='C')
        x_diff.flat =   x_diff.ravel()[d_indx] - x.ravel()
        return x_diff
    
def _create_kspace_sampling_density(nufft):
        """
        Compute k-space sampling density
        """    
        y = numpy.ones(nufft.st['M'],dtype = numpy.complex64)
#         w = numpy.abs(nufft.xx2k(nufft.adjoint(y)))
        
        if nufft.parallel_flag is 1:
            w =  numpy.abs( nufft.xx2k(nufft.adjoint(y)))[..., 0]#**2) ))
        else:
            w =  numpy.abs( nufft.xx2k(nufft.adjoint(y)))#**2) ))
        nufft.st['w'] = w#self.nufftobj.vec2k(w)
        RTR=nufft.st['w'] # see __init__() in class "nufft"
#         print('RTR.shape = ', RTR.shape)
        return RTR
   
    
def L1TVOLS(nufft, y, maxiter, rho ): # main function of solver
    """
    L1-total variation regularized ordinary least square 
    """
    mu = 1.0
    LMBD = rho*mu

    def AHA(x):
        x2 = nufft.selfadjoint(x)
        return x2
    def AH(y):
        
        x2 = nufft.adjoint(y.reshape(nufft.st['M'], order='C'))
        return x2
    
        
    
    uker = mu*_create_kspace_sampling_density(nufft)
#     print('uker.shape', uker.shape) 
    uker = uker - LMBD* helper.create_laplacian_kernel(nufft)
#     print('uker.shape', uker.shape)
#     import matplotlib.pyplot
#     matplotlib.pyplot.imshow(abs(uker))
#     matplotlib.pyplot.show()
#     print('y.shape=', y.shape)
    AHy = AH(y)
#     print('AHy.shape = ', AHy.shape)
    
    xkp1 = numpy.zeros_like(AHy)
    AHyk = numpy.zeros_like(AHy)
           
#         self._allo_split_variables()        
    zz= []
    bb = []
    dd = []
    d_indx, dt_indx = helper.indxmap_diff(nufft.st['Nd'])
    z=numpy.zeros(nufft.st['Nd'], dtype = nufft.dtype, order='C')
    
    ndims = len(nufft.st['Nd'])
    s_tmp = []
    for pp in range(0,ndims):
            s_tmp += [0, ]
            
            
    for jj in range(    0,  ndims): # n_dims + 1 for wavelets
        
        zz += [z.copy(),]
        bb += [z.copy(),]
        dd +=  [z.copy(),]
#     zf = z.copy()
#     bf = z.copy()
#     df = z.copy()

    n_dims = len(nufft.st['Nd'])#numpy.size(uf.shape)
    
    for outer in numpy.arange(0, maxiter):
#             for inner in numpy.arange(0,nInner):
            
        # solve Ku = rhs
            
        rhs = mu*AHyk# + df - bf)   # right hand side
        for pp in range(0,ndims):
#             diff1 =  dd[pp] - bb[pp]
#             diff1 = numpy.roll( diff1, 1, axis = pp) - diff1
            rhs += LMBD*(cDiff(dd[pp] - bb[pp], dt_indx[pp]))
#             del diff1
#                 LMBD*(cDiff(dd[1] - bb[1],  dt_indx[1]))  )          
        # Note K = F' uker F
        # so K-1 ~ F
        xkp1 = nufft.k2xx_one2one( (nufft.xx2k_one2one(rhs)+1e-7) / (uker+1e-7)) 
#                 self._update_d(xkp1)
        for pp in range(0,ndims):
            
#             zz[pp] = numpy.roll( xkp1, -1, axis = pp) - xkp1            
            zz[pp] = cDiff(xkp1, d_indx[pp])
#         zz[0] = cDiff(xkp1,  d_indx[0])
#         zz[1] = cDiff(xkp1,  d_indx[1])
        zf = AHA(xkp1)  -AHy 

        '''
        soft-thresholding the edges
        '''
        s = numpy.zeros_like(zz[pp]) # complex
        for pp in range(0,ndims):
            s_tmp[pp] = zz[pp] + bb[pp]

            s_r = numpy.hypot(s_tmp[pp].real, s_tmp[pp].imag)
            s = numpy.hypot(s_r, s.real)
#             s = numpy.sqrt(s**2 + s_tmp[pp]**2)
#             else:
#                 s =  s_tmp[pp]
#         s1 = zz[0] + bb[0]
#         s2 = zz[1] + bb[1]
#         s = sum((s_tmp[pp])**2 for pp in range(0,n_dims))
#         s = s1**2 + s2**2
        s += 1e-5

        threshold_value = 1/LMBD
        r =(s > threshold_value)*(s-threshold_value)/s#numpy.maximum(s - threshold_value ,  0.0)/s
        for pp in range(0,ndims):
            dd[pp] = s_tmp[pp]*r
#         dd[0] = s1*r
#         dd[1] = s2*r
#         df = zf+bf
#         
#         threshold_value=1.0/mu
#     
#         df.real =0.0+ (df.real>threshold_value)*(df.real - threshold_value) +(df.real<= - threshold_value)*(df.real+threshold_value)
#         df.imag = 0.0+(df.imag>threshold_value)*(df.imag - threshold_value) +(df.imag<= - threshold_value)*(df.imag+threshold_value) 
#                 df =     sy
        # end of shrinkage
        for pp in range(0,ndims):
            bb[pp] += zz[pp] - dd[pp]
#         bb[0] += zz[0] - dd[0] 
#         bb[1] += zz[1] - dd[1] 
#         bf += zf - df 
#                 self._update_b() # update b based on the current u

        AHyk = AHyk - zf # Linearized Bregman iteration f^k+1 = f^k + f - Au
#             print(outer)

    return xkp1 #(u,u_stack)

def _pipe_density(nufft, maxiter):
    '''
    Create the density function by iterative solution
    Generate pHp matrix
    '''
#         W = pipe_density(self.st['p'])
    # sampling density function
              
    W = numpy.ones(nufft.st['M'],dtype=nufft.dtype)
#         V1= self.st['p'].getH()
    #     VVH = V.dot(V.getH()) 
         
    for pp in range(0,maxiter):
#             E = self.st['p'].dot(V1.dot(W))

        E = nufft.forward(nufft.adjoint(W))
        W = (W/E)
   
    return W
 
def solve(nufft,   y,  solver=None, *args, **kwargs):
        """
        Solve NUFFT.
        The current version supports solvers = 'cg' or 'L1TVOLS' or 'L1TVLAD'.
        
        :param nufft: NUFFT_cpu object
        :param y: (M,) array, non-uniform data
        :return: x: image
            
        """
        if ('cgs' == solver) or ('qmr' ==solver) or ('minres'==solver):
            raise TypeError(solver +' requires real symmetric matrix')
        
    
        if None ==  solver:
    #         solver  =   'cg'
            print("solver must be one of the following solvers; \newline dc, cg, bicg, bicgstab, gmres, lgmres, L1OLS, L1LAD")
     
    
        if 'dc'   ==  solver:
            """
            Density compensation method
            nufft.st['W'] will be computed if doesn't exist
            If nufft.st['W'] exist then x2 = nufft.adjoint(nufft.st['W']*y)
            input: 
                y: (M,) array
            output:
                x2: Nd array
            """
            print(solver)
    #             try:
    # 
    #                 x = nufft.adjoint(nufft.st['W']*y)
    # 
    #             except: 
    
            W = _pipe_density( nufft, *args, **kwargs)
    
            x = nufft.adjoint(W*y)
    
            return x
        elif ('lsmr'==solver) or ('lsqr'==solver):
            """
            Assymetric matrix A
            Minimize L2 norm of |y-Ax|_2 or |y-Ax|_2+|x|_2
            Very stable 
            """
#                 A = nufft.sp
            def sp(k):
                k2 = k.reshape(nufft.Kd, order='C')
                return nufft.k2y(k2).ravel()
            def spH(y):
                y2 = y.reshape(nufft.st['M'], order='C')
                return nufft.y2k(y2).ravel()                
#                 return nufft.spH.dot(nufft.sp.dot(x))
#                 print('shape', (nufft.st['M']*nufft.batch, nufft.Kdprod*nufft.batch))
#                 print('spH ')
            A = scipy.sparse.linalg.LinearOperator((nufft.st['M']*nufft.batch, nufft.Kdprod*nufft.batch), matvec = sp, rmatvec = spH, )
            
            methods={'lsqr':scipy.sparse.linalg.lsqr,
                                'lsmr':scipy.sparse.linalg.lsmr,}
            k2 = methods[solver](A,  y.flatten(), *args, **kwargs)#,show=True)
            vec = k2[0]
            vec.shape = nufft.Kd
            xx = nufft.k2xx(vec)
            x= xx/nufft.sn
            return x#, k2[1:]        
        elif 'L1TVOLS' == solver:
            return  L1TVOLS(nufft, y, *args, **kwargs)
#         elif 'L1TVLAD' == solver:
#             return  L1TVLAD(nufft, y, *args, **kwargs)
        else:
            """
            Hermitian matrix A
            cg: conjugate gradient
            bicgstab: biconjugate gradient stablizing
            bicg: biconjugate gradient
            gmres: 
            lgmres:
            """
#             A = nufft.spHsp#nufft.st['p'].getH().dot(nufft.st['p'])
            def spHsp(x):
                k = x.reshape(nufft.Kd, order='C')
                return nufft.k2y2k(k).ravel()
#                 return nufft.spH.dot(nufft.sp.dot(x))
            
            A = scipy.sparse.linalg.LinearOperator((nufft.Kdprod*nufft.batch, nufft.Kdprod*nufft.batch), matvec = spHsp, rmatvec = spHsp, )

            
            methods={'cg':scipy.sparse.linalg.cg,   
                                 'bicgstab':scipy.sparse.linalg.bicgstab, 
                                 'bicg':scipy.sparse.linalg.bicg, 
    #                                  'cgs':scipy.sparse.linalg.cgs, 
                                 'gmres':scipy.sparse.linalg.gmres, 
                                 'lgmres':scipy.sparse.linalg.lgmres, 
    #                                  'minres':scipy.sparse.linalg.minres, 
    #                                  'qmr':scipy.sparse.linalg.qmr, 
                                 }
            k2 = methods[solver](A,  nufft.y2k(y).ravel(), *args, **kwargs)#,show=True)
    
    
            xx = nufft.k2xx(k2[0].reshape(nufft.Kd))
            x= xx/nufft.sn
            return x#     , k2[1:]