####################################
#     HSA code
####################################

import numpy
# import reikna
from functools import wraps as _wraps
from ..src._helper import helper#, helper1

def push_cuda_context(hsa_method):
    """
    Decorator: Push cuda context to the top of the stack before calling the context
    Add @push_cuda_context before the methods of NUFFT_device()
    """
    @_wraps(hsa_method)
    def wrapper(*args, **kwargs):
        try:
            args[0].thr._context.push()
        except:
            pass
        return hsa_method(*args, **kwargs)
    return wrapper


def _init__device(self, device_indx=None):
    """
    Constructor.

    :param API: The API for the heterogeneous system. API='cuda'
                or API='ocl'
    :param platform_number: The number of the platform found by the API.
    :param device_number: The number of the device found on the platform.
    :param verbosity: Defines the verbosity level, default value is 0
    :type API: string
    :type platform_number: integer
    :type device_number: integer
    :type verbosity: integer
    :returns: 0
    :rtype: int, float

    :Example:

    >>> from pynufft import NUFFT
    >>> NufftObj = NUFFT_device(API='cuda', platform_number=0,
                                     device_number=0, verbosity=0)
    """

    self.dtype = numpy.complex64
    self.verbosity = 0#verbosity

    from reikna import cluda
    import reikna.transformations
    from reikna.cluda import functions, dtypes
#         try:  # try to create api/platform/device using the given parameters
    API = device_indx[0]
    self.API = API
#             if 'cuda' == API:
#                 api = cluda.cuda_api()
#             elif 'ocl' == API:
#                 api = cluda.ocl_api()
#         api = device_indx[0]
    platform_number = device_indx[1]
    device_number = device_indx[2]
    
#         NUFFT_hsa.__init__(self, API, platform_number, device_number)
    
    platform = device_indx[3]
    device = device_indx[4]
#             self.thr = api.Thread(device)
    self.device = device
    self.thr = device_indx[5]
#             api = device_indx[6]
#             self.thr = api.Thread(device)  # pyopencl.create_some_context()
#         self.device = device  # : device name

#         """
#         Wavefront: as warp in cuda. Can control the width in a workgroup
#         Wavefront is required in spmv_vector as it improves data coalescence.
#         see cCSR_spmv and zSparseMatVec
#         """
    self.wavefront = device_indx[6]
#         if self.verbosity > 0:
#             print('Wavefront of OpenCL (as wrap of CUDA) = ', self.wavefront)
#             print('API = ', API)
#             print('thr = ',  self.thr
#                   )
    from ..src import re_subroutine  # import create_kernel_sets
    kernel_sets = re_subroutine.create_kernel_sets(self.API)

    prg = self.thr.compile(kernel_sets,
                           render_kwds=dict(LL=str(self.wavefront)),
                           fast_math=False)
    self.prg = prg
    self.processor = 'hsa'



def _plan_device(self, om, Nd, Kd, Jd, ft_axes=None, radix=None):
    """
    Design the multi-coil or single-coil memory reduced interpolator.


    :param om: The M off-grid locations in the frequency domain.
               Normalized between [-pi, pi]
    :param Nd: The matrix size of 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: The dimensions to be transformed by FFT.
               Example: ft_axes = (0, 1) for 2D,
               ft_axes = (0, 1, 2) for 3D;
               ft_axes = None for all dimensions.
    :param radix: expert mode.
                If provided, the shape is Nd.
                The last axis is the number of parallel coils.
    :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: tuple, selected axes to be transformed.
    :returns: 0
    :rtype: int, float
    :Example:

    >>> import pynufft
    >>> device=pynufft.helper.device_list()[0]
    >>> NufftObj = pynufft.NUFFT(device)
    >>> NufftObj.plan(om, Nd, Kd, Jd)

    """

    self.ndims = len(Nd)  # dimension
    if ft_axes is None:
        ft_axes = range(0, self.ndims)
    self.ft_axes = ft_axes

    self.st = helper.plan(om, Nd, Kd, Jd, ft_axes=ft_axes,
                          format='pELL', radix=radix)
