import pyfftw
import numpy as np
import tomviz.operators

class ReconDFMOperator(tomviz.operators.CancelableOperator):

    def transform_scalars(self, dataset):
        """3D Reconstruct from a tilt series using Direct Fourier Method"""

        from tomviz import utils
        import numpy as np

        # Get Tilt angles
        tiltAngles = utils.get_tilt_angles(dataset)

        tiltSeries = utils.get_array(dataset)
        if tiltSeries is None:
            raise RuntimeError("No scalars found!")

        Npad = np.size(tiltSeries, 0) * 2

        tiltSeries = np.double(tiltSeries)
        (Nx, Ny, Nproj) = tiltSeries.shape
        tiltAngles = np.double(tiltAngles)
        pad_pre = int(np.ceil((Npad - Ny) / 2.0))
        pad_post = int(np.floor((Npad - Ny) / 2.0))

        # Initialization
        Nz = Ny
        w = np.zeros((Nx, Ny, np.int(Nz / 2 + 1))) #store weighting factors
        v = pyfftw.n_byte_align_empty((Nx, Ny, Nz / 2 + 1), 16, dtype='complex128')
        v = np.zeros(v.shape) + 1j * np.zeros(v.shape)
        recon = pyfftw.n_byte_align_empty((Nx, Ny, Nz), 16, dtype='float64')
        recon_fftw_object = pyfftw.FFTW(
                v, recon, direction='FFTW_BACKWARD', axes=(0, 1, 2))

        p = pyfftw.n_byte_align_empty((Nx, Npad), 16, dtype='float64')
        pF = pyfftw.n_byte_align_empty((Nx, Npad / 2 + 1), 16, dtype='complex128')
        p_fftw_object = pyfftw.FFTW(p, pF, axes=(0, 1))

        dk = np.double(Ny) / np.double(Npad)
        NUMBER_OF_CHUNKS = int(10)
        self.progress.maximum = NUMBER_OF_CHUNKS
        step = 0
        for a in range(Nproj):
            if self.canceled:
                return

            #print angles[a]
            ang = tiltAngles[a] * np.pi / 180
            projection = tiltSeries[:, :, a] #2D projection image
            p = np.lib.pad(projection, ((0, 0), (pad_pre, pad_post)),
                           'constant', constant_values=(0, 0)) #pad zeros
            p = np.fft.ifftshift(p)
            p_fftw_object.update_arrays(p, pF)
            p_fftw_object()

            probjection_f = pF.copy()
            if ang < 0:
                probjection_f = np.conj(pF.copy())
                probjection_f[1:, :] = np.flipud(probjection_f[1:, :])
                ang = np.pi + ang

            # Bilinear extrapolation
            for i in range(0, np.int(np.ceil(Npad / 2)) + 1):
                ky = i * dk
                #kz = 0
                ky_new = np.cos(ang) * ky #new coord. after rotation
                kz_new = np.sin(ang) * ky
                sy = abs(np.floor(ky_new) - ky_new) #calculate weights
                sz = abs(np.floor(kz_new) - kz_new)
                for b in range(1, 5): #bilinear extrapolation
                    pz, py, weight = bilinear(kz_new, ky_new, sz, sy, Ny, b)
                    if (py >= 0 and py < Ny and pz >= 0 and pz < Nz / 2 + 1):
                        w[:, py, pz] = w[:, py, pz] + weight
                        v[:, py, pz] = v[:, py, pz] + weight * probjection_f[:, i]
            step += 1
            self.progress.update(step)

        v[w != 0] = v[w != 0] / w[w != 0]
        recon_fftw_object.update_arrays(v, recon)
        recon_fftw_object()
        recon = np.fft.fftshift(recon)
        step += 1
        self.progress.update(step)
        
        # Set the result as the new scalars.
        utils.set_array(dataset, recon)
    
        # Mark dataset as volume
        utils.mark_as_volume(dataset)


# Bilinear extrapolation


def bilinear(kz_new, ky_new, sz, sy, N, p):
    if p == 1:
        py = np.floor(ky_new)
        pz = np.floor(kz_new)
        weight = (1 - sy) * (1 - sz)
    elif p == 2:
        py = np.ceil(ky_new)
        pz = np.floor(kz_new)
        weight = sy * (1 - sz)
    elif p == 3:
        py = np.floor(ky_new)
        pz = np.ceil(kz_new)
        weight = (1 - sy) * sz
    elif p == 4:
        py = np.ceil(ky_new)
        pz = np.ceil(kz_new)
        weight = sy * sz
    if py < 0:
        py = N + py
    else:
        py = py
    return (int(pz), int(py), weight)

@@ -51,16 +51,17 @@
 rators%0A%0A
+%0A
 class Re
@@ -925,32 +925,45 @@
 yte_align_empty(
+%0A            
 (Nx, Ny, Nz / 2 
@@ -1177,20 +1177,16 @@
         
-    
 v, recon
@@ -1333,32 +1333,45 @@
 yte_align_empty(
+%0A            
 (Nx, Npad / 2 + 
@@ -1504,42 +1504,8 @@
 ad)%0A
-        NUMBER_OF_CHUNKS = int(10)
 %0A   
@@ -1538,23 +1538,12 @@
  = N
-UMBER_OF_CHUNKS
+proj
 %0A   
@@ -1552,24 +1552,44 @@
    step = 0%0A
+        print step%0A%0A
         for 
@@ -1659,16 +1659,77 @@
   return
+%0A            step += 1%0A            self.progress.update(step)
 %0A%0A      
@@ -3059,37 +3059,12 @@
 %5D + 
-weight * probjection_f%5B:, i%5D%0A
+%5C%0A  
     
@@ -3067,34 +3067,25 @@
             
-step += 1%0A
+ 
             
@@ -3088,34 +3088,37 @@
     
-self.progress.update(step)
+ weight * probjection_f%5B:, i%5D
 %0A%0A  
@@ -3278,69 +3278,8 @@
 on)%0A
-        step += 1%0A        self.progress.update(step)%0A        
 %0A   
@@ -3360,20 +3360,16 @@
  recon)%0A
-    
 %0A       
@@ -3459,18 +3459,16 @@
 olation%0A
-%0A%0A
 def bili

fix formatting