# -*- coding: utf-8 -*-
"""
Created on Mon Apr 25 20:27:42 2016

@author: Benedikt Hermann
"""
import numpy as np
import matplotlib.pyplot as plt

from scipy import misc
from scipy import optimize
import os 
import scipy as sp


def sampleavg(path, Basler, g, N=500.):
    #return avg. Image of N Samples, g is an int indicating the summation: 
    #g=1: 1, 3, 5, ... 
    #g=2: 2, 4, 6, ...
    #g=3: all.
    N=np.double(N)    
    i=1
    if g==1:
        I1 = misc.imread(os.path.join(path,Basler+str(2*i-1).zfill(4)+'.bmp'), flatten= 0)
        I=I1/N
        for j in range(int(N)-1):
            i=j+2
            print(i)
            I1 = misc.imread(os.path.join(path,Basler+str(2*i-1).zfill(4)+'.bmp'), flatten= 0)
            I=I+I1/N
    if g==2:
        I1 = misc.imread(os.path.join(path,Basler+str(2*i).zfill(4)+'.bmp'), flatten= 0)
        I=I1/N
        for j in range(int(N)-1):
            i=j+2
            print(i)
            I1 = misc.imread(os.path.join(path,Basler+str(2*i).zfill(4)+'.bmp'), flatten= 0)
            I=I+I1/N
    if g==3:
        I1 = misc.imread(os.path.join(path,Basler+str(i).zfill(4)+'.bmp'), flatten= 0)
        I=I1/N
        for j in range(2*int(N)-1):
            i=j+2
            print(i)
            I1 = misc.imread(os.path.join(path,Basler+str(i).zfill(4)+'.bmp'), flatten= 0)
            I=I+I1/N
    return I
    
def savitzky_golay(y, window_size, order, deriv=0, rate=1):
    from math import factorial
    
    try:
        window_size = np.abs(np.int(window_size))
        order = np.abs(np.int(order))
    except ValueError as msg:
        raise ValueError("window_size and order have to be of type int")
    if window_size % 2 != 1 or window_size < 1:
        raise TypeError("window_size size must be a positive odd number")
    if window_size < order + 2:
        raise TypeError("window_size is too small for the polynomials order")
    order_range = range(order+1)
    half_window = (window_size -1) // 2
    # precompute coefficients
    b = np.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)])
    m = np.linalg.pinv(b).A[deriv] * rate**deriv * factorial(deriv)
    # pad the signal at the extremes with
    # values taken from the signal itself
    firstvals = y[0] - np.abs( y[1:half_window+1][::-1] - y[0] )
    lastvals = y[-1] + np.abs(y[-half_window-1:-1][::-1] - y[-1])
    y = np.concatenate((firstvals, y, lastvals))
    return np.convolve( m[::-1], y, mode='valid')
    
def smooth(Z1, ws=51, order=1):#in: profiles
    #smoothen
    ysg1 = savitzky_golay(Z1, ws, order)
    return ysg1
    
def normalize(Z, plot=0):
    # get G(x) gaussian profile of laser, then normalize
    #Fourier Analysis
    N = len(Z);a=0.;b=N
    t, step = np.linspace(a,b,N,endpoint=False,retstep=True)
    y = Z-np.mean(Z)
    f = 1./N*abs(np.fft.fft(y)) 
    f2 = np.fft.fftshift(f) 
    k = np.r_[-len(y)/2.:len(y)/2.]/(b-a) 
    if plot:
        plt.plot(k,f2,"r-")
        plt.xlim((-0.05,0.05))
        plt.ylabel("DFT $I_{px="+"}-<I_{px="+"}>$", fontsize=16)
        plt.title("Discrete Fourier Transform", fontsize=16)
        plt.xlabel("ky")
        plt.tight_layout()
    m=np.argmax(f2[N/2+2:])
    f= (k[N/2+2+m])
    if plot:
        plt.figure("n")
        plt.plot(Z)
    def GaussSin(x, p1, p2, p3, p4, p5,f):
        return p1*np.exp(-1*p2*(x-p3)**2)*np.sin(2.*np.pi*f/2*x+p4)**2+p5
    p5=np.min(Z)
    p1=np.max(Z)-p5
    p3=np.argmax(Z)
    if plot: 
        plt.figure("n")
        plt.plot(t, GaussSin(t, p1,1e-5,p3,+0., p5,f), "k--") # find start params  
    
    F=optimize.curve_fit(GaussSin,t,Z,p0=[p1,1e-5,p3,+0., p5,f])[0]
    if plot: plt.plot(t, GaussSin(t, F[0],F[1],F[2],F[3],F[4],F[5]))
    def Gauss(x,p1,p2,p3):
        return p1*np.exp(-1*p2*(x-p3)**2)
    if plot: plt.plot(t, Gauss(t, F[0],F[1],F[2])+F[4], "r--")
    
    # Normalisation    
    NZ=(Z-F[4])/Gauss(t, F[0],F[1],F[2])
    if plot: 
        plt.figure()    
        plt.plot(NZ)
        
    return NZ, t, F[5], F[3] #normalized profile, x interval, freq, konstant phase of undist.
    