#     if batch is None:
    self.parallel_flag = 0
#     else:
#         self.parallel_flag = 1

#     if batch is None:
    self.batch = numpy.uint32(1)

#     else:
#         self.batch = numpy.uint32(batch)

    self.Nd = self.st['Nd']  # backup
    self.Kd = self.st['Kd']
    #  self.sn = numpy.asarray(self.st['sn'].astype(self.dtype),
    #                            order='C')# backup
#     if self.batch == 1:
#         self.multi_Nd = self.Nd
#         self.multi_Kd = self.Kd
#         self.multi_M = (self.st['M'], )
#         # Broadcasting the sense and scaling factor (Roll-off)
#         # self.sense2 = self.sense*numpy.reshape(self.sn, self.Nd + (1, ))
#     else:  # self.batch is 0:
#         self.multi_Nd = self.Nd + (self.batch, )
#         self.multi_Kd = self.Kd + (self.batch, )
#         self.multi_M = (self.st['M'], ) + (self.batch, )

    self.Kdprod = numpy.uint32(numpy.prod(self.st['Kd']))
    self.Jdprod = numpy.uint32(numpy.prod(self.st['Jd']))
    self.Ndprod = numpy.uint32(numpy.prod(self.st['Nd']))

    self.Nd_elements, self.invNd_elements = helper.strides_divide_itemsize(
                                                self.st['Nd'])
    # only return the Kd_elements
    self.Kd_elements = helper.strides_divide_itemsize(self.st['Kd'])[0]

    self._offload_device()

    return 0    

@push_cuda_context
def _offload_device(self):  # API, platform_number=0, device_number=0):
    """
    self.offload():

    Off-load NUFFT to the opencl or cuda device(s)

    :param API: define the device type, which can be 'cuda' or 'ocl'
    :param platform_number: define which platform to be used.
                            The default platform_number = 0.
    :param device_number: define which device to be used.
                        The default device_number = 0.
    :type API: string
    :type platform_number: int
    :type device_number: int
    :return: self: instance
    """

    self.pELL = {}  # dictionary

    self.pELL['nRow'] = numpy.uint32(self.st['pELL'].nRow)
    self.pELL['prodJd'] = numpy.uint32(self.st['pELL'].prodJd)
    self.pELL['sumJd'] = numpy.uint32(self.st['pELL'].sumJd)
    self.pELL['dim'] = numpy.uint32(self.st['pELL'].dim)
    self.pELL['Jd'] = self.thr.to_device(
        self.st['pELL'].Jd.astype(numpy.uint32))
    self.pELL['meshindex'] = self.thr.to_device(
        self.st['pELL'].meshindex.astype(numpy.uint32))
    self.pELL['kindx'] = self.thr.to_device(
        self.st['pELL'].kindx.astype(numpy.uint32))
    self.pELL['udata'] = self.thr.to_device(
        self.st['pELL'].udata.astype(self.dtype))

    self.volume = {}

    self.volume['Nd_elements'] = self.thr.to_device(
        numpy.asarray(self.Nd_elements, dtype=numpy.uint32))
    self.volume['Kd_elements'] = self.thr.to_device(
        numpy.asarray(self.Kd_elements, dtype=numpy.uint32))
    self.volume['invNd_elements'] = self.thr.to_device(
        self.invNd_elements.astype(numpy.float32))
    self.volume['Nd'] = self.thr.to_device(numpy.asarray(
        self.st['Nd'], dtype=numpy.uint32))
