"""
NUFFT class
=======================================
"""
from __future__ import absolute_import
import numpy
import warnings
import scipy.sparse
import numpy.fft
import scipy.signal
import scipy.linalg
import scipy.special
from functools import wraps as _wraps
# from ..linalg.nufft_cpu import NUFFT_cpu
# from ..linalg.nufft_hsa import NUFFT_hsa





def push_cuda_context(hsa_method):
    """
    Decorator: Push cude context to the top of the stack for current use
    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



class NUFFT:
    """
    NUFFT class
    =======================================
    A super class of cpu and gpu NUFFT functions. 
    
    Note: NUFFT does not inherit NUFFT_cpu (deprecated) and NUFFT_hsa (deprecated).
    """
    #import cpu codes
    from ._nufft_class_methods_cpu import _init__cpu, _plan_cpu, _precompute_sp_cpu, _solve_cpu, _forward_cpu, _adjoint_cpu, _selfadjoint_cpu, _selfadjoint2_cpu, _x2xx_cpu, _xx2k_cpu, _xx2k_one2one_cpu, _k2vec_cpu, _vec2y_cpu, _k2y_cpu, _y2vec_cpu, _vec2k_cpu, _y2k_cpu, _k2xx_cpu, _k2xx_one2one_cpu, _xx2x_cpu, _k2y2k_cpu
    # import host codes
    from ._nufft_class_methods_cpu import  _forward_host,  _adjoint_host, _selfadjoint_host,  _solve_host, _xx2k_host, _k2xx_host, _x2xx_host, _xx2x_host, _k2y_host,  _y2k_host
    # import device codes
    from ._nufft_class_methods_device import _init__device, _plan_device,  _set_wavefront_device, _offload_device, to_device, to_host, _x2xx_device, _xx2k_device, _k2y_device, _y2k_device, _k2xx_device, _xx2x_device,  _selfadjoint_device, _forward_device,  _adjoint_device, release, _solve_device, release
    
    # legacy codes (csr format for device)
    from ._nufft_class_methods_device import _y2k_legacy, _k2y_legacy, _forward_legacy, _adjoint_legacy, _selfadjoint_legacy, _plan_legacy, _offload_legacy, _solve_legacy
    from ._nufft_class_methods_cpu import _k2y_legacy_host,  _y2k_legacy_host, _selfadjoint_legacy_host, _forward_legacy_host, _adjoint_legacy_host,  _solve_legacy_host
    
    def __init__(self, device_indx=None, legacy=None):
        """
        Constructor.

        :param None:
        :type None: Python NoneType
        :return: NUFFT: the pynufft.NUFFT instance
        :rtype: NUFFT: the pynufft.NUFFT class
        :Example:

        >>> from pynufft import NUFFT
        >>> NufftObj = NUFFT()
        
        or 
        
        >>> from pynufft import NUFFT, helper
        >>> device = helper.device_list()[0]
        >>> NufftObj = NUFFT(device) # for first acceleration device in the system
        
        """
        if device_indx is None:
            self._init__cpu()
            self.processor = 'cpu'
        else:
            if legacy is True:
                self._init__device(device_indx)
                self.processor = 'hsa_legacy'
            else:
                self._init__device(device_indx)
                self.processor = 'hsa'
            
    def __del__(self):
        if self.processor is 'hsa' or 'hsa_legacy':
            self.release()
        else:
            pass
                
    def plan(self,  *args, **kwargs):
        """
        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.
        :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)

        """        
        func = {'cpu': self._plan_cpu,
                    'hsa': self._plan_device,
                    'hsa_legacy': self._plan_legacy}
        return func.get(self.processor)(*args, **kwargs)
    
    def forward(self, *args, **kwargs):
        """
        Forward NUFFT (host code)

        :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
        """        
        func = {'cpu': self._forward_cpu, 
                    'hsa': self._forward_host,
                    'hsa_legacy': self._forward_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def adjoint(self, *args, **kwargs):
        """
        Adjoint NUFFT (host code)

        :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
        """        
        func = {'cpu': self._adjoint_cpu,
                'hsa': self._adjoint_host,
                'hsa_legacy': self._adjoint_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def selfadjoint(self, *args, **kwargs):
        """
        selfadjoint NUFFT (host code)

        :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
        """
        func = {'cpu': self._selfadjoint_cpu,
                'hsa': self._selfadjoint_host,
                'hsa_legacy': self._selfadjoint_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def solve(self, *args, **kwargs):
        """
        Solve NUFFT (host code)
        :param y: data, numpy.complex64. The shape = (M,) 
        :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 Nd.
        """        
        func = {'cpu': self._solve_cpu,
                'hsa': self._solve_host,
                'hsa_legacy': self._solve_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def xx2k(self, *args, **kwargs):
        func = {'cpu': self._xx2k_cpu,
                'hsa': self._xx2k_host,
                'hsa_legacy': self._xx2k_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def k2xx(self, *args, **kwargs):
        func = {'cpu': self._k2xx_cpu,
                'hsa': self._k2xx_host,
                'hsa_legacy': self._k2xx_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def x2xx(self, *args, **kwargs):
        func = {'cpu': self._x2xx_cpu,
                'hsa': self._x2xx_host,
                'hsa_legacy': self._x2xx_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def xx2x(self, *args, **kwargs):
        func = {'cpu': self._xx2x_cpu,
                'hsa': self._xx2x_host,
                'hsa_legacy': self._xx2x_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def k2y(self, *args, **kwargs):
        func = {'cpu': self._k2y_cpu,
                'hsa': self._k2y_host,
                'hsa_legacy': self._k2y_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def y2k(self, *args, **kwargs):
        func = {'cpu': self._y2k_cpu,
                'hsa': self._y2k_host,
                'hsa_legacy': self._y2k_legacy_host}
        return func.get(self.processor)(*args, **kwargs)
    
    def k2yk2(self, *args, **kwargs):
        func = {'cpu': self._k2yk2_cpu,
                'hsa': self._k2yk2_host,
                'hsa_legacy': self._k2yk2_host}
        return func.get(self.processor)(*args, **kwargs)
    
#     def adjoint_many2one(self, *args, **kwargs):
#         func = {'cpu': self._adjoint_many2one_cpu,
#                 'hsa': self._adjoint_many2one_host,
#                 'hsa_legacy': self._adjoint_many2one_legacy_host}
#         return func.get(self.processor)(*args, **kwargs)
#     
#     def forward_one2many(self, *args, **kwargs):
#         func = {'cpu': self._forward_one2many_cpu,
#                 'hsa': self._forward_one2many_host,
#                 'hsa_legacy': self._forward_one2many_legacy_host}
#         return func.get(self.processor)(*args, **kwargs)
    
#     def selfadjoint_one2many2one(self, *args, **kwargs):
#         func = {'cpu': self._selfadjoint_one2many2one_cpu,
#                 'hsa': self._selfadjoint_one2many2one_host,
#                 'hsa_legacy': self._selfadjoint_one2many2one_legacy_host}
#         return func.get(self.processor)(*args, **kwargs)   
    
    def k2xx_one2one(self, *args, **kwargs):
        func = {'cpu':self._k2xx_one2one_cpu}
        return func.get(self.processor)(*args, **kwargs)
    
    def xx2k_one2one(self, *args, **kwargs):
        func = {'cpu':self._xx2k_one2one_cpu}
        return func.get(self.processor)(*args, **kwargs) 
    
    def k2y2k(self, *args, **kwargs):
        func = {'cpu': self._k2y2k_cpu}
        return func.get(self.processor)(*args, **kwargs)
#     def set_sense(self, *args, **kwargs):
#         func = {'cpu': self._set_sense_cpu,
#                 'hsa': self._set_sense_host,
#                 'hsa_legacy': self._set_sense_host}
#         return func.get(self.processor)(*args, **kwargs)
#     def reset_sense(self, *args, **kwargs):
#         func = {'cpu': self._reset_sense_cpu,
#                 'hsa': self._reset_sense_host,
#                 'hsa_legacy': self._reset_sense_host}
#         return func.get(self.processor)(*args, **kwargs)