clear;
clc;

load_const;
load_au;
load_cm;
load_color;

%---- INPUT ----
in.make_plots = 0;
in.recoll_on = 1;                  % Include recollisions at x = 0? yes - 1
in.tunnel_exit = 1;                % Start electrons at tunnel exit (x_te = W/eE) instead of x = 0? yes - 1

in.plot_trajectory = 0;

in.recoll_prob = @(E) 1.0;

% for calc spectrum
in.use_dir = 1;
in.use_recoll = 1;
in.use_phases = 1;
scale = 0.5;
eSMM_spec = 1;

% Set timestep and propagation time list [fs]
dt = 0.0005;
t_list = -3:dt:15;

tbirth = -3:dt:3; 

in.gamma_0 = 7;               % Nearfield enhancement parameters
in.lambda_nf = 100; 

in.U_0   = 0;
in.r_tip = in.lambda_nf;

in.W = 6.6;

pulse.wavelength = 800;
pulse.intensity = 1e12;
pulse.phi = 1.1;
pulse.width = -1;
pulse.t0 = 0;

Pulse = get_laserfield("intensity", pulse.intensity, "phase", pulse.phi, "t0", pulse.t0, "tlist", t_list, "wavelength", pulse.wavelength, "width", pulse.width);
cycles = Pulse.cycles;
Up = Pulse.Up*in.gamma_0^2;
keldysh = sqrt(in.W / (2*Pulse.Up*in.gamma_0^2))

%% Compare Near-field profile
d = 0:600;
Espec = 0:0.005*Up:30*Up;

in.lambda_nf = 100;
pulse.intensity = 1e12;
pulse.wavelength = 800;

near_field_list = 0:1; % 0 - exponential decay, 1 - power law decay
for idx = 1:numel(near_field_list)
    near_field_type = near_field_list(idx);

    [traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth, near_field_power_law=near_field_type);
    
    if near_field_type
        R = in.lambda_nf / 0.4;
        w_foc = 8000;
        a = 2*log(2)/(2*w_foc)^2;
    
        gamma = @(x) (x>=0).*((in.gamma_0-1)*(R./(x+R)).^3 + exp(-2*log(2)*x.^2 /(2*w_foc)^2));
    
    else
        gamma = @(x) (x>=0).*(1 + (in.gamma_0-1)*exp(-x/in.lambda_nf));
    end

    Up = Pulse.Up*in.gamma_0^2;
    Espec = linspace(0,30*Up,5000);

    spectrum = calc_spec_with_eff_phase(Espec, pulse.phi, traj, weight_scale = 1, scale = 0.5);
    
    Espec_NF(:,idx) = (Espec + traj.U_0)/Up;
    spec_NF(:,idx) = spectrum.all;
    near_field(:,idx) = gamma(d);
    co_dir_nf(idx) = max(traj.Ekin(traj.roi_direct)/Up);
    co_rec_nf(idx) = max(traj.Ekin(traj.roi_recoll)/Up);

end

%% Compare Decay lengths
Espec = 0:0.005*Up:30*Up;

pulse.intensity = 1e12;
pulse.wavelength = 800;

decay_list = [200, 50];
for idx = 1:numel(decay_list)
    in.lambda_nf = decay_list(idx);

    [traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth);
    
    Up = Pulse.Up*in.gamma_0^2;
    Espec = linspace(0,10*Up,5000);
    
    spectrum = calc_spec_with_eff_phase(Espec, pulse.phi, traj, weight_scale = 1, scale = 0.5);
    
    Espec_decay(:,idx) = (Espec + traj.U_0)/Up;
    spec_decay(:,idx) = spectrum.all;
    co_dir_dec(idx) = max(traj.Ekin(traj.roi_direct)/Up);
    co_rec_dec(idx) = max(traj.Ekin(traj.roi_recoll)/Up);
    
end

%% Compare intensities

Espec = 0:0.005*Up:30*Up;

in.lambda_nf = 100;
pulse.wavelength = 800;

