import numpy
from ..src._helper import helper#, helper1

def _init__cpu(self):
        """
        Constructor.
 
        :param None:
        :type None: Python NoneType
        :return: NUFFT: the pynufft_hsa.NUFFT instance
        :rtype: NUFFT: the pynufft_hsa.NUFFT class
        :Example:
 
        >>> from pynufft import NUFFT
        >>> NufftObj = NUFFT()
        """
        self.dtype = numpy.complex64  # : initial value: numpy.complex64
        self.debug = 0  #: initial value: 0
        self.Nd = ()  # : initial value: ()
        self.Kd = ()  # : initial value: ()
        self.Jd = ()  #: initial value: ()
        self.ndims = 0  # : initial value: 0
        self.ft_axes = ()  # : initial value: ()
        self.batch = None  # : initial value: None
        pass    

def _plan_cpu(self, om, Nd, Kd, Jd, ft_axes=None):
    """
    Plan the NUFFT object with the geometry provided.

    :param om: The M off-grid locates in the frequency domain,
                which is normalized between [-pi, pi]
    :param Nd: The matrix size of the equispaced image.
               Example: Nd=(256,256) for a 2D image;
                         Nd = (128,128,128) for a 3D image
    :param Kd: The matrix size of the oversampled frequency grid.
               Example: Kd=(512,512) for 2D image;
                        Kd = (256,256,256) for a 3D image
    :param Jd: The interpolator size.
               Example: Jd=(6,6) for 2D image;
                        Jd = (6,6,6) for a 3D image
    :param ft_axes: (Optional) The axes for Fourier transform.
                    The default is all axes if 'None' is given.
    :param batch: (Optional) Batch mode.
                 If the batch is provided, the last appended axis is the number
                 of identical NUFFT to be transformed.
                 The default is 'None'.
    :type om: numpy.float array, matrix size = M * ndims
    :type Nd: tuple, ndims integer elements.
    :type Kd: tuple, ndims integer elements.
    :type Jd: tuple, ndims integer elements.
    :type ft_axes: None, or tuple with optional integer elements.
    :returns: 0
    :rtype: int, float

    :ivar Nd: initial value: Nd
    :ivar Kd: initial value: Kd
    :ivar Jd: initial value: Jd
    :ivar ft_axes: initial value: None

    :Example:

    >>> from pynufft import NUFFT
    >>> NufftObj = NUFFT()
    >>> NufftObj.plan(om, Nd, Kd, Jd)

    or

    >>> NufftObj.plan(om, Nd, Kd, Jd, ft_axes)

    """

    self.ndims = len(Nd)  # : initial value: len(Nd)
    if ft_axes is None:
        ft_axes = range(0, self.ndims)
    self.ft_axes = ft_axes  # default: all axes (range(0, self.ndims)

    self.st = helper.plan(om, Nd, Kd, Jd, ft_axes=ft_axes,
                          format='CSR')

    self.Nd = self.st['Nd']  # backup
    self.Kd = self.st['Kd']
    # backup
    self.sn = numpy.asarray(self.st['sn'].astype(self.dtype), order='C')

#     if batch is None:  # single-coil
    self.parallel_flag = 0
    self.batch = 1

#     else:  # multi-coil
#         self.parallel_flag = 1
#         self.batch = batch

#     if self.parallel_flag is 1:
#         self.multi_Nd = self.Nd + (self.batch, )
#         self.uni_Nd = self.Nd + (1, )
#         self.multi_Kd = self.Kd + (self.batch, )
#         self.multi_M = (self.st['M'], ) + (self.batch, )
#         self.multi_prodKd = (numpy.prod(self.Kd), self.batch)
#         self.sn = numpy.reshape(self.sn, self.Nd + (1,))

#     elif self.parallel_flag is 0:
    self.Nd = self.Nd  # + (self.Reps, )
    self.Nd = self.Nd
    self.Kd = self.Kd  # + (self.Reps, )
    self.M = (self.st['M'], )
    self.prodKd = (numpy.prod(self.Kd), )

    # Calculate the density compensation function
    self.sp = self.st['p'].copy().tocsr()
    
    self.spH = (self.st['p'].getH().copy()).tocsr()
    
    self.Kdprod = numpy.int32(numpy.prod(self.st['Kd']))
    self.Jdprod = numpy.int32(numpy.prod(self.st['Jd']))
    del self.st['p'], self.st['sn']

    self.NdCPUorder, self.KdCPUorder, self.nelem = helper.preindex_copy(
        self.st['Nd'],
        self.st['Kd'])
#     self.volume = {}
#     self.volume['cpu_coil_profile'] = numpy.ones(self.multi_Nd)

    return 0


