clear all;
close all;
clc;

format long g;

%% Check the resonance configurations in wavelength and frequency

%%% Wg-definitions
alpha = 0.995; 
c = 299792458;
Tera = 1e12;
um = 1e-6;
mm = 1e-3;
lambda0_neff = 1.55; % micron 
n1 = 2.4473; n2 = -1.1327; n3 = -0.0440;
param_neff = [n1 n2 n3];
n_eff = @(param, lambda) (param(1) + param(2)*(lambda - lambda0_neff) + param(3)*(lambda - lambda0_neff).^2);
n_g = @(param, lambda) ( param(1) - param(2)*lambda0_neff - param(3)*(lambda.^2 - lambda0_neff.^2) );
n_eff_omega = @(param, omega) (n_eff(param, 2*pi*c./omega/Tera/um));
n_g_omega = @(param, omega) (n_g(param, 2*pi*c./omega/Tera/um));
L_wg = 14*mm;

%%% Pump-definitions
pump_envelope_lambda = @(param, lambda) (sech(((lambda - param(1))./(param(2)))));
pump_envelope_omega = @(param, omega) (sech(((omega - param(1))./(param(2)))));

%%%% DWDM channels definitions: Intensity profiles
DWDM_filtering_omega = @(param, omega) ( exp(-abs(omega - param(1)).^8/2/param(2).^8) );

%%%% Tuneable filter definitions: Intensity profiles
TF_filter_lambda = @(param, lambda) ( exp(-abs(lambda - param(1)).^8/2/param(2).^8) );
TF_filter_omega = @(param, omega) ( exp(-abs(omega - param(1)).^8/2/param(2).^8) );


%%% Dispersion Parameters
lambda0 = 1546.12;
omega0 = 2*pi*c/lambda0/1e3; % Tera rad/s
lambda_fit = linspace(1538,1556,10000);
omega_fit = linspace((2*pi*c/1556/1e3),(2*pi*c/1538/1e3), 10000); % 1 pm = 0.129 GHz
Emax = 1;
ng_calc = n_g(param_neff, lambda_fit/1000);
D_calc = diff(ng_calc)./diff(lambda_fit)*1e9/c * 1e6; % the unit becomes ps/(nm.km)
GVD = -(lambda_fit(2:end).^2*1e-18)./(2*pi*c).*D_calc.*1e-6*1e30; % the unit becomes fs^2/m
D_lambda0 = interp1(lambda_fit(2:end), D_calc, lambda0);
GVD_lambda0 = interp1(lambda_fit(2:end), GVD, lambda0)*1e-30;
delta_k = ((omega_fit - omega0)*Tera).^2*GVD_lambda0*3;
Phase_matching = sinc(delta_k*L_wg/2).^2;

%%%% Pump
lambda_p0 = lambda0;
omega_p0 = omega0;
SigmaP_lambda = 1.104/1.763; % For solitonic sech intensity FWHM = 1.763*sigma
SigmaP_omega =  2*pi*138*1e-3/1.763; % T rad/s
param_pump_lambda = [lambda_p0/1000 SigmaP_lambda/1000];
alphaP_lambda = pump_envelope_lambda(param_pump_lambda, lambda_fit/1000);
[PKS5,LOCS5,W_p,P5] = findpeaks(abs(alphaP_lambda).^2, lambda_fit, 'WidthReference', 'halfprom');
fprintf('The FWHM of the pump envelope is %f pm.\n', W_p*1000);
param_pump_omega = [omega_p0 SigmaP_omega];
alphaP_omega = pump_envelope_omega(param_pump_omega, omega_fit);
[PKS6,LOCS6,W_p_f,P6] = findpeaks(abs(alphaP_omega).^2, omega_fit, 'WidthReference', 'halfprom');
fprintf('The FWHM of the pump envelope is %f GHz.\n', W_p_f*1000/(2*pi));

Pump_data = csvread('Pritel_FROG_01March2019.csv', 1, 0);
P_freq = Pump_data(:, 1) - 0.02; % to match the measurement discrepencies between instruments
P_amp = Pump_data(:, 2);
P_phase = Pump_data(:, 3);
P_phase = P_phase - mean(P_phase);