intensity_list = [2.5e12, 5e11];
for idx = 1:numel(intensity_list)
    pulse.intensity = intensity_list(idx);

    [traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth);
    
    Up = Pulse.Up*in.gamma_0^2;
    Espec = linspace(0,10*Up,5000);

    spectrum = calc_spec_with_eff_phase(Espec, pulse.phi, traj, weight_scale = 1, scale = 0.5);
    
    Espec_intensity(:,idx) = (Espec + traj.U_0)/Up;
    spec_intensity(:,idx) = spectrum.all;
    co_dir_int(idx) = max(traj.Ekin(traj.roi_direct)/Up);
    co_rec_int(idx) = max(traj.Ekin(traj.roi_recoll)/Up);
    
end

%% Compare wavelengths

in.lambda_nf = 100;
pulse.intensity = 1e12;

wavelength_list = [1100, 500];
for idx = 1:numel(wavelength_list)
    pulse.wavelength = wavelength_list(idx);

    [traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth);
    
    Up = Pulse.Up*in.gamma_0^2;
    Espec = linspace(0,10*Up,5000);
        
    spectrum = calc_spec_with_eff_phase(Espec, pulse.phi, traj, weight_scale = 1, scale = 0.5);
    
    Espec_wavelength(:,idx) = (Espec + traj.U_0)/Up;
    spec_wavelength(:,idx) = spectrum.all;
    co_dir_wavelength(idx) = max(traj.Ekin(traj.roi_direct)/Up);
    co_rec_wavelength(idx) = max(traj.Ekin(traj.roi_recoll)/Up);
       
end

%% Plot

figure('Color', 'w', 'Units', 'normalized', 'Position', [0.1 0.1 0.8 0.6])

col_tdse = [my_Green; 0.5 0.9 0.7];
col_smm = [0 0 0; 0.5 0.5 0.5];

for idx = 1:numel(near_field_list)
    Espec = Espec_NF(:,idx);
    
    subplot(3,3,2)
    plot(co_dir_nf(idx)*[1 1], [-10 2], 'Color', my_B)
    hold on
    plot(co_rec_nf(idx)*[1 1], [-10 2], 'Color', my_R)
    
    % plot(out.E_list/Up, log10(out.spec), 'Color', my_Green, 'LineWidth', 1.5)
    
    plot(Espec(Espec<co_dir_nf(idx)), log10(abs(spec_NF((Espec<co_dir_nf(idx)),idx)).^2)-2, 'Color', col_smm(idx,:), 'LineWidth', 1.5)
    plot(Espec(Espec>co_dir_nf(idx)), log10(abs(spec_NF((Espec>co_dir_nf(idx)),idx)).^2)-2, '--', 'Color', col_smm(idx,:), 'LineWidth', 1.5)

    ylabel 'log10(Yield)'
    xlim([0 4])
    ylim([-9 0])
    set(gca, 'FontSize', 12)

    subplot(3,3,3)
    plot(d, near_field(:,idx), 'LineWidth', 1.5, 'Color', col_tdse(idx,:))
    hold on
    xlabel 'distance from surface (Å)'
    ylabel 'enhancement'
    set(gca, 'FontSize', 12)
end

for idx = 1:numel(decay_list)
    Espec = Espec_decay(:,idx);
    
    subplot(3,3,(idx-1)*3+4)
    plot(co_dir_dec(idx)*[1 1], [-10 2], 'Color', my_B)
    hold on
    plot(co_rec_dec(idx)*[1 1], [-10 2], 'Color', my_R)
    
    % plot(out.E_list/Up, log10(out.spec), 'Color', my_Green, 'LineWidth', 1.5)

    plot(Espec(Espec<co_dir_dec(idx)), log10(abs(spec_decay((Espec<co_dir_dec(idx)),idx)).^2)-2, 'Color', col_smm(1,:), 'LineWidth', 1.5)
    plot(Espec(Espec>co_dir_dec(idx)), log10(abs(spec_decay((Espec>co_dir_dec(idx)),idx)).^2)-2, '--', 'Color', col_smm(1,:), 'LineWidth', 1.5)

    ylabel 'log10(Yield)'
    xlim([0 4])
    ylim([-9 0])
    set(gca, 'FontSize', 12)
