clear;
clc;
%% Select the output flag for the CCD camera= 1: NIRCCD; 2: VISCCD
CCD_flag=1;
if CCD_flag==1
    Nw=256;    % Number of pixels of the NIR CCD    
elseif CCD_flag==2
    Nw=1340;   % Number of pixels of the Visible light CCD
end
%% Output file name for simulation result
 SI_data_output='Dimer_2DCS.txt';
 LO_spectrum_output='LOfield.txt' ;
%% Parameter assignments of laser source
w_0=230. ;  % [THz] Central frequency of the exciting pulse
t_d=60. ;   % e-1 pulse width [fs]
delw=440./t_d ; % e-1 spectral width [THz] of the laser pulse
%% Model parameter definition and assignment for the Heterodimer
% For more information, please see arXiv:1307.5855v1.
w_a=225. ; % [THz] Energy level at site a
w_b=230. ; % [THz] Energy level at site b
mu_a=0.8 ; % [D] Transition moment at site a
mu_b=1.5 ; % [D] Transition moment at site b
delE=0.5*(w_b-w_a)  ;
J=3.  ;    % J/h [THz]: Dimer's coupling strength
avew=0.5*(w_a+w_b) ;
theta2=atan(J/delE) ;
theta=0.5*theta2 ;
w_1=avew-delE/cos(theta2); % Energy of the excited level 1 of the dimer
w_2=avew+delE/cos(theta2); % Energy of the excited level 2 of the dimer
w_f1=w_2 ;% Energy from the excited level 1 to the final level of the dimer
w_f2=w_1 ;% Energy from the excited level 2 to the final level of the dimer
omega_1 = w_1*(pi/500.) ;
omega_2 = w_2*(pi/500.) ;
omega_f1=omega_2 ;
omega_f2=omega_1 ;
mu_1g=mu_a*cos(theta)+mu_b*sin(theta) ; % Transition moment from g to 1
mu_2g=-mu_a*sin(theta)+mu_b*cos(theta) ; % Transition moment from g to 2
mu_f1=mu_a*sin(theta)+mu_b*cos(theta) ; % Transition moment from 1 to f
mu_f2=mu_a*cos(theta)-mu_b*sin(theta) ; % Transition moment from 2 to f
gam=4.0 ; % Damping constant [THz] of the excited levels
decay=gam/1000. ; % [fs-1]
% Terms of the echo field
term_d1=2.*mu_1g^4 ;
term_d2=2.*mu_1g^2*mu_2g^2 ;
term_c1=2.*mu_2g^4 ;
term_c2=term_d2 ;
term_2q1=-mu_f1*mu_1g*(mu_f1*mu_1g+mu_f2*mu_2g) ;
term_2q2=-mu_f2*mu_2g*(mu_f2*mu_2g+mu_f1*mu_1g) ;
%% Some scalings
% Lambda: wavelength [in nm] and the corresponding freq [in THz]
% c=299.792458 [nm/fs], freq[THz]=1000*c[nm/fs]/lambda[nm].
% omega[rad/fs] = 2000*pi*freq[THz]
c = 299.792458;   % [nm/fs]
chng_f2nm = @(x) 1000*c./x;  % Change from [THz] to [nm]
%% Prepare the pulse spectrum of the exciting pulse
% freq =frequency [THz], lambda=wavelength [nm]
delf=4.0*delw/Nw ;
freq=(w_0+2.*delw)-(0:Nw-1)*delf ; % frequency mesh [THz] for SI
% omega = freq.*(pi/500.) ;     % Angular frequency mesh [rad/fs] for SI
lambda = chng_f2nm(freq);       % wavelength mesh [nm]
% Create tau mesh and the spectral mesh of echo field
delt=3. ; % Delay time stepsize [fs]
nt=512 ;  % Delay time mesh size
nt12=nt/2;
tmax = delt*nt12;
time_delay = -tmax+(1:nt12)*delt ; % tau ranges from -tmax+delt to 0 fs 
%% Simulation of Echo field E(tau, T=0, w)
E_echo=zeros(Nw,nt);
for i=1:nt12
    tau=-time_delay(i) ; % tau shall be positive for the pulse 1 arriving first
    arg1=term_d1*exp(sqrt(-1)*omega_1*tau-decay*tau);
    arg2=term_d2*exp(sqrt(-1)*omega_2*tau-decay*tau);
    arg3=term_c1*exp(sqrt(-1)*omega_2*tau-decay*tau);
    arg4=term_c2*exp(sqrt(-1)*omega_1*tau-decay*tau);
    arg5=term_2q1*exp(sqrt(-1)*omega_1*tau-decay*tau);
    arg6=term_2q2*exp(sqrt(-1)*omega_2*tau-decay*tau);
    for j=1:Nw
        G_d(j)=1./(freq(j)-w_1+sqrt(-1)*gam) ;
        G_c(j)=1./(freq(j)-w_2+sqrt(-1)*gam) ;
        G_2q1(j)=1./(freq(j)-w_f1+sqrt(-1)*gam) ;
        G_2q2(j)=1./(freq(j)-w_f2+sqrt(-1)*gam) ;
        E_echo(j,i)=(arg1+arg2)*G_d(j)+(arg3+arg4)*G_c(j)+arg5*G_2q1(j)+arg6*G_2q2(j) ;
    end
end
for i=1:nt12
    E_echo(:,i)= 10.*E_echo(:,i)./max(abs(E_echo(:,nt12))) ;
end

E_LO=zeros(Nw,1);
for j=1:Nw
    E_LO(j,1)=100.*exp(-((freq(j)-w_0)/delw)^2) ;