def _precompute_sp_cpu(self):
    """

    Private: Precompute adjoint (gridding) and Toepitz interpolation
             matrix.

    :param None:
    :type None: Python Nonetype
    :return: self: instance
    """
    try:
        W0 = numpy.ones((self.st['M'],), dtype=numpy.complex64)
        W = self.xx2k(self.adjoint(W0))
        self.W = (W*W.conj())**0.5
        del W0
        del W
    except:
        print("errors occur in self.precompute_sp()")
        raise




def _solve_cpu(self, y, solver=None, *args, **kwargs):
    """
    Solve NUFFT.
    :param y: data, numpy.complex64. The shape = (M,) or (M, batch)
    :param solver: 'cg', 'L1TVOLS', 'lsmr', 'lsqr', 'dc', 'bicg',
                   'bicgstab', 'cg', 'gmres','lgmres'
    :param maxiter: the number of iterations
    :type y: numpy array, dtype = numpy.complex64
    :type solver: string
    :type maxiter: int
    :return: numpy array with size.
            The shape = Nd ('L1TVOLS') or  Nd 
            ('lsmr', 'lsqr', 'dc','bicg','bicgstab','cg', 'gmres','lgmres')
    """
    from ..linalg.solve_cpu import solve
    x2 = solve(self,  y,  solver, *args, **kwargs)
    return x2  # solve(self,  y,  solver, *args, **kwargs)

def _forward_cpu(self, x):
    """
    Forward NUFFT on CPU

    :param x: The input numpy array, with the size of Nd
    :type: numpy array with the dtype of numpy.complex64
    :return: y: The output numpy array, with the size of (M,)
    :rtype: numpy array with the dtype of numpy.complex64
    """
    y = self._k2y_cpu(self._xx2k_cpu(self._x2xx_cpu(x)))

    return y

def _adjoint_cpu(self, y):
    """
    Adjoint NUFFT on CPU

    :param y: The input numpy array, with the size of (M,) 
    :type: numpy array with the dtype of numpy.complex64
    :return: x: The output numpy array,
                with the size of Nd or Nd 
    :rtype: numpy array with the dtype of numpy.complex64
    """
    x = self._xx2x_cpu(self._k2xx_cpu(self._y2k_cpu(y)))

    return x

def _selfadjoint_one2many2one_cpu_deprecated(self, x):
    y2 = self._forward_one2many_cpu(x)
    x2 = self._adjoint_many2one_cpu(y2)
    del y2
    return x2

def _selfadjoint_cpu(self, x):
    """
    selfadjoint NUFFT on CPU

    :param x: The input numpy array, with size=Nd
    :type: numpy array with dtype =numpy.complex64
    :return: x: The output numpy array, with size=Nd
    :rtype: numpy array with dtype =numpy.complex64
    """
    # x2 = self.adjoint(self.forward(x))

    x2 = self._xx2x_cpu(self._k2xx_cpu(self._k2y2k_cpu(self._xx2k_cpu(self._x2xx_cpu(x)))))

    return x2

def _selfadjoint2_cpu(self, x):
    try:
        x2 = self._k2xx_cpu(self.W * self._xx2k_cpu(x))
    except:
        self._precompute_sp_cpu()
        x2 = self._k2xx_cpu(self.W * self._xx2k_cpu(x))
    return x2

def _x2xx_cpu(self, x):
    """
    Private: Scaling on CPU
    Inplace multiplication of self.x_Nd by the scaling factor self.sn.
    """
    xx = x * self.sn
    return xx

def _xx2k_cpu(self, xx):
    """
    Private: oversampled FFT on CPU

    Firstly, zeroing the self.k_Kd array
    Second, copy self.x_Nd array to self.k_Kd array by cSelect
    Third, inplace FFT
    """
    
    output_x = numpy.zeros(self.Kd, dtype=self.dtype, order='C')

    for bat in range(0, self.batch):
        output_x.ravel()[self.KdCPUorder * self.batch + bat] = xx.ravel()[
            self.NdCPUorder * self.batch + bat]

    k = numpy.fft.fftn(output_x, axes=self.ft_axes)
    
    return k

def _xx2k_one2one_cpu(self, xx):
    """
    Private: oversampled FFT on CPU

    First, zeroing the self.k_Kd array
    Second, copy self.x_Nd array to self.k_Kd array by cSelect
    Third, inplace FFT
    """
    
    output_x = numpy.zeros(self.st['Kd'], dtype=self.dtype, order='C')

    # for bat in range(0, self.batch):
    output_x.ravel()[self.KdCPUorder] = xx.ravel()[self.NdCPUorder]

    k = numpy.fft.fftn(output_x, axes=self.ft_axes)
    
    return k

def _k2vec_cpu(self, k):
    k_vec = numpy.reshape(k, self.prodKd, order='C')
    return k_vec

def _vec2y_cpu(self, k_vec):
    '''
    gridding:
    '''
    y = self.sp.dot(k_vec)
    # y = self.st['ell'].spmv(k_vec)

    return y