end

for idx = 1:numel(intensity_list)
    Espec = Espec_intensity(:,idx);
    
    subplot(3,3,(idx-1)*3+5)
    plot(co_dir_int(idx)*[1 1], [-10 2], 'Color', my_B)
    hold on
    plot(co_rec_int(idx)*[1 1], [-10 2], 'Color', my_R)
    
    % plot(out.E_list/Up, log10(out.spec), 'Color', my_Green, 'LineWidth', 1.5)

    plot(Espec(Espec<co_dir_int(idx)), log10(abs(spec_intensity((Espec<co_dir_int(idx)),idx)).^2)-2, 'Color', col_smm(1,:), 'LineWidth', 1.5)
    plot(Espec(Espec>co_dir_int(idx)), log10(abs(spec_intensity((Espec>co_dir_int(idx)),idx)).^2)-2, '--', 'Color', col_smm(1,:), 'LineWidth', 1.5)

    xlim([0 4])
    ylim([-9 0])
    set(gca, 'FontSize', 12)
end

for idx = 1:numel(wavelength_list)
    Espec = Espec_wavelength(:,idx);
    
    subplot(3,3,(idx-1)*3+6)
    plot(co_dir_wavelength(idx)*[1 1], [-10 2], 'Color', my_B)
    hold on
    plot(co_rec_wavelength(idx)*[1 1], [-10 2], 'Color', my_R)
    
    % plot(out.E_list/Up, log10(out.spec), 'Color', my_Green, 'LineWidth', 1.5)

    plot(Espec(Espec<co_dir_wavelength(idx)), log10(abs(spec_wavelength((Espec<co_dir_wavelength(idx)),idx)).^2)-2, 'Color', col_smm(1,:), 'LineWidth', 1.5)
    plot(Espec(Espec>co_dir_wavelength(idx)), log10(abs(spec_wavelength((Espec>co_dir_wavelength(idx)),idx)).^2)-2, '--', 'Color', col_smm(1,:), 'LineWidth', 1.5)

    xlim([0 4])
    ylim([-9 0])
    set(gca, 'FontSize', 12)
end

subplot(3,3,7)
xlabel 'Energy (Up)'
subplot(3,3,8)
xlabel 'Energy (Up)'
subplot(3,3,9)
xlabel 'Energy (Up)'

%% add TDSE comparison
load('TDSE_comparison.mat')

subplot(3,3,2)
plot(tdse.E_exp, log10(tdse.spec_exp), 'LineWidth', 1.5, 'Color', col_tdse(1,:))
plot(tdse.E_nf, log10(tdse.spec_nf), 'LineWidth', 1.5, 'Color', col_tdse(2,:))

subplot(3,3,4)
plot(tdse.E_lnf200, log10(tdse.spec_lnf200), 'LineWidth', 1.5, 'Color', col_tdse(1,:))

subplot(3,3,5)
plot(tdse.E_int12, log10(tdse.spec_int12), 'LineWidth', 1.5, 'Color', col_tdse(1,:))

subplot(3,3,6)
plot(tdse.E_wave1100, log10(tdse.spec_wave1100), 'LineWidth', 1.5, 'Color', col_tdse(1,:))

subplot(3,3,7)
plot(tdse.E_lnf50, log10(tdse.spec_lnf50), 'LineWidth', 1.5, 'Color', col_tdse(1,:))

subplot(3,3,8)
plot(tdse.E_int11, log10(tdse.spec_int11), 'LineWidth', 1.5, 'Color', col_tdse(1,:))

subplot(3,3,9)
plot(tdse.E_wave500, log10(tdse.spec_wave500), 'LineWidth', 1.5, 'Color', col_tdse(1,:))
