# Copyright (c) 2016, Imperial College London
# Copyright (c) 2016, Ghislain Antony Vaillant
# All rights reserved.
#
# Distributed under the terms of the new BSD license.
# See the accompanying LICENSE file or read the terms at
# https://opensource.org/licenses/BSD-3-Clause.

from ._nfft import ffi, lib
from functools import reduce
import numpy

__all__ = ("Plan")


class Plan(object):
  
    """The NFFT plan class."""
    
    def __init__(self, N, M, *args, **kwargs):
        "Instantiate the NFFT plan."        
        d = len(N)
        n = [2 * Nt for Nt in N]
        # TODO: auto-adjust with desired dtype.
        m = 6
        # TODO: auto-adjust with desired precomputation.
        nfft_flags = (lib.PRE_PHI_HUT, lib.PRE_PSI, lib.FFT_OUT_OF_PLACE,
                      lib.FFTW_INIT)
        fftw_flags = (lib.FFTW_ESTIMATE, lib.FFTW_DESTROY_INPUT)
        # Create plan handle.
        plan = ffi.new("nfft_plan *")
        lib.nfft_init_guru(
            plan,
            d,
            ffi.new("const int []", N),
            M,
            ffi.new("const int []", n),
            m,
            reduce(lambda x, y: x | y, nfft_flags, 0),
            reduce(lambda x, y: x | y, fftw_flags, 0),
        )
        # Create and bind f_hat array.
        f_hat = numpy.empty(N, dtype=numpy.complex128)
        plan.f_hat = ffi.cast("fftw_complex *", f_hat.ctypes.data)       
        # Create and bind f array.
        f = numpy.empty(M, dtype=numpy.complex128)
        plan.f = ffi.cast("fftw_complex *", f.ctypes.data)
        # Create and bind x array.
        if d == 1:
            x = numpy.empty(M, dtype=numpy.float64)
        else:
            x = numpy.empty([M, d], dtype=numpy.float64)
        plan.x = ffi.cast("double *", x.ctypes.data)
        # Hold plan handle and interface arrays together.
        self.__handle = plan
        self.f_hat = f_hat
        self.f = f
        self.x = x
        # TODO: support precompute on instantiation.

    def __del__(self):
        lib.nfft_finalize(self.__handle)

    def forward(self, direct=False):
        """Compute and return the forward transform."""
        if direct:
            self.execute_forward_direct()
        else:
            self.execute_forward()
        return self.f

    def execute_forward(self):
        """Perform the foward transform (fast)."""
        lib.nfft_trafo(self.__handle)

    def execute_forward_direct(self):
        """Perform the forward transform (direct)."""
        lib.nfft_trafo_direct(self.__handle)

    def adjoint(self, direct=False):
        """Compute and return the adjoint transform."""
        if direct:
            self.execute_adjoint_direct()
        else:
            self.execute_adjoint()
        return self.f_hat

    def execute_adjoint(self):
        """Perform the adjoint transform (fast)."""
        lib.nfft_adjoint(self.__handle)

    def execute_adjoint_direct(self):
        """Compute the adjoint transform (direct)."""
        lib.nfft_adjoint_direct(self.__handle)

    def precompute(self):
        "Precompute the plan."
        lib.nfft_precompute_one_psi(self.__handle)
@@ -899,20 +899,22 @@
         
-plan
+handle
  = ffi.n
@@ -962,53 +962,18 @@
 uru(
-%0A            plan,%0A            d,%0A           
+handle, d,
  ffi
@@ -996,28 +996,16 @@
 %5B%5D%22, N),
-%0A           
  M,%0A    
@@ -1039,28 +1039,16 @@
 %5B%5D%22, n),
-%0A           
  m,%0A    
@@ -1151,26 +1151,16 @@
 lags, 0)
-,%0A        
 )%0A      
@@ -1255,20 +1255,22 @@
         
-plan
+handle
 .f_hat =
@@ -1315,23 +1315,16 @@
 es.data)
-       
 %0A       
@@ -1410,20 +1410,22 @@
         
-plan
+handle
 .f = ffi
@@ -1648,20 +1648,22 @@
         
-plan
+handle
 .x = ffi
@@ -1777,20 +1777,22 @@
 andle = 
-plan
+handle
 %0A       

Use handle attribute consistently.