
import numpy as np
from skimage.io import imshow, show
import matplotlib.pyplot as plt
#im Moment mit eta = 2*eta und gamma = 0.5*gamma

class gabortsfm:
    
    def __init__(self,  num_scales, num_rotations, filter_num_rows=28, filter_num_columns=28):
        self._filter_num_rows = filter_num_rows
        self._filter_num_columns = filter_num_columns
        self._num_scales = num_scales
        self._num_rotations = num_rotations
        self._gamma = 0.6713
        self._eta = 0.35
        self._sigma_phi = 0.7
        self._sigma_psi = 0.5158
        self._scaling_psi = 0.84
        self._scaling_overall = 1.0
        self._psi_filter_bank = self._generate_psi_filter_bank()
        self._phi_filter = self._generate_phi_filter()
        
        
        #uncomment to plot filters
        
        psifilters = [val for one_filter_bank in self._psi_filter_bank for val in one_filter_bank]
        phifilter = np.reshape(self._phi_filter,(len(self._phi_filter),len(self._phi_filter)))
        fig, axes = plt.subplots(num_scales, num_rotations, figsize=(10, 10), subplot_kw={'xticks': [], 'yticks': []})
        fig.subplots_adjust(hspace=0.3, wspace=0.05)
        for ax, s in zip(axes.flat, psifilters):
            im = ax.imshow(s, interpolation='nearest')
        plt.show()
        
    
    def apply_filter(self, img, meta=None):
        return _gabortsfm(img, meta, self._num_scales, self._num_rotations, self._phi_filter, self._psi_filter_bank)
    

    def _generate_psi_filter_bank(self): #for propagated images
        psi_filter_bank = [[0]*self._num_rotations]*self._num_scales
        for j in range(0, self._num_scales):
            a = [0]*self._num_rotations
            for ir in range(0, self._num_rotations):
                theta = ir*np.pi/self._num_rotations
                a[ir]=self._generate_psi_single_filter(j+1, theta)
            psi_filter_bank[j]=a
        return psi_filter_bank


    def _generate_psi_single_filter(self, j, theta):
        xi_x, xi_y = self._generate_grid(j, theta)
        for i in range(0, len(xi_x)):
            xi_x[i]=xi_x[i]-2*self._eta
        for i in range(0, len(xi_y)):
            xi_y[i]=xi_y[i]/self._gamma*2
        filter = np.exp(-2*((np.pi*self._sigma_psi)**2)*((xi_x**2)+(xi_y**2)))
        filter *= self._scaling_psi*self._scaling_overall*2*self._sigma_psi*np.sqrt(np.pi/self._gamma)/np.sqrt(1+np.exp(-4*(np.pi*self._sigma_psi*self._eta)**2)-np.exp(-6*(np.pi*self._sigma_psi*self._eta)**2))
        return filter
    

    def _generate_phi_filter(self): #for output image
        j = 1
        theta = 0
        xi_x, xi_y = self._generate_grid(j, theta)
        phi_filter = np.exp((-2*(np.pi*self._sigma_phi)**2)*(xi_x**2+xi_y**2))
        phi_filter *= self._scaling_overall*2*self._sigma_phi*np.sqrt(np.pi)
        return phi_filter

    
    def _generate_grid(self, j, theta):
        scale = 2**(j-1)
        x = self._filter_num_rows
        temp_x = list(range(int(np.floor(-self._filter_num_rows/2 + 0.5)),int(np.floor(self._filter_num_rows/2 + 0.5))))
        temp_y = list(range(int(np.floor(-self._filter_num_columns/2 + 0.5)),int(np.floor(self._filter_num_columns/2 + 0.5))))
        temp_x = temp_x/max(np.abs(temp_x))
        temp_y = temp_y/max(np.abs(temp_y))
        xi_x, xi_y = np.meshgrid(temp_x, temp_y)
        xi_x = xi_x/scale
        xi_y = xi_y/scale
        xi_x_temp = xi_x*np.cos(theta) + xi_y*np.sin(theta)
        xi_y_temp = -xi_x*np.sin(theta) +xi_y*np.cos(theta)
        xi_x = xi_x_temp
        xi_y = xi_y_temp
        return [xi_x, xi_y]

def _gabortsfm(img, meta, num_scales, num_rotations, phi_filter, psi_filter_bank):
    prop = []
    out = []
    inp_img = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(img)))
    if(phi_filter.size > inp_img.size):
        raise NameError('Wavelet_Transform::ApplyTransform: Filter size larger than input image size')
    oup_img = inp_img*phi_filter
    oup_img = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(oup_img)))
    oup_img = conv_comp_img_2_oup_img(oup_img)
    out.extend([oup_img])

    if meta is not None and 'scale' in meta:
        current_scale = min(meta['scale'], num_scales)
    else:
        current_scale = num_scales

    new_meta = []
    for i in range(0, current_scale):
        for j in range(0, num_rotations):
            if(psi_filter_bank[i][j].size > inp_img.size):
                raise NameError('Wavelet_Transform::ApplyTransform: Filter size larger than input image size')
            prop_img = inp_img*psi_filter_bank[i][j]
            prop_img = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(prop_img)))
            prop.extend([prop_img])
            m = {'scale': i}
            new_meta.extend([m])

    return {'prop': prop, 'out': out, 'meta': new_meta}

def conv_comp_img_2_oup_img(img):
#    re = np.real(img.flatten('F'))
#    im = np.imag(img.flatten('F'))
#    res_img = np.array([re, im])
    res_img = np.real(img.flatten('F'))
    res_img.flatten('F')
    res_img = np.reshape(res_img,(np.sqrt(len(res_img)),np.sqrt(len(res_img))))
    return res_img










                      
                      
