
import numpy as np
from skimage.io import imshow, show
import matplotlib.pyplot as plt

class gabortrafo:

    def __init__(self,  num_scales, 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._sigma_psi = 0.5*num_scales    #rule of thumb, filter sehen einigermassen okay aus mit diesem Faktor
                                            #_sigma_psi bestimmt die breite der Gausskerne
                                            #(-3dB-Zonen der verschiedenen Filter sind direkt aneinander angrenzend)
        self._psi_filter_bank = self._generate_psi_filter_bank()
        
        #falls direkt ein Output generiert werden soll
        self._phi_filter = self._generate_phi_filter()


        #comment to hide filter plots
        psifilters = [val for one_filter_bank in self._psi_filter_bank for val in one_filter_bank]
        fig, axes = plt.subplots(2*num_scales+1, 2*num_scales+1, 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 _gabortrafo(img, meta, self._num_scales, self._phi_filter, self._psi_filter_bank)


    #for propagated images
    def _generate_psi_filter_bank(self):
        psi_filter_bank = [[0]*(2*self._num_scales+1)]*(2*self._num_scales+1)
        for i in range(0, len(psi_filter_bank)):
            a = [0]*len(psi_filter_bank[i])
            for j in range(0, len(psi_filter_bank)):
                a[j]=self._generate_psi_single_filter(j-self._num_scales, i-self._num_scales)
            psi_filter_bank[i]=a
        return psi_filter_bank


    def _generate_psi_single_filter(self, xp, yp):
        xi_x, xi_y = self._generate_grid(xp, yp)
        filter = np.exp(-2*((np.pi*self._sigma_psi)**2)*((xi_x**2)+(xi_y**2)))
        return filter
    
    
    #for output image
    def _generate_phi_filter(self):
        xi_x, xi_y = self._generate_grid(0, 0)
        return np.exp(-2*((np.pi*self._sigma_psi)**2)*((xi_x**2)+(xi_y**2)))


    def _generate_grid(self, xp, yp):
        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)
        for k in range(0, xi_x.shape[0]):
            for l in range(0, xi_x.shape[1]):
                xi_x[k,l]=xi_x[k,l]-xp/(self._num_scales+1)
        for k in range(0, xi_y.shape[0]):
            for l in range(0, xi_y.shape[1]):
                xi_y[k,l]=xi_y[k,l]-yp/(self._num_scales+1)
        return [xi_x, xi_y]




def _gabortrafo(img, meta, num_scales, phi_filter, psi_filter_bank):
    prop = []
    out = []
    inp_img = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(img)))
    
    #uncomment, falls direkt ein output generiert werden soll
    """
    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])
    """
    
    #es wird jeweils höchstens bis zur äussersten Corona des vorherigen Filters iteriert
    if meta is not None and 'corona' in meta:
        current_scale = min(meta['corona'], num_scales)
    else:
        current_scale = num_scales

    new_meta = []
    for i in range(num_scales-current_scale, num_scales+current_scale+1):
        for j in range(num_scales-current_scale, num_scales+current_scale+1):
            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]

            if(i+j<2*num_scales+1): #Wegen Symmetrie von reellen Signalen müssen andere Fälle nicht berücksichtigt werden
                prop_img = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(prop_img)))
                prop.extend([prop_img])
                current_corona = max(abs(i-num_scales), abs(j-num_scales))
                m = {'corona': current_corona}
                new_meta.extend([m])

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



#wird nur für direkte Output-Generation benötigt
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())
    res_img.flatten('F')
    print(len(res_img))
    res_img = np.reshape(res_img,(np.sqrt(len(res_img)),np.sqrt(len(res_img))))
    return res_img










