"""
Author: Samim Zahoor Taray <zsameem@gmail.com>
Copyright: All rights reserved.
"""


import numpy as np
import torch

from .torch_utils import *
from .utils_deblur import *

class VaryingConvOperator():
    
    def __init__(self, otf_array, k_h, k_w, p_nh, p_nw, h, w, initfrom=None):
        self.otf_array = otf_array
        self.k_h = k_h
        self.k_w = k_w
        self.p_nh = p_nh
        self.p_nw = p_nw
        self.h = h
        self.w = w
        self.o_hw = 10
        self.mask_array = self.gen_windows(self.h, self.w, self.p_nh, self.p_nw, self.o_hw)
        # self.is2D = True if len(kernel.shape) > 1 else False
        # self.mode = 'reflect' #{'reflect', 'constant', 'nearest', 'mirror', 'wrap'}
    def gen_windows(self, h, w, p_nh, p_nw, o_hw):
        """ Generate window functions with overlap to implement Efficient
        Filter Flow Algorithms
        Args:
            h, w (int): height and width of the input image
            p_nh, p_nw (int): Number of regions in which input image will be
            divided by the windowing function
            o_hw: Degree of overlap between patches

        Returns:
            bool: 4-D array of dimension [p_nh, p_nw, h, w] containing p_nh*p_nw
            normalized window functions for each rectangular overlapping region 
            of the image.

        """
        mask_array = np.zeros((p_nh, p_nw, h, w))
        # print("generating masks of shape {}".format(mask_array.shape))
        for i in range(p_nh):
            for j in range(p_nw):
                # calculate the size of overlapping window.
                # If the patch is along the borders (i==0, or j==0) then
                # the size will be smaller. We calculate this below.
                # vert_pos = 0 = top boundary, 1 = middle, 2 = bottom boundary
                # hor_pos = 0 = left boundary, 1 = middle, 2 = right boundary
                vert_pos = 0
                if i == 0: # On top boundary
                    ph_start = 0 
                    ph_end = ph_start + int(h/p_nh) + o_hw
                elif i== p_nh-1: # On bottom boundary
                    ph_end = h
                    ph_start = h - int(h/p_nh) - o_hw
                    vert_pos = 2
                else:
                    ph_start = i* int(h/p_nh) - o_hw
                    ph_end = (i+1)* int(h/p_nh) + o_hw
                    vert_pos = 1

                hor_pos = 0
                if j == 0: # On left boundary
                    pw_start = 0 
                    pw_end = int(w/p_nw) + o_hw
                elif j == p_nw-1: # On bottom boundary
                    pw_end = w
                    pw_start = w - int(w/p_nw) - o_hw
                    hor_pos = 2
                else:
                    pw_start = j * int(w/p_nw) - o_hw
                    pw_end = (j+1) * int(w/p_nw) + o_hw
                    hor_pos = 1

                # Index into pre-allocated mask_array and set entries to 1 corrosponding to
                # the current patch
                # Apply damping to borders.
                mask_array[i, j, ph_start:ph_end, pw_start: pw_end] = 1
                for border_i in range(2*o_hw):   

                    if vert_pos != 0:
                        # Upper border
                        mask_array[i,j, ph_start+border_i, pw_start+border_i:pw_end-border_i] = border_i/(2*o_hw)

                    # Right border
                    if hor_pos != 2:
                        mask_array[i,j,ph_start+border_i:ph_end-border_i, pw_end-border_i] = border_i/(2*o_hw)

                    # Lower border
                    if vert_pos != 2:
                        mask_array[i,j, ph_end-border_i, pw_start+border_i:pw_end-border_i+1] = border_i/(2*o_hw)

                    # Left border
                    if hor_pos != 0:
                        mask_array[i,j,ph_start+border_i:ph_end-border_i, pw_start+border_i] = border_i/(2*o_hw)

        sum_array = np.zeros((h,w))
        for i in range(p_nh):
            for j in range(p_nw):
                sum_array += mask_array[i, j, :,:]
        normalized_masks = mask_array[:,:]/sum_array
        return normalized_masks

    def __mul__(self, img):
        '''
        do the actual convolution.
        If the kernel is 1D, a separable convolution is done.
        '''
        return self.eval_Ax(img, self.mask_array, self.otf_array)


    def T(self, img):
        # res = ConvolutionOperator(self, initfrom=self)
        # res.kern = res.kernel.T
        # return res
        return self.eval_At_x(img, self.mask_array, self.otf_array)

    def eval_Ax(self, x, window_array, otf_array):
        """Efficient Filter Flow implementation for evaluation of y = Ax using
        the Spatially Varying Overlap and Add Method.
        Args:
        
        x: 2D Numpy array containing input image as normalized float values
        
        window_array: 4D window function of dimension [p_nh, p_nw, h, w] where
        the last 2 dims represent a window function over the entire image. 
        
        psf_array: 4D Array of dimension [p_nh, p_nw, h, w] containing the FFT of
        the point spread functions for each rectangular overlapping region in the 
        image. (dtype is complex)
        
        Paper: http://suvrit.de/papers/cvpr10.pdf
        
        """
        # # Fourier transform of each masked image
        # masked_ffts = np.zeros([p_nh, p_nw, h, w]).astype(np.complex)
        sum_fft = np.zeros((self.h,self.w), dtype=np.complex)
        for i in range(self.p_nh):
            for j in range(self.p_nw):
                masked_ij = x * window_array[i,j]
                patch_ij = np.fft.ifft2( np.fft.fft2 (masked_ij) * otf_array[i,j])
        #       masked_ffts[i, j, :, :] = patch_ij
                sum_fft += patch_ij
        return np.real(sum_fft) 

    def eval_At_x(self, x, window_array, otf_array):
        """Efficient Filter Flow implementation for evaluation of y = A(transpose)x using
        the Spatially Varying Overlap and Add Method.
        Args:
        
        x: 2D Numpy array containing input image as normalized float values
        
        window_array: 4D window function of dimension [p_nh, p_nw, h, w] where
        the last 2 dims represent a window function over the entire image. 
        
        psf_array: 4D Array of dimension [p_nh, p_nw, h, w] containing the FFT of
        the point spread functions for each rectangular overlapping region in the 
        image. (dtype is complex)
        
        Paper: http://suvrit.de/papers/cvpr10.pdf
        
        """
        # # Fourier transform of each masked image
        # masked_ffts = np.zeros([p_nh, p_nw, h, w]).astype(np.complex)
        sum_fft = np.zeros((self.h,self.w), dtype=np.complex)
        for i in range(self.p_nh):
            for j in range(self.p_nw):
    #             masked_ij = x * window_array[i,j]
                patch_ij = (np.fft.ifft2( np.fft.fft2 (x) * np.conj(otf_array[i,j]))) * window_array[i,j]
        #       masked_ffts[i, j, :, :] = patch_ij
                sum_fft += patch_ij
        return np.real(sum_fft)