figure,
yyaxis left;
plot(P_freq, P_amp.^2, 'color', 'k');
set(gca, 'YColor', 'k');
xlabel('\nu (THz)'); ylabel('Intensity (arb. units)'); 
yyaxis right;
plot(P_freq, P_phase, 'color', 'b');
set(gca, 'YColor', 'b');
ylabel('Phase (\pi^c)'); title('Spectral profile');
% xlim([193 195]);
set(gca, 'box', 'on', 'linewidth', 1.0); 


%%%% DWDM channels
lambda_i0 = 1552.52; % nm
lambda_s0 = 1539.77; % nm
omega_i0 = 2*pi*c/lambda_i0/1e3; % Tera rad/s
omega_s0 = 2*pi*c/lambda_s0/1e3; % Tera rad/s
sigma_DWDM_lambda = 1.248/2.083; % nm
sigma_DWDM_omega = 2*pi*155*1e-3/2.083; % T rad/s = 2*pi*155 GHz at 1552.52 nm
param_DWDM_idler = [omega_i0 sigma_DWDM_omega];
param_DWDM_signal = [omega_s0 sigma_DWDM_omega];
param_DWDM_pump = [omega_p0 sigma_DWDM_omega];
Filter_idler_lambda = abs(Emax)^2*exp(-abs(lambda_fit - lambda_i0).^8/2/sigma_DWDM_lambda.^8);
Filter_signal_lambda = abs(Emax)^2*exp(-abs(lambda_fit - lambda_s0).^8/2/sigma_DWDM_lambda.^8);
Filter_pump_lambda = abs(Emax)^2*exp(-abs(lambda_fit - lambda_p0).^8/2/sigma_DWDM_lambda.^8);
Filter_idler_omega = abs(Emax)^2*DWDM_filtering_omega(param_DWDM_idler, omega_fit);
Filter_signal_omega = abs(Emax)^2*DWDM_filtering_omega(param_DWDM_signal, omega_fit);
Filter_pump_omega = abs(Emax)^2*DWDM_filtering_omega(param_DWDM_pump, omega_fit);

%%%% Tuneable filter 
sigma_TF_lambda = 0.224/2.083; % nm
sigma_TF_omega =  2*pi*28.1*1e-3/2.083; % T rad/s
param_TF_i_lambda = [lambda_i0 sigma_TF_lambda];
param_TF_i_omega = [omega_i0 sigma_TF_omega];
TF_i_lambda = TF_filter_lambda(param_TF_i_lambda, lambda_fit);
TF_i_omega = TF_filter_omega(param_TF_i_omega, omega_fit);

figure,
set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0);
hold on;
plot(omega_fit, abs(alphaP_omega).^2, 'linewidth', 1.0, 'color', 'g');
plot(P_freq*2*pi, P_amp.^2, 'linewidth', 1.0, 'color', 'r');
plot(omega_fit, TF_i_omega, 'linewidth', 1.0, 'color', 'b');
plot(omega_fit, Filter_idler_omega, 'linewidth', 1.0, 'color', 'k');
plot(omega_fit, Filter_signal_omega, 'linewidth', 1.0, 'color', 'k');
plot(omega_fit, Filter_pump_omega, 'linewidth', 1.0, 'color', 'k');
hold off;
xlabel('\omega (T^c/s)'); ylabel('Intensity (arb. units)');
legend('Pump envelope', 'TF signal', 'DWDM channels');


figure,
set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0);
hold on;
plot(lambda_fit, abs(alphaP_lambda).^2, 'linewidth', 1.0, 'color', 'r');
plot(lambda_fit, TF_i_lambda, 'linewidth', 1.0, 'color', 'b');
plot(lambda_fit, Filter_idler_lambda, 'linewidth', 1.0, 'color', 'k');
plot(lambda_fit, Filter_signal_lambda, 'linewidth', 1.0, 'color', 'k');
plot(lambda_fit, Filter_pump_lambda, 'linewidth', 1.0, 'color', 'k');
hold off;
xlabel('\lambda (nm)'); ylabel('Intensity (arb. units)');
legend('Pump envelope', 'TF signal', 'DWDM channels');