#     self.volume['gpu_coil_profile'] = self.thr.array(
#         self.multi_Nd, dtype=self.dtype).fill(1.0)

    Nd = self.st['Nd']

    self.tSN = {}
    self.tSN['Td_elements'] = self.thr.to_device(
        numpy.asarray(self.st['tSN'].Td_elements, dtype=numpy.uint32))
    self.tSN['invTd_elements'] = self.thr.to_device(
        self.st['tSN'].invTd_elements.astype(numpy.float32))
    self.tSN['Td'] = self.thr.to_device(
        numpy.asarray(self.st['tSN'].Td, dtype=numpy.uint32))
    self.tSN['Tdims'] = self.st['tSN'].Tdims
    self.tSN['tensor_sn'] = self.thr.to_device(
        self.st['tSN'].tensor_sn.astype(numpy.float32))

    self.Ndprod = numpy.int32(numpy.prod(self.st['Nd']))
    self.Kdprod = numpy.int32(numpy.prod(self.st['Kd']))
    self.M = numpy.int32(self.st['M'])

    import reikna.fft
    self.fft = reikna.fft.FFT(
            numpy.empty(self.Kd, dtype=self.dtype),
            self.ft_axes).compile(self.thr, fast_math=False)

    self.zero_scalar = self.dtype(0.0+0.0j)
    del self.st['pELL']
    if self.verbosity > 0:
        print('End of offload')

@push_cuda_context
def _set_wavefront_device(self, wf):
#         try:
    self.wavefront = int(wf)#api.DeviceParameters(device).warp_size
    if self.verbosity > 0:
        print('Wavefront of OpenCL (as wrap of CUDA) = ', self.wavefront)

    from ..src import re_subroutine  # import create_kernel_sets
    kernel_sets = re_subroutine.create_kernel_sets(self.API)

    prg = self.thr.compile(kernel_sets,
                           render_kwds=dict(LL=str(self.wavefront)),
                           fast_math=False)
    self.prg = prg

# @push_cuda_context
# def _reset_sense_device_deprecated(self):
#     self.volume['gpu_coil_profile'].fill(1.0)
# 
# @push_cuda_context
# def _set_sense_device_deprecated(self, coil_profile_device):
#     self.volume['gpu_coil_profile'] = coil_profile_device


    # if coil_profile.shape == self.Nd + (self.batch, ):

@push_cuda_context
def to_device(self, x, shape=None):
    gx = self.thr.to_device(x.copy().astype(self.dtype))
#         g_image = self.thr.array(image.shape, dtype=self.dtype)
#         self.thr.to_device(image.astype(self.dtype), dest=g_image)
    return gx
 
@push_cuda_context
def to_host(self, data):
    return data.get() 

# @push_cuda_context
# def _s2x_device(self, s):
#     x = self.thr.array(self.multi_Nd, dtype=self.dtype)
# 
#     self.prg.cPopulate(
#         self.batch,
#         self.Ndprod,
#         s,
#         x,
#         local_size=None,
#         global_size=int(self.batch * self.Ndprod))
# 
#     self.prg.cMultiplyVecInplace(
#         numpy.uint32(1),
#         self.volume['gpu_coil_profile'],
#         x,
#         local_size=None,
#         global_size=int(self.batch*self.Ndprod))
# 
#     return x    

@push_cuda_context
def _x2xx_device(self, x):

    xx = self.thr.array(x.shape, dtype=self.dtype)
    self.thr.copy_array(x, dest=xx, )

    self.prg.cTensorMultiply(numpy.uint32(self.batch),
                             numpy.uint32(self.tSN['Tdims']),
                             self.tSN['Td'],
                             self.tSN['Td_elements'],
                             self.tSN['invTd_elements'],
                             self.tSN['tensor_sn'],
                             xx,
                             numpy.uint32(0),
                             local_size=None,
                             global_size=int(self.batch*self.Ndprod))
    # self.thr.synchronize()
    return xx    