class VaryingConvOperatorTorch():
    
    def __init__(self, otf_array, k_h, k_w, p_nh, p_nw, h, w, initfrom=None):
        if otf_array.dtype == np.complex:
            otf_tensor_real = torch.from_numpy(otf_array.real)
            otf_tensor_imag = torch.from_numpy(otf_array.imag)
            # print(otf_tensor_imag.shape)
            self.otf_array = torch.stack((otf_tensor_real, otf_tensor_imag), dim=4).float()
        # self.otf_array = otf_array
        self.k_h = k_h
        self.k_w = k_w
        self.p_nh = p_nh
        self.p_nw = p_nw
        self.h = h
        self.w = w
        self.o_hw = 10
        self.mask_array = self.gen_windows(self.h, self.w, self.p_nh, self.p_nw, self.o_hw)
        # self.is2D = True if len(kernel.shape) > 1 else False
        # self.mode = 'reflect' #{'reflect', 'constant', 'nearest', 'mirror', 'wrap'}
    def gen_windows(self, h, w, p_nh, p_nw, o_hw):
        """ Generate window functions with overlap to implement Efficient
        Filter Flow Algorithms
        Args:
            h, w (int): height and width of the input image
            p_nh, p_nw (int): Number of regions in which input image will be
            divided by the windowing function
            o_hw: Degree of overlap between patches

        Returns:
            bool: 4-D array of dimension [p_nh, p_nw, h, w] containing p_nh*p_nw
            normalized window functions for each rectangular overlapping region 
            of the image.

        """
        mask_array = np.zeros((p_nh, p_nw, h, w))
        # print("generating masks of shape {}".format(mask_array.shape))
        for i in range(p_nh):
            for j in range(p_nw):
                # calculate the size of overlapping window.
                # If the patch is along the borders (i==0, or j==0) then
                # the size will be smaller. We calculate this below.
                # vert_pos = 0 = top boundary, 1 = middle, 2 = bottom boundary
                # hor_pos = 0 = left boundary, 1 = middle, 2 = right boundary
                vert_pos = 0
                if i == 0: # On top boundary
                    ph_start = 0 
                    ph_end = ph_start + int(h/p_nh) + o_hw
                elif i== p_nh-1: # On bottom boundary
                    ph_end = h
                    ph_start = h - int(h/p_nh) - o_hw
                    vert_pos = 2
                else:
                    ph_start = i* int(h/p_nh) - o_hw
                    ph_end = (i+1)* int(h/p_nh) + o_hw
                    vert_pos = 1

                hor_pos = 0
                if j == 0: # On left boundary
                    pw_start = 0 
                    pw_end = int(w/p_nw) + o_hw
                elif j == p_nw-1: # On bottom boundary
                    pw_end = w
                    pw_start = w - int(w/p_nw) - o_hw
                    hor_pos = 2
                else:
                    pw_start = j * int(w/p_nw) - o_hw
                    pw_end = (j+1) * int(w/p_nw) + o_hw
                    hor_pos = 1

                # Index into pre-allocated mask_array and set entries to 1 corrosponding to
                # the current patch
                # Apply damping to borders.
                mask_array[i, j, ph_start:ph_end, pw_start: pw_end] = 1
                for border_i in range(2*o_hw):   

                    if vert_pos != 0:
                        # Upper border
                        mask_array[i,j, ph_start+border_i, pw_start+border_i:pw_end-border_i] = border_i/(2*o_hw)

                    # Right border
                    if hor_pos != 2:
                        mask_array[i,j,ph_start+border_i:ph_end-border_i, pw_end-border_i] = border_i/(2*o_hw)

                    # Lower border
                    if vert_pos != 2:
                        mask_array[i,j, ph_end-border_i, pw_start+border_i:pw_end-border_i+1] = border_i/(2*o_hw)

                    # Left border
                    if hor_pos != 0:
                        mask_array[i,j,ph_start+border_i:ph_end-border_i, pw_start+border_i] = border_i/(2*o_hw)

        sum_array = np.zeros((h,w))
        for i in range(p_nh):
            for j in range(p_nw):
                sum_array += mask_array[i, j, :,:]
        normalized_masks = mask_array[:,:]/sum_array
        return torch.from_numpy(normalized_masks).float()

    def __mul__(self, img):
        '''
        do the actual convolution.
        If the kernel is 1D, a separable convolution is done.
        '''
        return self.eval_Ax(img, self.mask_array.cuda(), self.otf_array.cuda())


    def T(self, img):
        # res = ConvolutionOperator(self, initfrom=self)
        # res.kern = res.kernel.T
        # return res
        return self.eval_At_x(img, self.mask_array.cuda(), self.otf_array.cuda())

    def eval_Ax(self, x, window_array, otf_array):
        """Efficient Filter Flow implementation for evaluation of y = Ax using
        the Spatially Varying Overlap and Add Method.
        Args:
        
        x: 2D Numpy array containing input image as normalized float values
        
        window_array: 4D window function of dimension [p_nh, p_nw, h, w] where
        the last 2 dims represent a window function over the entire image. 
        
        psf_array: 4D Array of dimension [p_nh, p_nw, h, w] containing the FFT of
        the point spread functions for each rectangular overlapping region in the 
        image. (dtype is complex)
        
        Paper: http://suvrit.de/papers/cvpr10.pdf
        
        """
        # # Fourier transform of each masked image
        # masked_ffts = np.zeros([p_nh, p_nw, h, w]).astype(np.complex)
        sum_fft = torch.zeros((self.h,self.w, 2)).cuda()
        for i in range(self.p_nh):
            for j in range(self.p_nw):
                masked_ij = x * window_array[i,j]
                fft_conv = tensor_complex_mulx ( torch.rfft(masked_ij, 2, onesided=False), otf_array[i,j])
                patch_ij = torch.ifft( fft_conv, 2)
        #       masked_ffts[i, j, :, :] = patch_ij
                sum_fft += patch_ij
        return sum_fft[:,:,0] 

    def eval_At_x(self, x, window_array, otf_array):
        """Efficient Filter Flow implementation for evaluation of y = A(transpose)x using
        the Spatially Varying Overlap and Add Method.
        Args:
        
        x: 2D Numpy array containing input image as normalized float values
        
        window_array: 4D window function of dimension [p_nh, p_nw, h, w] where
        the last 2 dims represent a window function over the entire image. 
        
        psf_array: 4D Array of dimension [p_nh, p_nw, h, w] containing the FFT of
        the point spread functions for each rectangular overlapping region in the 
        image. (dtype is complex)
        
        Paper: http://suvrit.de/papers/cvpr10.pdf
        
        """
        # # Fourier transform of each masked image
        # masked_ffts = np.zeros([p_nh, p_nw, h, w]).astype(np.complex)
        sum_fft = torch.zeros((self.h,self.w, 2), dtype=x.dtype).cuda()
        for i in range(self.p_nh):
            for j in range(self.p_nw):
    #             masked_ij = x * window_array[i,j]
                fft_conv = tensor_complex_mulx(torch.rfft(x, 2,onesided=False), self.make_conjugate(otf_array[i,j]))
                patch_ij = torch.ifft( fft_conv, 2) * torch.stack((window_array[i,j], window_array[i,j]), dim=2) 
        #       masked_ffts[i, j, :, :] = patch_ij
                sum_fft += patch_ij
        return sum_fft[:,:,0]
    
    def get_otfarray(self):
        return self.otf_array

    def get_maskarray(self):
        return self.mask_array

    def make_complex(self, arr):
        new_arr = torch.stack((arr, torch.zeros(arr.shape, dtype=arr.dtype)), dim=2)
        return new_arr

    def make_conjugate(self, arr):
        new_arr = torch.stack((arr[:,:,0], -arr[:,:,1]), dim=2)
        return new_arr


