"""
This py file is to extract pulse duration from Interferometric Autocorrelation (IAC) 
Trace of Chirped Pulse by Fourier Analysis
"""
from numpy.fft import fft, ifft
from numpy import real, linspace, log, exp, cosh, sqrt, pi, conj, zeros, ones, where, array
from scipy import signal
import matplotlib.pyplot as plt

def fft_filter(y,frange,mode=0,pack=0):
    ''' 
    y = data to be analyzed \n
   
    frange = list f1 and f2 frequencies index to be included or excluded \n
    mode = 0 for inclusive band, else for exclusive band \n
    pack = 1 for unpacking outputs mainly for debugging purpose and intermediate steps \n
    '''
     
    fft_y=fft(y)
          
    #getting conditionals statements
    frange = array(frange)
    f1 = frange[0:len(frange):2]
    f2 = frange[1:len(frange):2]
   
    for i,j in zip (f1,f2):
               
    #band inclusion or rejection mode
        if mode == 0: 
            freq_filter = zeros(len(y))
            freq_filter[i:j] = 1
            freq_filter_r = freq_filter[::-1]
            freq_filter_r[i+1:j+1]=1 #symmetrize
        else :
            freq_filter = ones(len(y))
            freq_filter[i:j] = 0
            freq_filter_r = freq_filter[::-1]
            freq_filter_r[i+1:j+1]=0 #summetrize

    #filtered fft data
    ffty_f=fft_y*freq_filter_r[::-1]
    ry=real(ifft(ffty_f))
         
    if pack == 0:
        return ry/max(ry)
        
    else:
        return ry/max(ry),fft_y/max(fft_y),freq_filter/max(freq_filter),ffty_f/max(ffty_f) #packing is only for debugging mode

        

#direct autocorrelator analysis 
def pulse(t,t0, wl, b=0 ,shape='gauss'): 
    w0=2*pi*3e2/wl #in 1/fs as t is in fs
    if (shape=='sech'):
        E_t=1/cosh(2*log(1+sqrt(2))*t/t0)**(1+1j*b)*exp(-1j*w0*t)
    else:
        E_t=exp(-2*log(2)*(t/t0)**2)*exp(-1j*w0*t)*exp(-1j*b*2*log(2)*(t/t0)**2) #linear chirp defined by b
    return E_t

#interoferometric autocorrelation signal of given pulse
def iac(E_t):
    I_t=abs(E_t)**2
    IAC_constant=2*sum(I_t**2)
    
    IAC_intensity=4*signal.correlate(I_t,I_t,mode="same")
    IAC_1w_dump=signal.correlate(I_t*E_t,E_t,mode="same")+signal.correlate(E_t,I_t*E_t,mode="same")

    IAC_1w=2*IAC_1w_dump+2*conj(IAC_1w_dump)
    IAC_2w_dump=signal.correlate(E_t**2,E_t**2,mode="same")
    IAC_2w = IAC_2w_dump+conj(IAC_2w_dump)
    IAC=IAC_constant+IAC_intensity+IAC_1w+IAC_2w
    IAC=abs(IAC)
    IAC=8*IAC/max(IAC)
    return IAC,8*IAC_intensity/max(IAC_intensity)
    #IAC_intensity is the intensity autocorrelation, later used to check the code.
    
#pulse
t=linspace(-200,200,1000) #in femtosecond
e_t=pulse(t,25,800,2,'sech') #sech pulse
dt = t[1]-t[0]
a_corr,a_intensity=iac(e_t)

#plotting figures
plt.figure(1)
plt.plot(t,a_corr,t,a_intensity)
#Getting rid  of dc offset 1 from interferometric autocorrelation
a_corr = a_corr -1
a_w = fft(a_corr)
plt.figure(2)
plt.plot(a_w)
#to get the idea of index of frequencies to be selected

#fourier filter
y = fft_filter(a_corr,[0, 25])

def fwhm(x,y):
    half=max(y)/2.0
    xt = x[where(y>=half)]
    return xt[-1]-xt[0]

#checking
if abs(fwhm(t,y)-fwhm(t,a_intensity))<1e-8:
    print("Good!")
