# -*- coding: utf-8 -*-
"""
Created on Fri May 15 09:04:45 2020

@author: Chathu Dilrukshi and TJ Hammond
https://www.uwindsor.ca/people/thammond/302/projects
If you have any comments or questions, please let us know at
thammond 'at' uwindsor.ca

This solves for the gate and pulse from a frequency resolved optical gating (FROG)
trace using singular value decomposition (SVD).

The input uses a symmetric signal, simulating the SHG FROG setup. This setup doesn't
discriminate forwards or backwards time, and so the reconstructed phase can be the
opposite of the real phase. You can account for this ambiguity by changing Tsign from 1 to -1
if the pulse direction is wrong (and the phase signs).

This code is based on the paper DJ Kane "Recent progress toward real-time measurements of ultrashort pulses" JQE (1999)
and Matlab FROG code from Adam Wyatt (2020). Frequency-resolved optical gating (FROG) (https://www.mathworks.com/matlabcentral/fileexchange/16235-frequency-resolved-optical-gating-frog), MATLAB Central File Exchange. Retrieved May 31, 2020. 

"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from drawnow import drawnow

plt.close('all')

c = 2.9979e8    #speed of light

NI = 201        #number of iterations for solving FROG
skipI = 5      #number of iterations to skip when plotting
#%% Pulse character

tFWHM = 25e-15          #full width half maximum pulse duration
lambda1 = 800e-9        #central wavelength

GDD = 200e-30           #group delay dispersion
TOD = 2000e-45          #third order dispersion
FOD = 20000e-60         #fourth order dispersion

Tsign = 1               #1 if sign is right, -1 if sign is wrong

#%% Set up the grid, time and angular frequency

Nt = 2**6
T = 200e-15
t = np.linspace(-T,T,Nt)
dt = t[1]-t[0]
w = 2*np.pi/(2*dt)*np.linspace(-1,1,Nt)
dw = w[1]-w[0]

#%% The original pulse

w1 = 2*np.pi*c/lambda1  #central angular frequency, unused since we just use envelope
tau = tFWHM/(2*np.sqrt(np.log(2))) #converting from FWHM pulse duration to 1/e
Et0 = np.exp(-t**2/(2*tau**2)) #transform limited pulse

PHI0 = w**2*GDD/2 + w**3*TOD/6 #the phase in frequency
Ew0 = np.fft.fftshift(np.fft.fft(np.fft.ifftshift(Et0))) #convert to frequency
Ew0p = Ew0*np.exp(-1j*PHI0) #add phase
Ew0p = Ew0p/np.max(np.abs(Ew0p)) #normalize

Et0p = np.fft.fftshift(np.fft.ifft(np.fft.ifftshift(Ew0p))) #convert back to time
Et0p = Et0p/np.max(np.abs(Et0p)) #normalize

Et0pnan = Et0p*1 #copy the time array without reference, could also use Et0p.copy()
Et0pnan[np.abs(Et0pnan)<1e-3]=np.nan #make small values undefined for phase plots
Ew0pnan = Ew0p*1 #copy the frequency array without reference, could also use Ew0p.copy()
Ew0pnan[np.abs(Ew0pnan)<1e-3]=np.nan #do the same thing for frequency space

#%% Set up the FROG trace/spectrogram; here create the outer product

Pt0 = Et0p #the pulse
Gt0 = Pt0 #the gate

outer = np.outer(Pt0,Gt0) + np.outer(Gt0,Pt0)

#%% Shift outer matrix, FFT to get spectrogram

shifted = np.zeros((Nt,Nt),dtype='complex') #you need a shifted and unshifted matrix
unshift = np.zeros((Nt,Nt),dtype='complex')

for n in range(0,Nt):
    shifted[n,:] = np.roll(outer[n,:], -n, axis=0)

pfrog = np.fft.ifftshift(np.fft.ifft(shifted,axis=0))
spectrogram = np.abs(pfrog)/np.max(np.max(np.abs(pfrog)))

#%% Create the original guesses; using the transform limited pulses one way to do it

Pt = Et0
Gt = Pt

#%% Set up the figure

fig, axs = plt.subplots(3, 2, figsize=(6.4,7))

#%% Now the business

for countI in range(0,NI):

    outer = np.outer(Pt,Gt) + np.outer(Gt,Pt) #create outer product from guess

    #%% Do the shifting and the FFT
    
    for n in range(0,Nt):
        shifted[n,:] = np.roll(outer[n,:], -n, axis=0)    
    
    pfrog = np.fft.fftshift(np.fft.ifft(shifted,axis=0)) #this is the FROG trace
    pfrog = pfrog/np.max(np.max(np.abs(pfrog))) 
    afrog = np.abs(pfrog) 
    
    pfnew = spectrogram*np.exp(1j*np.angle(pfrog)) #use the spectrogram amplitude with the FROG phase
    
    #%% Go back to SVD matrix
    
    timer = np.fft.fft(np.fft.fftshift(pfnew), axis=0) #go back to time domain
    
    for n in range(0,Nt):
        unshift[n,:] = np.roll(timer[n,:], n, axis=0)
    
    [U, S, V] = np.linalg.svd(unshift) #use SVD magic
    Pt = U[:,0]
    Gt = V[0,:]
    
    Pt = Pt/max(abs(Pt));
    Gt = Gt/max(abs(Gt));
    
    #%% Plotting
    
    if np.mod(countI,skipI)==0:
        #plot of spectrogram
        axs[0,0].clear()
        my_hot = matplotlib.cm.get_cmap('hot')
        my_hot.set_under('w')
        axs[0,0].pcolormesh(spectrogram,cmap=my_hot,vmin=1e-3)
        axs[0,0].set_title('Iteration %d' % countI)
        
        #plot of FROG trace
        axs[0,1].clear()
        error = np.sum(np.sum(np.abs(spectrogram-afrog)))/Nt
        my_cmr = matplotlib.cm.get_cmap('CMRmap')
        my_cmr.set_under('w')
        axs[0,1].pcolormesh(afrog,cmap=my_cmr,vmin=1e-3)
        axs[0,1].set_title('Error %1.2e' % error)
        
        #plot of Et
        axs[1,0].clear()
        Ptnan = Pt*1 #could also use Pt.copy()
        Ptnan[np.abs(Ptnan)<1e-3]=np.nan
        axs[1,0].plot(t*1e15,np.abs(Et0pnan),'b',Tsign*t*1e15,np.abs(Ptnan),'r')
        axs[1,0].set_xlim((t[0]*1e15,t[-1]*1e15))
        axs[1,0].set_title('|E(t)|')
        
        #plot of Ew
        Pw = np.fft.fftshift(np.fft.fft(np.fft.fftshift(Pt)))
        Pw = Pw/np.max(np.abs(Pw))
        Pwnan = Pw*1 #could also use Pw.copy()
        Pwnan[np.abs(Pwnan)<1e-3]=np.nan
        axs[1,1].clear()
        axs[1,1].plot(w*1e-15,np.abs(Ew0pnan),'b',w*1e-15,np.abs(Pwnan),'r')
        axs[1,1].set_xlim((w[0]*1e-15,w[-1]*1e-15))
        axs[1,1].set_title('|E(w)|')
        
        #phase of Et
        axs[2,0].clear()
        anglert0 = np.unwrap(np.angle(Et0p))
        anglert0[np.isnan(Et0pnan)]=np.nan
        anglert = Tsign*np.unwrap(np.angle(Pt))
        anglert[np.isnan(Ptnan)]=np.nan
        axs[2,0].plot(t*1e15,anglert0-anglert0[Nt//2],'b',Tsign*t*1e15,anglert-anglert[Nt//2],'r')
        axs[2,0].set_xlim((t[0]*1e15,t[-1]*1e15))
        axs[2,0].set_title('$\phi(t)$')
        
        #phase of Ew
        axs[2,1].clear()
        anglerw0 = np.unwrap(np.angle(Ew0p))
        anglerw0[np.isnan(Ew0pnan)]=np.nan
        anglerw = Tsign*np.unwrap(np.angle(Pw))
        anglerw[np.isnan(Pwnan)]=np.nan
        axs[2,1].plot(w*1e-15,anglerw0-anglerw0[Nt//2],'b',w*1e-15,anglerw-anglerw[Nt//2],'r')
        axs[2,1].set_xlim((w[0]*1e-15,w[-1]*1e-15))
        axs[2,1].set_title('$\phi(\omega)$')
        
        plt.subplots_adjust(top=0.9, bottom=0.1, left=0.1, right=0.9, hspace=0.6, wspace=0.2)
        
        plt.pause(0.01)
        drawnow