import time

from cryoem.data.cnvtools import pad, crop, circshift
import numpy as np
from numpy import shape, array
from numpy.fft import fftn, ifftn
from numpy import real, floor, conj





class Cnvarray:
    """
    Creates convolution matrix for n-dimensional numpy arrays.

    Mathematical model: y = x * f = Xf = Fx

    Call:  Z = cnvarray(z, sw, mode)

    Input: z        either PSF or image
           sw       shape of array the convolution matrix is applied to
           mode     a flag indicating the size of the output
                    'valid'  (0): The output consists only of those elements that
                                  do not rely on the zero-padding.
                                  sy = sx - sf + 1
                    'same'   (1): The output is the same size as the input centered
                                  with respect to the 'full' output.
                                  sy = sx
                    'full'   (2): The output is the full discrete linear convolution
                                  of the inputs. (*Default*)
                                  sy = sx + sf - 1

    Output: Class object which provides the following methods:

           cnv      convolution
           cnvtp    correlation, transpose operation

    Example:
          from pylab import imread
          from imagetools import rgb2gray

          x = imread('lena.png')
          f = imread('blur.png')

          x = rgb2gray(x)
          f = rgb2gray(f)

          F = cnvarray(f, shape(x), 'valid')
          y = F.cnv(x)

    Copyright (C) 2010 Michael Hirsch
    """

    def __init__(self, f, sx, mode):


        sf = array(shape(f))
        sx = array(sx)

        # Check what is f and what is x
        if all(sx > sf):
            self.f = f
            self.x = []
        elif all(sx < sf):
            self.f = []
            self.x = f
            sf = sx
            sx = array(shape(self.x))
        elif any(sf < sx) and any(sf > sx):
            raise IOError('Size missmatch')

        if mode == 'valid':
            sy   = sx - sf + 1
            sfft = sx
        elif mode == 'same':
            sy   = sx
            sfft = np.vectorize(next_pow2)(sx + sf - 1)
            ## sfft = sx + sf - 1
        elif mode == 'full':
            sy   = sx + sf - 1
            sfft = sx + sf - 1
        elif mode == 'circ':
            sy   = sx
            sfft = sx
        else:
            raise NotImplementedError('Not yet implemented!')

        # Precompute FFTs
        if self.f != []:
            self.fft = fftn(self.f, sfft)
            self.sz  = sx
        else:
            self.fft = fftn(self.x, sfft)
            self.sz  = sf

        self.sfft = sfft
        self.sf   = sf
        self.sx   = sx
        self.sy   = sy
        self.mode = mode

    def cnv(self, u):
        """
        Computes the convolution of the convolution matrix with an image or PSF,
        i.e. Xf or Fx respectively

        Input:  z   image or PSF
        Ouput:  y   blurred image
        """


        t = time.time()
        y  = real(ifftn(fftn(u, self.sfft) * self.fft))
        ## y1 = fftn(u, self.sfft)
        ## t2 = time.time()
        ## y2 = y1 * self.fft
        ## t3 = time.time()
        ## y = real(ifftn(y2))
        ## t4 = time.time()
        ## print u.shape
        ## print y1.shape
        ## print y.shape
        ## print 'cnv: %1.6f, %1.6f, %1.6f' % (t2 - t, t3 - t2, t4 - t3)

        if self.mode == 'valid':
            y = crop(y, self.sy, self.sf-1)
        elif self.mode == 'same':
            ## print y.shape, self.sy, self.sf

            ## y = crop(y, self.sy, floor(self.sf/2) - 1) # Added '-1'.
            y = crop(y, self.sy, floor((self.sf+1)/2) - 1) # Added '-1'.

            ## margin = floor((y.shape - self.sy)/2)
            ## print y.shape, self.sy, self.sf, (y.shape - self.sy)
            ## y = crop(y, self.sy, margin)
        elif self.mode == 'full':
            pass
        elif self.mode == 'circ':
            y = circshift(y, floor(self.sf/2))
        else:
            raise NotImplementedError('Not a valid mode!')

        return y

    def cnvtp(self, y):
        """
        Computes the correlation of the convolution matrix with an image or PSF,
        i.e. X'y or F'y respectively.

        Input:  y   blurred image
        Ouput:  z   image or PSF
        """


        if self.mode == 'valid':
            y = pad(y, self.sfft, self.sf-1)
        elif self.mode == 'same':
            y = pad(y, self.sfft, floor(self.sf/2) - 1) # Added '-1'.
        elif self.mode == 'full':
            pass
        elif self.mode == 'circ':
            y = circshift(y, -floor(self.sy/2))
        else:
            raise NotImplementedError('Not a valid mode!')

        y  = real(ifftn(fftn(y, self.sfft) * conj(self.fft)))
        y  = crop(y, self.sz)

        return y

    def deconv(self, y, z0 = None, mode = 'lbfgsb', maxfun = 100, iprint = 10, m = 5):
        """
        Deconvolution routine. Expects a blurry image and outputs blur kernel or sharp image.
        Currently only deconvolution via lbfgsb is implemented.
        """
        from numpy import prod, array, shape, zeros, infty, reshape, ones
        from scipy.optimize import fmin_l_bfgs_b

        if not all(array(y.shape) == self.sy):
            raise IOError ('Sizes incompatible. Expected blurred image!')

        if mode == 'lbfgsb':

            lf = prod(self.sz)
            if z0 == None:
                z0 = ones(lf)/lf    # initialisation
            else:
                z0 = z0.flatten()
            lb = 0.                 # lower bound
            ub = infty              # upper bound
            zhat = fmin_l_bfgs_b(func = self.cnvinv_objfun, x0 = z0, \
                                 fprime = self.cnvinv_gradfun,\
                                 args = [self.sz, y], factr = 10., pgtol = 10e-15, \
                                 maxfun = maxfun, bounds = [(lb, ub)] * lf,
                                 m = m, iprint = iprint)

        else:
            raise NotImplementedError('Not a valid mode!')

        return reshape(zhat[0],self.sz), zhat[1], zhat[2]


    def cnvinv_objfun(self, z, sz, y):

        from numpy import reshape, real, dot, transpose, ravel, array

        res = y - self.cnv(reshape(z,sz))
        obj = real(0.5 * dot(transpose(ravel(res)), ravel(res)))

        return array([obj])

    def cnvinv_gradfun(self, z, sz, y):

        from numpy import reshape, real

        res = y - self.cnv(reshape(z,sz))
        grad =  - real(self.cnvtp(res))

        return grad.flatten()