@push_cuda_context
def _xx2k_device(self, xx):
    """
    Private: oversampled FFT on the heterogeneous device

    First, zeroing the self.k_Kd array
    Second, copy self.x_Nd array to self.k_Kd array by cSelect
    Third, inplace FFT
    """
    k = self.thr.array(self.Kd, dtype=self.dtype)
    
    k.fill(0)

    self.prg.cTensorCopy(
         self.batch,
         numpy.uint32(self.ndims),
         self.volume['Nd_elements'],
         self.volume['Kd_elements'],
         self.volume['invNd_elements'],
         xx,
         k,
         numpy.int32(1),  # Directions: Nd -> Kd, 1; Kd -> Nd, -1
         local_size=None,
         global_size=int(self.Ndprod))
    self.fft(k, k, inverse=False)
#         self.thr.synchronize()
    return k    

@push_cuda_context
def _k2y_device(self, k):
    """
    Private: interpolation by the Sparse Matrix-Vector Multiplication
    """

    y = self.thr.array(self.M, dtype=self.dtype).fill(0)
    self.prg.pELL_spmv_mCoil(
                        self.batch,
                        self.pELL['nRow'],
                        self.pELL['prodJd'],
                        self.pELL['sumJd'],
                        self.pELL['dim'],
                        self.pELL['Jd'],
                        # self.pELL_currsumJd,
                        self.pELL['meshindex'],
                        self.pELL['kindx'],
                        self.pELL['udata'],
                        k,
                        y,
                        local_size=int(self.wavefront),
                        global_size=int(self.pELL['nRow'] *
                                        self.batch * self.wavefront)
                        )
#         self.thr.synchronize()
    return y    

@push_cuda_context
def _y2k_device(self, y):
    """
    Private: gridding by the Sparse Matrix-Vector Multiplication
    Atomic_twosum together provide better accuracy than generic atomic_add. 
    See: ocl_add and cuda_add code-strings in atomic_add(), inside the re_subroutine.py. 
    
    """

#         kx = self.thr.array(self.multi_Kd, dtype=numpy.float32).fill(0.0)
#         ky = self.thr.array(self.multi_Kd, dtype=numpy.float32).fill(0.0)
    k = self.thr.array(self.Kd, dtype=numpy.complex64).fill(0.0)
#     res = self.thr.array(self.Kd, dtype=numpy.complex64).fill(0.0)
    # array which saves the residue of two sum
    
    self.prg.pELL_spmvh_mCoil(
                        self.batch,
                        self.pELL['nRow'],
                        self.pELL['prodJd'],
                        self.pELL['sumJd'],
                        self.pELL['dim'],
                        self.pELL['Jd'],
                        self.pELL['meshindex'],
                        self.pELL['kindx'],
                        self.pELL['udata'],
#                             kx, ky,
                        k,
#                         res,
                        y,
                        local_size=int(self.wavefront),
                        global_size=int(self.pELL['nRow'] *self.wavefront)#*
#                                             int(self.pELL['prodJd']) * int(self.batch))
                        )

    return k# + res    


@push_cuda_context
def _k2xx_device(self, k):
    """
    Private: the inverse FFT and image cropping (which is the reverse of
    _xx2k() method)
    """

    self.fft(k, k, inverse=True)

    xx = self.thr.array(self.Nd, dtype=self.dtype)
    xx.fill(0)

    self.prg.cTensorCopy(
                         self.batch,
                         numpy.uint32(self.ndims),
                         self.volume['Nd_elements'],
                         self.volume['Kd_elements'],
                         self.volume['invNd_elements'],
                         k,
                         xx,
                         numpy.int32(-1),
                         local_size=None,
                         global_size=int(self.Ndprod))
    return xx    


@push_cuda_context
def _xx2x_device(self, xx):
    x = self.thr.array(xx.shape, dtype=self.dtype)
    self.thr.copy_array(xx, dest=x, )

    self.prg.cTensorMultiply(numpy.uint32(self.batch),
                             numpy.uint32(self.tSN['Tdims']),
                             self.tSN['Td'],
                             self.tSN['Td_elements'],
                             self.tSN['invTd_elements'],
                             self.tSN['tensor_sn'],
                             x,
                             numpy.uint32(0),
                             local_size=None,
                             global_size=int(self.batch *
                                             self.Ndprod))
    
    return x    

