import numpy as np
import pywt


class CwtTimeSuiteBase(object):
    """
    Set-up for CWT timing.
    """
    params = ([32, 128, 512, 2048],
              ['cmor', 'cgau4', 'fbsp', 'gaus4', 'mexh', 'morl', 'shan'],
              [16, 64, 256],
              ['conv', 'fft'])
    param_names = ('n', 'wavelet', 'max_scale', 'method')

    def setup(self, n, wavelet, max_scale, method):
        try:
            from pywt import cwt
        except ImportError:
            raise NotImplementedError("cwt not available")
        self.data = np.ones(n, dtype='float')
        self.scales = np.arange(1, max_scale+1)


class CwtTimeSuite(CwtTimeSuiteBase):
    def time_cwt(self, n, wavelet, max_scale, method):
        try:
            pywt.cwt(self.data, self.scales, wavelet, method=method)
        except TypeError:
            # older PyWavelets does not support use of the method argument
            if method == 'fft':
                raise NotImplementedError(
                    "fft-based convolution not available.")
            pywt.cwt(self.data, self.scales, wavelet)

@@ -259,22 +259,78 @@
    %5B
-'conv', 'fft'%5D
+np.float32, np.float64%5D,%0A              %5B'conv', 'fft'%5D,%0A              
 )%0A  
@@ -374,16 +374,25 @@
 _scale',
+ 'dtype',
  'method
@@ -429,32 +429,39 @@
 elet, max_scale,
+ dtype,
  method):%0A      
@@ -628,15 +628,13 @@
 ype=
-'float'
+dtype
 )%0A  
@@ -679,9 +679,11 @@
 cale
-+
+ + 
 1)%0A%0A
@@ -766,16 +766,23 @@
 x_scale,
+ dtype,
  method)

add float32 cases to cwt benchmarks