def alphafit2(Z1, Z2,B=0.015,m=550., b=150., s=50.):
    Z1, t, f1, dp1=normalize(Z1)
    Z2, t, f2, dp2=normalize(Z2)
    
    #fit for sin phase and freq
    def sin2(x, A, f, dp):
        return A*np.sin(2*np.pi*f/2.*x+dp)**2
        
    Para=optimize.curve_fit(sin2, t, Z1, p0=[1., f1, dp1])[0]
    A1=Para[0]; f1=Para[1]; dp1=Para[2]
    
    Para=optimize.curve_fit(sin2, t, Z2, p0=[1., f2, dp2])[0]
    A2=Para[0]
    Z1=Z1/A1; Z2=Z2/A2    
    Dif=Z1-Z2
    def alphaF(x, B, s, m, b):
        deg=2
        return B*(-1*np.exp(-1./s**deg*(x-(m-b))**deg)+np.exp(-1./s**deg*(x-(m+b))**deg))
    def fitF(x,B, s, m, b):
            A=1.
            return A*(np.sin(2*np.pi*f1/2.*x+dp1)**2-np.sin(2*np.pi*f1/2.*x+dp1+alphaF(x, B, s, m, b))**2)
    
    Para=optimize.curve_fit(fitF, t, Dif, p0=[B, s, m, b])[0]
    afit = alphaF(t, Para[0],Para[1],Para[2],Para[3])
    Diffit = fitF(t, Para[0],Para[1],Para[2],Para[3])
    return Dif,Diffit, afit
    
def phaseunwrap(Z1, Z2):
    """ 
    1. Fourier transform
    2. Gaussian window
    3. shift oszillation peak to zero
    3. inverse Fourier Transform -> A(x)
    4. A(x)/Aref(x)
    """
    #1 DFT
    N = len(Z1);a=0.;b=N
    t, step = np.linspace(a,b,N,endpoint=False,retstep=True)
    f = 1./N*(np.fft.fft(Z1)) 
    f2 = np.fft.fftshift(f) 
    k = np.r_[-len(Z1)/2.:len(Z1)/2.]/(b-a)
    #2 Gaussian Window
    def Ga(x):
        return np.exp(-1./x**2*(3.*(k[1]-k[0]))**2)
    W1=(f2)*Ga(k)
     
    W1[0:int(len(W1)/2)]=0.
    #3 shift oszillation peak to zero  
    b=np.argmax(W1)    
    m=int(b-len(W1)/2)
    S1=np.zeros(len(W1), dtype=np.complex64)
    #S1=np.roll(np.array(W1),-m)
    S1[int(len(W1)/2)-m:-m]=W1[int(len(W1)/2):]
#    plt.figure("test")
#    plt.plot(k, S1)
#    plt.plot(k, LP(k))
    
    D1=np.fft.ifft(N*S1)
    
    #same for ref. signal 
    f = 1./N*(np.fft.fft(Z2)) 
    f2 = np.fft.fftshift(f) 
    k = np.r_[-len(Z2)/2.:len(Z2)/2.]/(b-a)
    W2=(f2)*Ga(k)
    W2[0:int(len(W2)/2)]=0
    m=int(np.argmax(W2)-len(W2)/2)
    S2=np.zeros(len(W1), dtype=np.complex64)  
    #S2=np.roll(W2,-m)
    
    S2[int(len(W2)/2)-m:-m]=W2[int(len(W2)/2):]
    D2=np.fft.ifft(N*S2)
    
    #4 A(x)/Aref(x)
    E=D2/D1
    P1=np.zeros(len(S1))
    #print(E)
    
    P1=np.angle(E)
    
    # minus sign if I1 and I2 are interchanged (on<->off)
    #if np.argmax(smooth(P1[int(len(P1)/10):-int(len(P1)/10)]))<np.argmin(smooth(P1[int(len(P1)/10):-int(len(P1)/10)])): P1=P1*(-1.)
        
    return P1

def Ga(x):
    return np.exp(-1./x**2*(6.*(x[1]-x[0]))**2)
    