# @push_cuda_context
# def _x2s_device(self, x):
#     s = self.thr.array(self.st['Nd'], dtype=self.dtype)
# 
#     self.prg.cMultiplyConjVecInplace(
#         numpy.uint32(1),
#         self.volume['gpu_coil_profile'],
#         x,
#         local_size=None,
#         global_size=int(self.batch*self.Ndprod))
# 
#     self.prg.cAggregate(
#         self.batch,
#         self.Ndprod,
#         x,
#         s,
#         local_size=int(self.wavefront),
#         global_size=int(self.batch*self.Ndprod*self.wavefront))
# 
#     return s    

# @push_cuda_context
# def _selfadjoint_one2many2one_device_deprecated(self, gx):
#     """
#     selfadjoint_one2many2one NUFFT on the heterogeneous device
# 
#     :param gx: The input gpu array, with size=Nd
#     :type gx: reikna gpu array with dtype =numpy.complex64
#     :return: gx: The output gpu array, with size=Nd
#     :rtype: reikna gpu array with dtype =numpy.complex64
#     """
# 
#     gy = self._forward_one2many_device(gx)
#     gx2 = self._adjoint_many2one_device(gy)
#     del gy
#     return gx2    
# 
# @push_cuda_context
# def _selfadjoint_one2many2one_legacy_deprecated(self, gx):
#     """
#     selfadjoint_one2many2one NUFFT (Teplitz) on the heterogeneous device
# 
#     :param gx: The input gpu array, with size=Nd
#     :type gx: reikna gpu array with dtype =numpy.complex64
#     :return: gx: The output gpu array, with size=Nd
#     :rtype: reikna gpu array with dtype =numpy.complex64
#     """
# 
#     gy = self._forward_one2many_legacy(gx)
#     gx2 = self._adjoint_many2one_legacy(gy)
#     del gy
#     return gx2    

@push_cuda_context
def _selfadjoint_device(self, gx):
    """
    selfadjoint NUFFT on the heterogeneous device

    :param gx: The input gpu array, with size=Nd
    :type gx: reikna gpu array with dtype =numpy.complex64
    :return: gx: The output gpu array, with size=Nd
    :rtype: reikna gpu array with dtype =numpy.complex64
    """

    gy = self._forward_device(gx)
    gx2 = self._adjoint_device(gy)
    del gy
    return gx2    

@push_cuda_context
def _selfadjoint_legacy(self, gx):
    """
    selfadjoint NUFFT on the heterogeneous device

    :param gx: The input gpu array, with size=Nd
    :type gx: reikna gpu array with dtype =numpy.complex64
    :return: gx: The output gpu array, with size=Nd
    :rtype: reikna gpu array with dtype =numpy.complex64
    """

    gy = self._forward_legacy(gx)
    gx2 = self._adjoint_legacy(gy)
    del gy
    return gx2    

@push_cuda_context
def _forward_device(self, gx):
    """
    Forward NUFFT on the heterogeneous device

    :param gx: The input gpu array, with size = Nd
    :type gx: reikna gpu array with dtype = numpy.complex64
    :return: gy: The output gpu array, with size = (M,)
    :rtype: reikna gpu array with dtype = numpy.complex64
    """
    xx = self._x2xx_device(gx)


    k = self._xx2k_device(xx)
    del xx
    gy = self._k2y_device(k)
    del k
    return gy    

@push_cuda_context
def _forward_legacy(self, gx):
    """
    Forward NUFFT on the heterogeneous device

    :param gx: The input gpu array, with size = Nd
    :type gx: reikna gpu array with dtype = numpy.complex64
    :return: gy: The output gpu array, with size = (M,)
    :rtype: reikna gpu array with dtype = numpy.complex64
    """
    xx = self._x2xx_device(gx)


    k = self._xx2k_device(xx)
    del xx
    gy = self._k2y_legacy(k)
    del k
    return gy    

