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; 
tbirth_cycles = 0; % 0 - use tbirth from above

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))

%% Propagation

if tbirth_cycles > 0
    N_left = ceil(tbirth_cycles/2);
    N_right = tbirth_cycles - N_left;
    tstart = (2-pulse.phi)*pi/Pulse.w + Pulse.t0;
    tbirth = tstart-N_left*Pulse.T:dt:tstart+N_right*Pulse.T;
end

% Propagation
[traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth);

roi = traj.roi_direct|traj.roi_recoll;
min_rate = 1e-10*max(traj.rate);
E_cutoff = max(traj.Ekin(traj.rate > min_rate));
Ekin_list(traj.roi_tbirth) = traj.Ekin;
rate_list(traj.roi_tbirth) = traj.rate;
S_list(traj.roi_tbirth) = traj.S.tot;
roi_direct_list(traj.roi_tbirth) = traj.roi_direct;
roi_recoll_list(traj.roi_tbirth) = traj.roi_recoll;
Ebirth_list(traj.roi_tbirth) = traj.Ebirth;

mean_cutoff = mean(E_cutoff);

%% Calc fits and spectrum 

Espec = 0:0.005*Up:30*Up;
Espec_hist = (0:0.05:20)*Up;

in.Up = Up; in.Ephot = Pulse.Ephot;
spec_SMM = histw_phase(traj.Ekin, traj.rate, 0*traj.S.tot, Espec_hist);
spectrum = calc_spec_with_eff_phase(Espec, pulse.phi, traj, weight_scale = 1, scale = 0.5);

Espec = Espec + traj.U_0;
%%
spec_plot_all = spectrum.all;
spec_plot_dir = spectrum.direct;
spec_plot_rec = spectrum.recoll;
total_relative_phase = unwrap(angle(spec_plot_rec)-angle(spec_plot_dir));

direct_list = find(traj.roi_direct);
recoll_list = find(traj.roi_recoll);

[co_dir, idx_co_dir] = max(traj.Ekin(direct_list)/Up);
[co_rec, idx_co_rec] = max(traj.Ekin(recoll_list)/Up);
max_co = min(co_dir,co_rec);

% Phase function
theta_rec = spectrum.Theta_Fit_list_rec(:,1,1)/pi + spectrum.offset_rec(1,1);
theta_dir = spectrum.Theta_Fit_list_dir(:,1,1)/pi + spectrum.offset_dir(1,1);

% Weights, phases and energies of trajectories
weight_d = traj.rate(direct_list)/max(traj.rate);
weight_r = traj.rate(recoll_list)/max(traj.rate);
Edirplot = traj.Ekin(direct_list)/Up;
Sdirplot = traj.S.tot(direct_list)/const.hbar/pi;
wdirplot = 5+15*weight_d;
roi_d = 1:25:numel(Edirplot);
Erplot = traj.Ekin(recoll_list)/Up;
Srplot = traj.S.tot(recoll_list)/const.hbar/pi;
wrecplot = 5+15*weight_r;
roi_r = 1:25:numel(Erplot);

% energy list for effective phase (as in function)
Espec_hist_dir = linspace(0, 1.1*co_dir*Up, 200);
roi_dir = Espec_hist_dir < co_dir*Up;
Espec_hist_rec = linspace(0, 1.1*co_rec*Up, 200);
roi_rec = Espec_hist_rec < co_rec*Up;

% define equidistant dots
wE = 2.5/10;
wS = 15/4;
% for recollision
dS = diff(Srplot)*wE;
dE = diff(Erplot)*wS;
delta = sqrt(dE.^2 + dS.^2);
delta_cum = cumsum(delta);
Delta_sum = sum(delta);
delta_soll = 0.25;
delta_soll_list = 0:delta_soll:Delta_sum;
E_soll = interp1(delta_cum, Erplot(1:end-1), delta_soll_list);
S_soll = interp1(delta_cum, Srplot(1:end-1), delta_soll_list);
w_soll = interp1(delta_cum, weight_r(1:end-1), delta_soll_list);
% for direct
dS_d = diff(Sdirplot)*wE;
dE_d = diff(Edirplot)*wS;
delta_d = sqrt(dE_d.^2 + dS_d.^2);
delta_cum_d = cumsum(delta_d);
Delta_sum_d = sum(delta_d);
delta_soll_list_d = 0:delta_soll:Delta_sum_d;
E_solld = interp1(delta_cum_d, Edirplot(1:end-1), delta_soll_list_d);
S_solld = interp1(delta_cum_d, Sdirplot(1:end-1), delta_soll_list_d);
w_solld = interp1(delta_cum_d, weight_d(1:end-1), delta_soll_list_d);

% calculate contrasts
E_dir = traj.Ekin(traj.roi_direct);
E_rec = traj.Ekin(traj.roi_recoll);
S_dir = traj.S.tot(traj.roi_direct);
S_rec = traj.S.tot(traj.roi_recoll);
rate_dir = traj.rate(traj.roi_direct);
rate_rec = traj.rate(traj.roi_recoll);