def initialize_varying_operators(h, w, kernel_array):
    if len(kernel_array.shape) == 4:
        [p_nh, p_nw, k_h, k_w] = kernel_array.shape
        
        otf_r = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        otf_g = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        otf_b = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        
        for i in range(p_nh):
            for j in range(p_nw):
                k = kernel_array[i, j, :,:]
                k = np.float64(k)  # float64 !
                k = k/np.sum(k)
                otf_b[i, j] = psf2otf(k, (h,w))
                otf_g[i, j] = psf2otf(k, (h,w))
                otf_r[i, j] = psf2otf(k, (h,w))

        k_r = VaryingConvOperatorTorch(otf_r, k_h, k_w, p_nh, p_nw, h, w)
        k_g = VaryingConvOperatorTorch(otf_g, k_h, k_w, p_nh, p_nw, h, w)
        k_b = VaryingConvOperatorTorch(otf_b, k_h, k_w, p_nh, p_nw, h, w)
        return k_r, k_g, k_b

    elif len(kernel_array.shape) == 5:
        [p_nh, p_nw, k_h, k_w, _] = kernel_array.shape
        
        otf_r = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        otf_g = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        otf_b = np.zeros((p_nh, p_nw, h,w)).astype(np.complex)
        
        for i in range(p_nh):
            for j in range(p_nw):
                k = kernel_array[i, j]
                k = np.float64(k)  # float64 !
                # k = k/np.sum(k)
                otf_b[i, j] = psf2otf(k[:,:,0], (h,w))
                otf_g[i, j] = psf2otf(k[:,:,1], (h,w))
                otf_r[i, j] = psf2otf(k[:,:,2], (h,w))

        k_r = VaryingConvOperatorTorch(otf_r, k_h, k_w, p_nh, p_nw, h, w)
        k_g = VaryingConvOperatorTorch(otf_g, k_h, k_w, p_nh, p_nw, h, w)
        k_b = VaryingConvOperatorTorch(otf_b, k_h, k_w, p_nh, p_nw, h, w)
        return k_r, k_g, k_b
    
    else:
        raise("Don't know how to handle this kernel array")