@push_cuda_context
def _forward_one2many_device_deprecated(self, s):
    x = self._s2x_device(s)
    y = self._forward_device(x)
    return y    

@push_cuda_context
def _adjoint_many2one_device_deprecated(self, y):
    x = self._adjoint_device(y)
    s = self._x2s_device(x)
    return s    

@push_cuda_context
def _forward_one2many_legacy_deprecated(self, s):
    x = self._s2x_device(s)
    y = self._forward_legacy(x)
    return y    

@push_cuda_context
def _adjoint_many2one_legacy_deprecated(self, y):
    x = self._adjoint_legacy(y)
    s = self._x2s_device(x)
    return s    

@push_cuda_context
def _adjoint_device(self, gy):
    """
    Adjoint NUFFT on the heterogeneous device

    :param gy: The input gpu array, with size=(M,)
    :type: reikna gpu array with dtype =numpy.complex64
    :return: gx: The output gpu array, with size=Nd
    :rtype: reikna gpu array with dtype =numpy.complex64
    """
    k = self._y2k_device(gy)

    xx = self._k2xx_device(k)
    del k
    gx = self._xx2x_device(xx)
    del xx
    return gx   

@push_cuda_context
def _adjoint_legacy(self, gy):
    """
    Adjoint NUFFT on the heterogeneous device

    :param gy: The input gpu array, with size=(M,)
    :type: reikna gpu array with dtype =numpy.complex64
    :return: gx: The output gpu array, with size=Nd
    :rtype: reikna gpu array with dtype =numpy.complex64
    """
    k = self._y2k_legacy(gy)

    xx = self._k2xx_device(k)
    del k
    gx = self._xx2x_device(xx)
    del xx
    return gx   

@push_cuda_context
def release(self):
    try:
        del self.volume
    except:
        pass
    try:
        del self.tSN
    except:
        pass
    try:
        del self.prg
    except:
        pass
    try:
        del self.pELL
    except:
        pass
    try:
        del self.csr 
        del self.csrH
    except:
        pass
    try:
        self.thr.release()
    except:
        pass
    try:
        del self.thr     
    except:
        pass
    
@push_cuda_context
def _solve_device(self, gy, solver=None, *args, **kwargs):
    """
    The solver of NUFFT

    :param gy: data, reikna array, (M,) size
    :param solver: could be 'cg', 'L1TVOLS', 'L1TVLAD'
    :param maxiter: the number of iterations
    :type gy: reikna array, dtype = numpy.complex64
    :type solver: string
    :type maxiter: int
    :return: reikna array with size Nd
    """
    from ..linalg.solve_device import solve
    return solve(self,  gy,  solver, *args, **kwargs)

@push_cuda_context
def _solve_legacy(self, gy, solver=None, *args, **kwargs):
    """
    The solver of NUFFT

    :param gy: data, reikna array, (M,) size
    :param solver: could be 'cg', 'L1TVOLS', 'L1TVLAD'
    :param maxiter: the number of iterations
    :type gy: reikna array, dtype = numpy.complex64
    :type solver: string
    :type maxiter: int
    :return: reikna array with size Nd
    """
    from ..linalg.solve_legacy import solve as solve2
    return solve2(self,  gy,  solver, *args, **kwargs)

def _plan_legacy(self, om, Nd, Kd, Jd, ft_axes = None):
    """
    Design the min-max interpolator.
     
    :param om: The M off-grid locations in the frequency domain. Normalized between [-pi, pi]
    :param Nd: The matrix size of 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
    :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. 
    :returns: 0
    :rtype: int, float
    :Example:

    >>> import pynufft
    >>> NufftObj = pynufft.NUFFT_cpu()
    >>> NufftObj.plan(om, Nd, Kd, Jd) 
     
    """         
    self.ndims = len(Nd)  # dimension
    if ft_axes is None:
        ft_axes = range(0, self.ndims)
    self.ft_axes = ft_axes

    self.st = helper.plan(om, Nd, Kd, Jd, ft_axes=ft_axes,
                          format='CSR')