end
%% Plot the spectral profile of the E_echo field at tau=0.
figure(1)
plot(freq',abs(E_echo(:,nt12)),'-r','LineWidth',2);
%axis([freq(1),freq(Nw),min(tdprf),max(tdprf)]);
set(gca,'FontSize',16);
xlabel('\fontsize{18}Frequency(THz)');
ylabel('\fontsize{18}Echo Field (arb. units)');
title('\fontsize{18}\bfEcho field spectrum');
%% Simulation of SI pattern
% The frequency spacing of SI varies from 0.1 to 10 THz, corresponding
           % to a time separation ts [ps] from 50 ps to 0.5 ps. 
ts=6.0 ; % [ps]
Echo=zeros(Nw,nt12);
LO=zeros(Nw,1);
phase=zeros(Nw,1);
inten=zeros(Nw,1);
LO(:,1)=sqrt(E_LO(:,1));
for i=1:nt12
   for j=1:Nw
      phase(j,1)=unwrap(angle(E_echo(j,i))) ;
      intens(j,1)=abs(E_echo(j,i)) ;
      Echo(j,i)=intens(j,1)^2+E_LO(j,1)+2*LO(j,1)*intens(j,1)*cos(phase(j,1)+freq(j)*ts) ;
   end
end
%% Plotting SI at the specified delay times
figure(2)
imagesc(freq,time_delay,Echo');
set(gca,'FontSize',16);
xlabel('\fontsize{18}Emission Frequency (THz)');
ylabel('\fontsize{18}Delay Time (fs)');
title('\fontsize{20}\bfSpectral Interferograms');
colormap(jet) ;
cb=colorbar;
set(cb,'FontSize',16);
%
%% LO output
T_0=0.000 ;
% time_delay[mm] = time_delay[fs].*c/2*10^6 
time_delay = time_delay.*1.4990e-004 ;
fid=fopen(LO_spectrum_output, 'w');
fprintf(fid, '%s', 'Wavelength') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', lambda) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'T-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', T_0) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'Echo-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', T_0) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'LO-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', T_0) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'Intensity') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%f ', E_LO) ;
fprintf(fid, '\n') ;
fclose(fid);
%% Time-delayed SI output
fid=fopen(SI_data_output, 'w');
fprintf(fid, '%s', 'Wavelength') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', lambda) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'Echo-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', time_delay) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'T-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', T_0) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'LO-delay') ;
fprintf(fid, '\n') ;
fprintf(fid, '%f ', T_0) ;
fprintf(fid, '\n') ;
fprintf(fid, '\n') ;
%
fprintf(fid, '%s', 'Intensity') ;
fprintf(fid, '\n') ;
%
for i=1:nt12
    fprintf(fid, '%f ', Echo(:,i)) ;
    fprintf(fid, '\n') ;
end
fclose(fid);
%% tau-FFT of the echo field E_echo(:,i) i=1:nt
time_delay=[];
time_delay = -tmax+(1:nt)*delt ; % from -tmax+delt to tmax
echo_a=zeros(nt,Nw);
% echo_p=zeros(nt,Nw);
Efield=transpose(E_echo);
%
% Assemble the time-dependent echo field spectrum E(tau,freq), in which 
% tau is the time_delay between the first two pulses.
%
% According to Nyquist theorem, at a sampling rate 1/delt, the highest 
% frequency that can be correctly sampled is f_max=500(THz)/delt(fs).
% For delt= 3.34-fs, f_max=150 THz. Any signal with frequency f>fmax will 
% be mapped to f-2*fmax.
fmax=500./delt
dfreq=fmax/nt12
w_ex = -fmax+(0:nt-1)*dfreq ;
for n=1:Nw
    echo_a(:,n)=fftshift(fft(Efield(:,n),nt));
%    echo_p(:,n)= unwrap(angle(echo_a(:,n)));
end
twoD_sigr=zeros(nt,Nw);
twoD_sigm=zeros(nt,Nw);
twoD_yaxis=zeros(nt,1);
% Positive frequency w1 in Echo appears at -w1 in the w_ex mesh because of
% delta(w_ex+w1). However due to the aliasing effect from w1>fmax, the 
% spectral peak shows up at 2*fmax-w_ex. So the real frequench shall become
% -w1=w_ex-2*fmax. 
for ii=1:nt
    twoD_yaxis(ii,1)=w_ex(ii)-2*fmax;
    twoD_sigr(ii,:)=real(echo_a(ii,:));
    twoD_sigm(ii,:)=abs(echo_a(ii,:));
end
f_max=max(freq);
f_min=min(freq);
%% Plotting the 2D data
figure(3)
contourf(freq,twoD_yaxis,twoD_sigm,10);
set(gca,'YDir','normal');
axis([f_min f_max -f_max -f_min],'square');
set(gca,'FontSize',16);
title('\fontsize{18}\bfFFT: Abs(Echo)');
xlabel('\fontsize{18}Emission Frequency (THz)');
ylabel('\fontsize{18}Excitation Frequency (THz)');
colormap(jet) ;
cb=colorbar;
set(cb,'FontSize',16);

figure(4)
contourf(freq,twoD_yaxis,twoD_sigr,10);
set(gca,'YDir','normal');
axis([f_min f_max -f_max -f_min],'square');
set(gca,'FontSize',16);
title('\fontsize{18}\bfFFT: Real(Echo)');
xlabel('\fontsize{18}Emission Frequency (THz)');
ylabel('\fontsize{18}Excitation Frequency (THz)');
colormap(jet) ;
cb=colorbar;
set(cb,'FontSize',16);



    