figure,
set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0);
hold on;
plot(lambda_fit, abs(alphaP_lambda).^2, 'linewidth', 1.0, 'color', 'r');
plot(2*pi*c./omega_fit/Tera/1e-9, abs(alphaP_omega).^2, 'linewidth', 1.0, 'color', 'g', 'linestyle', ':');
plot(lambda_fit, TF_i_lambda, 'linewidth', 1.0, 'color', 'b');
plot(2*pi*c./omega_fit/Tera/1e-9, TF_i_omega, 'linewidth', 1.0, 'color', 'g', 'linestyle', 'none', 'marker', '.', 'markersize', 6);
% plot(lambda_fit, Filter_idler_lambda, 'linewidth', 1.0, 'color', 'r');
% plot(lambda_fit, Filter_signal_lambda, 'linewidth', 1.0, 'color', 'k');
plot(lambda_fit, Filter_pump_lambda, 'linewidth', 0.50, 'color', 'b');
plot(2*pi*c./omega_fit/Tera/1e-9, Filter_pump_omega, 'color', 'k', 'linestyle', 'none', 'marker', '.', 'markersize', 6);
hold off;
xlabel('\lambda (nm)'); ylabel('Intensity (arb. units)'); title('Comparing frequency to wavelength conversion');
legend('Pump envelope in \lambda', 'Pump envelope in \omega', 'TF \lambda', 'TF \omega', 'DWDM channels');
% xlim([1544 1548]);

r=corrcoef(abs(alphaP_omega).^2, fliplr(abs(alphaP_lambda).^2));
r2=r(1,2).^2;
disp (['Goodness of fit, r^2 value: ' num2str(r2) ])

% return

%% Plotting JSA
res_omega = 0.004; % Tera rad/s (0.00393 T rad/s = 0.005 pm)
omega_rang = 1.989; % Tera rad/s
omega_s = ((omega_s0-omega_rang):res_omega:(omega_s0+omega_rang));
omega_i = ((omega_i0-omega_rang):res_omega:(omega_i0+omega_rang));
omega_p = ((omega_p0-omega_rang):res_omega:(omega_p0+omega_rang));
Ns = length(omega_s);
Ni = length(omega_i);
Np = length(omega_p);
A1 = zeros(Ns,Ni);
A2 = zeros(Ns,Ni);

% Pump: SigmaP in T rad/s; 49.6 GHz is 395.626 pm; 138 GHz is 1.104 nm.
% SigmaP_omega = 2*pi*138*1e-3/1.763/(1); %2*pi*138*1e-3; % 2*pi*49.6*1e-3;  /(2*sqrt(2))
% param_pump_omega = [omega_p0 SigmaP_omega];
% alphaP_omega = pump_envelope_omega(param_pump_omega, omega_p);
% [PKS6,LOCS6,W_p_dash,P6] = findpeaks(abs(P_amp).^2, P_freq*2*pi, 'WidthReference', 'halfprom');
% fprintf('The FWHM of the pump envelope is %f GHz.\n', W_p_dash*1000/(2*pi));
alphaP1 = interp1(P_freq, P_amp, omega_p/(2*pi)).*exp(1i*interp1(P_freq, P_phase, omega_p/(2*pi)));

figure,
yyaxis left;
plot(omega_p/(2*pi), abs(alphaP1).^2, 'color', 'k');
set(gca, 'YColor', 'k');
xlabel('\nu (THz)'); ylabel('Intensity (arb. units)'); 
yyaxis right;
plot(omega_p/(2*pi), interp1(P_freq, P_phase, omega_p/(2*pi)), 'color', 'b');
set(gca, 'YColor', 'b');
ylabel('Phase (\pi^c)'); title('Spectral profile');
% xlim([193 195]);
set(gca, 'box', 'on', 'linewidth', 1.0); 

% DWDM: omega_i0, omega_s0 and omega_p0 are DWDM's center frequency
param_DWDM_idler = [omega_i0 sigma_DWDM_omega];
param_DWDM_signal = [omega_s0 sigma_DWDM_omega];
param_DWDM_pump = [omega_p0 sigma_DWDM_omega];

% Phase contribution
Phase_max = linspace(0, 1, 2);
N_P = length(Phase_max);
figure,
for kp = 1:N_P
    
    alphaP1 = interp1(P_freq, P_amp, omega_p/(2*pi)).*exp(-1i*interp1(P_freq, P_phase, omega_p/(2*pi))*Phase_max(kp));
   
    yyaxis left; hold on;
    plot(omega_p/(2*pi), abs(alphaP1).^2, 'color', 'k');
    
    yyaxis right; hold on;
    plot(omega_p/(2*pi), angle(alphaP1), 'color', 'b');
    