#     if batch is None:
    self.parallel_flag = 0
#     else:
#         self.parallel_flag = 1

#     if batch is None:
    self.batch = numpy.uint32(1)

#     else:
#         self.batch = numpy.uint32(batch)

    self.Nd = self.st['Nd']  # backup
    self.Kd = self.st['Kd']
    #  self.sn = numpy.asarray(self.st['sn'].astype(self.dtype),
    #                            order='C')# backup
#     if self.batch == 1:
#         self.multi_Nd = self.Nd
#         self.multi_Kd = self.Kd
#         self.multi_M = (self.st['M'], )
        # Broadcasting the sense and scaling factor (Roll-off)
        # self.sense2 = self.sense*numpy.reshape(self.sn, self.Nd + (1, ))
#     else:  # self.batch is 0:
#         self.multi_Nd = self.Nd + (self.batch, )
#         self.multi_Kd = self.Kd + (self.batch, )
#         self.multi_M = (self.st['M'], ) + (self.batch, )

    self.Kdprod = numpy.uint32(numpy.prod(self.st['Kd']))
    self.Jdprod = numpy.uint32(numpy.prod(self.st['Jd']))
    self.Ndprod = numpy.uint32(numpy.prod(self.st['Nd']))

    self.Nd_elements, self.invNd_elements = helper.strides_divide_itemsize(
                                                self.st['Nd'])
    # only return the Kd_elements
    self.Kd_elements = helper.strides_divide_itemsize(self.st['Kd'])[0]
    
    self.sp = self.st['p'].copy().tocsr()
    self.spH = (self.st['p'].getH().copy()).tocsr()        
    
    self._offload_legacy()

    
     
    return 0

@push_cuda_context
def _offload_legacy(self):
    """
    self.offload():
     
    Off-load NUFFT to the opencl or cuda device(s)
     
    :param API: define the device type, which can be 'cuda' or 'ocl'
    :param platform_number: define which platform to be used. The default platform_number = 0.
    :param device_number: define which device to be used. The default device_number = 0.
    :type API: string
    :type platform_number: int
    :type device_number: int
    :return: self: instance

    """


#     self.pELL = {}  # dictionary
# 
#     self.pELL['nRow'] = numpy.uint32(self.st['pELL'].nRow)
#     self.pELL['prodJd'] = numpy.uint32(self.st['pELL'].prodJd)
#     self.pELL['sumJd'] = numpy.uint32(self.st['pELL'].sumJd)
#     self.pELL['dim'] = numpy.uint32(self.st['pELL'].dim)
#     self.pELL['Jd'] = self.thr.to_device(
#         self.st['pELL'].Jd.astype(numpy.uint32))
#     self.pELL['meshindex'] = self.thr.to_device(
#         self.st['pELL'].meshindex.astype(numpy.uint32))
#     self.pELL['kindx'] = self.thr.to_device(
#         self.st['pELL'].kindx.astype(numpy.uint32))
#     self.pELL['udata'] = self.thr.to_device(
#         self.st['pELL'].udata.astype(self.dtype))

    self.volume = {}

    self.volume['Nd_elements'] = self.thr.to_device(
        numpy.asarray(self.Nd_elements, dtype=numpy.uint32))
    self.volume['Kd_elements'] = self.thr.to_device(
        numpy.asarray(self.Kd_elements, dtype=numpy.uint32))
    self.volume['invNd_elements'] = self.thr.to_device(
        self.invNd_elements.astype(numpy.float32))
    self.volume['Nd'] = self.thr.to_device(numpy.asarray(
        self.st['Nd'], dtype=numpy.uint32))