def next_pow2(a):
    """Return smallest power of 2, >= a."""

    return int(2**np.ceil(np.log2(a)))

if __name__ == '__main__':

    from pylab import *
    #from imagetools import *
    from numpy import *

    #x  = imread('lena.png')
    #x  = rgb2gray(x)
    x  = random.rand(512,512)
    x = np.zeros([512, 512])
    x[246:266, 246:266] = 1
    sx = shape(x)

    #f  = imread('blur.png')
    #f  = rgb2gray(f)
    f   = random.rand(16,16)
    f = np.ones([17, 17])
    f  /= sum(f.flatten())
    sf = shape(f)

    X = Cnvarray(x, sf, mode = 'same')
    y = X.cnv(f)

    ## figure(1)
    ## subplot(131)
    ## imshow(x,'gray', interpolation='nearest')
    ## title('x')
    ## subplot(132)
    ## imshow(f,'gray', interpolation='nearest')
    ## title('f')
    ## subplot(133)
    ## imshow(y,'gray', interpolation='nearest')
    ## title('y')

    # Comparison to scipy.signal

    import scipy.signal as scsi

    print('-------------------------------------------------------')
    print('Safety check, comparison with scipy.signal.fftconvolve ')
    print('-------------------------------------------------------')

    for mode in ['valid','same','full']:
    ## for mode in ['same']:
        Xcmp = Cnvarray(x, sf, mode)
        ycmp = Xcmp.cnv(f)
        z = scsi.fftconvolve(x,f,mode)
        print('Difference for %s convolution: %e' % (
            mode, sum(abs(ycmp-z))))
        ## figure(2)
        ## subplot(1, 2, 1)
        ## imshow(ycmp,'gray', interpolation='nearest')
        ## title('ycmp (us)')
        ## subplot(1, 2, 2)
        ## imshow(z,'gray', interpolation='nearest')
        ## title('z (scipy)')


    if True:
        print('-------------------------------------------------------')
        print('Time comparison for convolution computation')
        print('-------------------------------------------------------')
        import time as t

        ti = t.clock()
        for i in arange(10):
            z = scsi.fftconvolve(x,f,'valid')
        tf = t.clock()
        print('Time elapsed for scipy.signal.fftconvolve2d %.3f' %(tf-ti))

        ti = t.clock()
        X = Cnvarray(x, sf, mode = 'valid')
        for i in arange(10):
            z = X.cnv(f)
        tf = t.clock()
        print('Time elapsed for Cnvarray %.3f' % (tf-ti))

    ## raise

    ## print('-------------------------------------------------------')
    ## print('Deconvolution: PSF estimation')
    ## print('-------------------------------------------------------')

    ## iter = 100
    ## fhat = X.deconv(y, maxfun = iter)

    #figure(2)
    #imshow(fhat[0])
    #title('Estimated PSF after %d iterations' % iter)