end
yyaxis left; hold off; set(gca, 'YColor', 'k'); xlabel('\nu (THz)'); ylabel('Intensity (arb. units)');
yyaxis right; hold off; set(gca, 'YColor', 'b'); ylabel('Phase (\pi^c)'); title('Spectral profile');
set(gca, 'box', 'on', 'linewidth', 1.0);

% TF configuration
FWHM_TF_nm = linspace(0.05, 1.0, 20);
FWHM_TF_GHz = c*FWHM_TF_nm/lambda0.^2;
sigma_TF = 2*pi*FWHM_TF_GHz*1e-3/2.083; % 2*pi*154e-3;  % T rad/s
N_TF = length(sigma_TF);
Purity1 = zeros(N_P, N_TF);
Purity2 = zeros(N_P, N_TF);
eig_Schmidt1 = zeros(N_P, N_TF, 10);
eig_Schmidt2 = zeros(N_P, N_TF, 10);

% return

for kp = 1:N_P
    
    for k_TF = 1:N_TF
        
        param_TF_s = [omega_s0 sigma_TF(k_TF)];
        TF_s = TF_filter_omega(param_TF_s, omega_s);
        TF_s_A = sqrt(TF_s);
        [PKSTF,LOCSTF,W_TF,PTF] = findpeaks(TF_s, omega_s, 'WidthReference', 'halfprom');
        fprintf('TF = %f GHz.\n', W_TF*1e3/(2*pi));
        
        parfor ki = 1:Ni
            
            for ks = 1:Ns
                
                %%% Energy conservation
                omega_p_dash = omega_s(ks) + omega_i(ki) - omega_p;
                
                %%% Momentum conservation / phase matching with dispersion
                delta_k = Tera*( n_g_omega(param_neff, omega_p_dash).*omega_p_dash +  n_g_omega(param_neff, omega_p).*omega_p...
                    -  n_g_omega(param_neff, omega_s(ks)).*omega_s(ks) - n_g_omega(param_neff, omega_i(ki)).*omega_i(ki))/c;
                Phase_matching = sinc(delta_k*L_wg/2).^2;
               
                
                %%% Without any DWDM filtering
                alphaP1 = interp1(P_freq, P_amp, omega_p/(2*pi)).*exp(-1i*interp1(P_freq, P_phase, omega_p/(2*pi))*Phase_max(kp));
                alphaP2 = interp1(P_freq, P_amp, omega_p_dash/(2*pi)).*exp(-1i*interp1(P_freq, P_phase, omega_p_dash/(2*pi))*Phase_max(kp));
                A0_w = alphaP1.*alphaP2.*Phase_matching;
                A1(ks, ki) = trapz(omega_p, A0_w);
                
                %%% With DWDM + TF filtering
                A0_w_DWDM = A0_w.*sqrt(DWDM_filtering_omega(param_DWDM_pump, omega_p).*DWDM_filtering_omega(param_DWDM_pump, omega_p_dash)...
                    .*DWDM_filtering_omega(param_DWDM_signal, omega_s(ks))).*TF_s_A(ks);
                A2(ks, ki) = trapz(omega_p, A0_w_DWDM);
                
                
            end
            
        end
        
        rho_signal1 = A1*transpose(conj(A1));
        rho_signal1 = rho_signal1/trace(rho_signal1);
        Purity1(kp, k_TF) = trace(rho_signal1*rho_signal1);
        
        [U1,S1,V1] = svd(rho_signal1);
        diag_S = sqrt(diag(S1))/sqrt(sum(diag(S1)));
        eig_Schmidt1(kp, k_TF, :) = diag_S(1:10);
        
        rho_signal2 = A2*transpose(conj(A2));
        rho_signal2 = rho_signal2/trace(rho_signal2);
        Purity2(kp, k_TF) = trace(rho_signal2*rho_signal2);
        
        [U2,S2,V2] = svd(rho_signal2);
        diag_S = sqrt(diag(S2))/sqrt(sum(diag(S2)));
        eig_Schmidt2(kp, k_TF, :) = diag_S(1:10);
        
        
        figure(5),
        om_p_r = omega_rang;
        subplot(4, 5, k_TF);
        mesh(omega_i/(2*pi), omega_s/(2*pi), abs(A1)/max(max(abs(A1)))); view(2); colorbar;
        ylim([omega_s0-om_p_r omega_s0+om_p_r]/(2*pi)); xlim([omega_i0-om_p_r omega_i0+om_p_r]/(2*pi));
        title(['P = ' num2str(Purity1(kp,k_TF), '%.2f')]);
        set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
        
        figure(6),
        om_p_r = omega_rang/3;
        subplot(4, 5, k_TF);
        mesh((omega_s)/(2*pi), (omega_i)/(2*pi), abs(A2)/max(max(abs(A2)))); view(2); colorbar;
        xlim([omega_s0-om_p_r omega_s0+om_p_r]/(2*pi)); ylim([omega_i0-om_p_r omega_i0+om_p_r]/(2*pi)); zlim([-5 5]);
        title(['P = ' num2str(Purity2(kp,k_TF), '%.2f')]);
        set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
        
        figure(7),
        om_p_r = omega_rang/3;
        subplot(4, 5, k_TF);
        mesh((omega_s)/(2*pi), (omega_i)/(2*pi), angle(A2)); view(2); colorbar;
        xlim([omega_s0-om_p_r omega_s0+om_p_r]/(2*pi)); ylim([omega_i0-om_p_r omega_i0+om_p_r]/(2*pi));
        title(['P = ' num2str(Purity2(kp,k_TF), '%.2f')]);
        set(gca, 'FontName', 'Calibri', 'FontSize', 12, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
        
        
    end
    
end

figure,
subplot(221),
set(gca, 'FontName', 'Calibri', 'FontSize', 10, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
hold on;
plot(omega_s/(2*pi), V1(:,1), 'LineWidth', 1.0, 'color', 'k');
plot(omega_s/(2*pi), V1(:,2), 'LineWidth', 1.0, 'color', 'b');
plot(omega_s/(2*pi), V1(:,3), 'LineWidth', 1.0, 'color', 'r');
hold off;
legend('Mode 1', 'Mode 2', 'Mode 3');
xlim([omega_s(1) omega_s(end)]/(2*pi));
xlabel('f_s (THz)'); ylabel('Amplitude (arb. units)');

subplot(222),
set(gca, 'FontName', 'Calibri', 'FontSize', 10, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
hold on;
plot(omega_s/(2*pi), V2(:,1), 'LineWidth', 1.0, 'color', 'k');
plot(omega_s/(2*pi), V2(:,2), 'LineWidth', 1.0, 'color', 'b');
plot(omega_s/(2*pi), V2(:,3), 'LineWidth', 1.0, 'color', 'r');
hold off;
legend('Mode 1', 'Mode 2', 'Mode 3');
xlim([omega_s0-om_p_r omega_s0+om_p_r]/(2*pi)); ylim([-0.2 0.2]);
xlabel('f_s (THz)'); ylabel('Amplitude (arb. units)');

subplot(223),
bar((1:length(eig_Schmidt1(1,:,1))), eig_Schmidt1(1,:,:));
set(gca, 'FontName', 'Calibri', 'FontSize', 10, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
xlabel('Schmidt mode #'); ylabel('\lambda_k');
xlim([1 10]); ylim([0 1]);
 
subplot(224),
bar((1:length(eig_Schmidt2(1,:,1))), eig_Schmidt2(1,:,:));
set(gca, 'FontName', 'Calibri', 'FontSize', 10, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
xlabel('Schmidt mode #'); ylabel('\lambda_k');
xlim([1 10]); ylim([0 1]);

figure,
set(gca, 'FontName', 'Calibri', 'FontSize', 10, 'LineWidth', 1.0, 'box', 'on', 'Color', 'none');
hold on;
plot(sigma_TF/(2*pi)*1e3, Purity1(1,:,:), 'marker', '.', 'markersize', 30, 'linewidth', 1.0);
plot(sigma_TF/(2*pi)*1e3, Purity2(1,:,:), 'marker', '.', 'markersize', 30, 'linewidth', 1.0, 'linestyle', '--');
xlabel('TF FWHM (GHz)'); ylabel('Purity'); legend('w/o DWDM', 'w/ DWDM');
hold off;