#     self.volume['gpu_coil_profile'] = self.thr.array(
#         self.multi_Nd, dtype=self.dtype).fill(1.0)

    Nd = self.st['Nd']

    self.tSN = {}
    self.tSN['Td_elements'] = self.thr.to_device(
        numpy.asarray(self.st['tSN'].Td_elements, dtype=numpy.uint32))
    self.tSN['invTd_elements'] = self.thr.to_device(
        self.st['tSN'].invTd_elements.astype(numpy.float32))
    self.tSN['Td'] = self.thr.to_device(
        numpy.asarray(self.st['tSN'].Td, dtype=numpy.uint32))
    self.tSN['Tdims'] = self.st['tSN'].Tdims
    self.tSN['tensor_sn'] = self.thr.to_device(
        self.st['tSN'].tensor_sn.astype(numpy.float32))

    self.Ndprod = numpy.int32(numpy.prod(self.st['Nd']))
    self.Kdprod = numpy.int32(numpy.prod(self.st['Kd']))
    self.M = numpy.int32(self.st['M'])

    import reikna.fft
    self.fft = reikna.fft.FFT(
            numpy.empty(self.Kd, dtype=self.dtype),
            self.ft_axes).compile(self.thr, fast_math=False)

    self.zero_scalar = self.dtype(0.0+0.0j)
#     del self.st['pELL']
#     if self.verbosity > 0:
#         print('End of offload')
     
#         self.SnGPUArray = self.thr.to_device(  self.sn)
    self.csr = {}
    self.csrH = {}
    self.csr['data'] = self.thr.to_device( self.sp.data.astype(self.dtype))
    self.csr['indices'] = self.thr.to_device( self.sp.indices.astype(numpy.uint32))
    self.csr['indptr'] =  self.thr.to_device( self.sp.indptr.astype(numpy.uint32))
    self.csr['numrow'] = self.M
    self.csr['numcol'] = self.Kdprod
     

    del self.sp
    
    self.csrH['data'] = self.thr.to_device(  self.spH.data.astype(self.dtype))
    self.csrH['indices'] =  self.thr.to_device(  self.spH.indices.astype(numpy.uint32))
    self.csrH['indptr'] =  self.thr.to_device(  self.spH.indptr.astype(numpy.uint32))
    self.csrH['numrow'] = self.Kdprod
     
    del self.spH
    
#     import reikna.fft
# 
#     self.fft = reikna.fft.FFT(numpy.empty(self.st['Kd'], dtype=self.dtype), self.ft_axes).compile(self.thr, fast_math=False)
# 
#     self.zero_scalar=self.dtype(0.0+0.0j)
     

@push_cuda_context  
def _k2y_legacy(self, k):
    """
    Private: interpolation by the Sparse Matrix-Vector Multiplication
    """
    y =self.thr.array( self.st['M'], dtype=self.dtype).fill(0)
    self.prg.cCSR_spmv_vector(    
                     self.batch,                            
                       self.csr['numrow'], 
                       self.csr['indptr'],
                       self.csr['indices'],
                       self.csr['data'], 
                       k,
                       y,
                       local_size=int(self.wavefront),
                       global_size=int(self.csr['numrow']*self.wavefront*self.batch) 
                        )

#     self.thr.synchronize()
    return y    
@push_cuda_context
def _y2k_legacy(self, y):
    """
    Private: gridding by the Sparse Matrix-Vector Multiplication
    """
    k = self.thr.array(self.Kd, dtype = self.dtype)

    self.prg.cCSR_spmv_vector(
                        self.batch,
                       self.csrH['numrow'], 
                       self.csrH['indptr'],
                       self.csrH['indices'],
                       self.csrH['data'], 
                       y,
                       k,
                       local_size=int(self.wavefront),
                       global_size=int(self.csrH['numrow']*self.wavefront) 
                        )#,g_times_l=int(csrnumrow))

    self.thr.synchronize()
    return k