def ft(D):
    lx=len(D[:,0])
    kx = np.r_[-lx/2.:lx/2.]/lx
    F0=np.fft.fft(D,axis=0)
    F0=np.fft.fftshift(F0, axes=0)
    F0[int(lx/2):,:]=0.
    F0=np.transpose(np.transpose(F0)*Ga(kx))
    IF0=np.fft.ifft(F0,axis=0)
    return IF0

def unwrap(D0,D1,fringe_dir='ver'):
    if fringe_dir=="hor":
        F0=ft(D0)
        F1=ft(D1)
    if fringe_dir=="ver":
        F0=ft(np.transpose(D0))
        F1=ft(np.transpose(D1))
        
    if fringe_dir=="hor": return np.angle(F1/F0)
    if fringe_dir=="ver": return np.transpose(np.angle(F1/F0))

def correct_phase_jumps(P):
    Pc=np.copy(P)
    lx=len(P[:,0])
    ly=len(P[0,:])

    for i in range(ly-1):
        for j in range(lx):
            diff=Pc[j,i]-Pc[j,i+1]
            if abs(diff)>4.:
                g=np.round(diff/np.pi)
                Pc[j:,i+1]+=g*np.pi                 

    return Pc
                
def phase(A,As,shift=217):
    lx=len(A[:,0])
    ly=len(A[0,:]) 
    Ph=np.copy(A)
    Phs=np.copy(As)
    for i in range(lx):
        Ph[i,:]-=np.average(Ph[i,0:50])
        Phs[i,:]-=np.average(Ph[i,0:50])
        for j in range(shift,ly):
            Phs[i,j]=Phs[i,j]+Phs[i,j-shift]
            Ph[i,j]=Ph[i,j]+Phs[i,j-shift]
    return Ph

def MLEM(X,theta, meas):
    from scipy import ndimage
    #rotate
    Xp=ndimage.rotate(X, -theta, reshape=False)
    Lxt=len(X[0,:])
    X1=np.zeros((Lxt,Lxt))
    
    ysim=np.sum(Xp,axis=0)
    ZZ=(len(ysim)-len(meas))/2
    meas=np.append(ysim[:ZZ],meas)
    meas=np.append(meas,ysim[-ZZ:]) 

    T=np.abs(meas/ysim)
    
    X1[:,Lxt/4:-Lxt/4]=Xp[:,Lxt/4:-Lxt/4]*T[Lxt/4:-Lxt/4]
    return ndimage.rotate(X1,theta, reshape=False)
    
def iabel(dfdx, x):
    nx = len(x)

    integral = sp.zeros(nx-1, dtype=float)
    
    for i in range(0, nx-1):
        divisor = sp.sqrt(x[i:]**2 - x[i]**2)
        integrand = dfdx[i:] / divisor
        integrand[0] = integrand[1] # deal with the singularity at x=r
        integral[i] = - sp.trapz(integrand, x[i:]) / sp.pi
    return integral
    
def abel(f,x):
    #f is the the distribution w.r.t. r
    nx=len(x)
    integral= sp.zeros(nx,dtype=float)
    for i in range(0, nx-1):
        divisor = sp.sqrt(x[i:]**2 - x[i]**2)
        integrand = f[i:]*x[i:] / divisor
        integrand[0] = integrand[1] # deal with the singularity at x=r
        integral[i] = 2*sp.trapz(integrand, x[i:])
    return integral


def periodic_gaussian_deriv(x, sigma): 
    nx = len(x) 
    # put the peak in the middle of the array: 
    mu = x[nx/2] 
    g = dgaussiandx(x, mu, sigma) 
    # need to normalize, the discrete approximation will be a bit off: 
    g = g / (-sp.sum(x * g)) 
    # reorder to split the peak across the boundaries: 
    return(sp.append(g[nx/2:nx], g[0:nx/2]))
    
def dgaussiandx(x, mu, sigma):
    return(
        -(x-mu)*sp.exp(-(x-mu)**2/sigma**2/2.0)/(sigma**2*sp.sqrt(2*sp.pi*sigma**2))
         )
    
def noise_img(img):
    nx=len(img[:,0])
    ny=len(img[0,:])
    sx=20
    sy=20
    noise=np.zeros((nx/sx-2, ny/sy-1))
    for i in range(ny/sx-1):
        p=img[:,i*sx+sx/2]
        n=p-smooth(p, ws=21, order=3)
    
        for j in range(nx/sy-2):
            noise[j,i]=np.std(n[j*sy:(j+1)*sy])
            
    return noise
    
def gauss(x,a,x0,sigma):
    return a*np.exp(-(x-x0)**2/(2*sigma**2))    
def gauss_fit(x, y):
    n = len(x)                        
    mean = n/2
    sigma = n/4
    popt,pcov = sp.optimize.curve_fit(gauss,x,y,p0=[1,mean,sigma])
    return popt, pcov