roi_dir_1 = rate_dir > rate_dir(idx_co_dir);
Dir_1 = histw_phase(E_dir(roi_dir_1), rate_dir(roi_dir_1), S_dir(roi_dir_1), Espec_hist_dir);
Dir_2 = histw_phase(E_dir(~roi_dir_1), rate_dir(~roi_dir_1), S_dir(~roi_dir_1), Espec_hist_dir);
contrast_dir = (2*abs(Dir_1).*abs(Dir_2))./(abs(Dir_1).^2 + abs(Dir_2).^2);

[~, idx_dir_contrast] = min(abs(contrast_dir-0.1));

roi_rec_1 = rate_rec > rate_rec(idx_co_rec);
Rec_1 = histw_phase(E_rec(roi_rec_1), rate_rec(roi_rec_1), S_rec(roi_rec_1), Espec_hist_rec);
Rec_2 = histw_phase(E_rec(~roi_rec_1), rate_rec(~roi_rec_1), S_rec(~roi_rec_1), Espec_hist_rec);
contrast_rec = (2*abs(Rec_1).*abs(Rec_2))./(abs(Rec_1).^2 + abs(Rec_2).^2);
[~, idx_rec_contrast] = min(abs(contrast_rec-0.1));

%%
figure('Color', 'w', 'Units', 'normalized', 'Position', [0.3 0.1 0.3 0.8])
subplot(3,1,1)
plot(co_dir*[1 1], [-10 50], 'Color', my_B)
hold on
plot(co_rec*[1 1], [-10 50], 'Color', my_R)
plot(Espec_hist_dir(idx_dir_contrast)/Up*[1 1], [-10 50], 'Color', my_G)
plot(Espec_hist_rec(idx_rec_contrast)/Up*[1 1], [-10 50], 'Color', my_G)
scatter(E_soll, S_soll, 3+50*w_soll, my_R, 'filled')
scatter(E_solld, S_solld, 3+50*w_solld, my_B, 'filled')
plot(Espec(Espec>co_rec*Up)/Up, theta_rec(Espec>co_rec*Up), '--', 'Color', my_G, 'LineWidth', 1.5)
plot(Espec(Espec<co_rec*Up)/Up, theta_rec(Espec<co_rec*Up), 'Color', my_G, 'LineWidth', 1.5)
plot(Espec(Espec>co_dir*Up)/Up, theta_dir(Espec>co_dir*Up), '--', 'Color', my_G, 'LineWidth', 1.5)
plot(Espec(Espec<co_dir*Up)/Up, theta_dir(Espec<co_dir*Up), 'Color', my_G, 'LineWidth', 1.5)
ylabel 'Phase [pi]'
box on
xlim([0 2.5])
ylim([0 20])
set(gca, 'FontSize', 12)

subplot(3,1,2)
plot(co_dir*[1 1], [-35 0], 'Color', my_B)
hold on
plot(co_rec*[1 1], [-35 0], 'Color', my_R)
plot(Espec_hist_dir(idx_dir_contrast)/Up*[1 1], [-30 0], 'Color', my_G)
plot(Espec_hist_rec(idx_rec_contrast)/Up*[1 1], [-30 0], 'Color', my_G)
plot([1.5 2.25], log10(rate_rec(idx_co_rec))*[1 1], 'Color', my_G)
plot([0.2 1], log10(rate_dir(idx_co_dir))*[1 1], 'Color', my_G)
plot(E_dir/Up, log10(rate_dir), 'Color', my_B, 'LineWidth', 1)
plot(E_rec/Up, log10(rate_rec), 'Color', my_R, 'LineWidth', 1)
scatter(E_soll, log10(w_soll*max(traj.rate)), 3+50*w_soll, my_R, 'filled')
scatter(E_solld, log10(w_solld*max(traj.rate)), 3+50*w_solld, my_B, 'filled')
ylabel 'rate'
box on
ylim([-20 -5])
xlim([0 2.5])
set(gca, 'FontSize', 12)

subplot(3,1,3)
plot(co_dir*[1 1], [0 1], 'Color', my_B)
hold on
plot(co_rec*[1 1], [0 1], 'Color', my_R)
plot(Espec_hist_dir(idx_dir_contrast)/Up*[1 1], [0 1], 'Color', my_G)
plot(Espec_hist_rec(idx_rec_contrast)/Up*[1 1], [0 1], 'Color', my_G)
plot([0 2.5], 0.1*[1 1], 'Color', my_Greylight)
plot(Espec_hist_dir/Up, contrast_dir, 'Color', my_B, 'LineWidth', 1)
plot(Espec_hist_rec/Up, contrast_rec, 'Color', my_R, 'LineWidth', 1)
xlabel 'Energy [Up]'
ylabel 'contrast'
box on
xlim([0 2.5])
ylim([0 0.15])
set(gca, 'FontSize', 12)