def _k2y_cpu(self, k):
    """
    Private: interpolation by the Sparse Matrix-Vector Multiplication
    """
    y = self._vec2y_cpu(self._k2vec_cpu(k))
    # numpy.reshape(self.st['p'].dot(Xk), (self.st['M'], ), order='F')
    return y

def _y2vec_cpu(self, y):
    '''
   regridding non-uniform data (unsorted vector)
    '''
    # k_vec = self.st['p'].getH().dot(y)
    k_vec = self.spH.dot(y)
    # k_vec = self.st['ell'].spmvH(y)

    return k_vec

def _vec2k_cpu(self, k_vec):
    '''
    Sorting the vector to k-spectrum Kd array
    '''
    k = numpy.reshape(k_vec, self.Kd, order='C')

    return k

def _y2k_cpu(self, y):
    """
    Private: gridding by the Sparse Matrix-Vector Multiplication
    """
    k_vec = self._y2vec_cpu(y)
    k = self._vec2k_cpu(k_vec)
    return k

def _k2xx_cpu(self, k):
    """
    Private: the inverse FFT and image cropping (which is the reverse of
             _xx2k() method)
    """
#         dd = numpy.size(self.Kd)

    k = numpy.fft.ifftn(k, axes=self.ft_axes)
    xx = numpy.zeros(self.Nd, dtype=self.dtype, order='C')
    for bat in range(0, self.batch):
        xx.ravel()[self.NdCPUorder * self.batch + bat] = k.ravel()[
            self.KdCPUorder * self.batch + bat]
#         xx = xx[crop_slice_ind(self.Nd)]
    return xx


def _k2xx_one2one_cpu(self, k):
    """
    Private: the inverse FFT and image cropping
             (which is the reverse of _xx2k() method)
    """
#         dd = numpy.size(self.Kd)

    k = numpy.fft.ifftn(k, axes=self.ft_axes)
    xx = numpy.zeros(self.st['Nd'], dtype=self.dtype, order='C')
    # for bat in range(0, self.batch):
    xx.ravel()[self.NdCPUorder] = k.ravel()[self.KdCPUorder]
    # xx = xx[crop_slice_ind(self.Nd)]
    return xx

def _xx2x_cpu(self, xx):
    """
    Private: rescaling, which is identical to the  _x2xx() method
    """
    x = self._x2xx_cpu(xx)
    return x

def _k2y2k_cpu(self, k):
    """
    Private: the integrated interpolation-gridding by the Sparse
             Matrix-Vector Multiplication
    """

    Xk = self._k2vec_cpu(k)
    k = self._y2vec_cpu(self._vec2y_cpu(Xk))
    k = self._vec2k_cpu(k)
    return k


## host code
def _forward_host(self, x):
    gx = self.to_device(x)
    y = self._forward_device(gx).get()
    return y

def _forward_legacy_host(self, x):
    gx = self.to_device(x)
    y = self._forward_legacy(gx).get()
    return y

def _adjoint_host(self, y):
    gy = self.to_device(y)
    gx = self._adjoint_device(gy)
    return gx.get()

def _adjoint_legacy_host(self, y):
    gy = self.to_device(y)
    gx = self._adjoint_legacy(gy)
    return gx.get()
    
def _selfadjoint_host(self, x):
    gx = self.to_device(x)
    gx2 = self._selfadjoint_device(gx)
    return gx2.get()
    
def _selfadjoint_legacy_host(self, x):
    gx = self.to_device(x)
    gx2 = self._selfadjoint_legacy(gx)
    return gx2.get()
    
def _solve_host(self, y, *args, **kwargs):
    gy = self.to_device(y)
    x = self._solve_device(gy, *args, **kwargs)
    return x.get()

def _solve_legacy_host(self, y, *args, **kwargs):
    gy = self.to_device(y)
    x = self._solve_legacy(gy, *args, **kwargs)
    return x.get()

def _xx2k_host(self, xx):
    gxx = self.to_device(xx)
    gk = self._xx2k_device(gxx)
    return gk.get()


def _k2xx_host(self, k):
    gk = self.to_device(k)
    gxx = self._k2xx_device(gk)
    return gxx.get()

def _x2xx_host(self, x):
    gx = self.to_device(x)
    gxx = self._x2xx_device(gx)
    return gxx.get()

def _xx2x_host(self, x):
    gx = self.to_device(x)
    gxx = self._xx2x_device(gx)
    return gxx.get()

def _k2y_host(self, k):
    gk = self.to_device(k)
    gy = self._k2y_device(gk)
    return gy.get()

def _y2k_host(self, y):
    gy = self.to_device(y)
    gk = self._y2k_device(gy)
    return gk.get()

def _k2y_legacy_host(self, k):
    gk = self.to_device(k)
    gy = self._k2y_legacy(gk)
    return gy.get()

def _y2k_legacy_host(self, y):
    gy = self.to_device(y)
    gk = self._y2k_legacy(gy)
    return gk.get()