clear;
clc;

load_const;
load_au;
load_cm;
load_color;

% Additional plots?
in.make_plots = 0;
in.plot_trajectory = 0;

% Propagation effects?
in.recoll_on = 1;           % propagation including recollisions (yes/no)
in.tunnel_exit = 1;         % Start electrons at scaled tunnel exit (x_te = in.tunnelexit * W/eE)
in.recoll_prob = @(E) 1.0;  % Recollision probability as function of return energy

% for calculation of energy spectrum
in.use_dir = 1;     % include direct electrons?
in.use_recoll = 1;  % include recollision electrons?
in.use_phases = 1;  % include phase in energy spectrum? 
scale = 0.5;        % scale factor for energy width
eSMM_spec = 1;      % use extended calculation? 

% simulation variables
dt = 0.0005;                % timestep [fs]
t_list = -2:dt:50; %-3:dt:5;           % propagation time [fs] % -2:dt:50 for phase extrapolation plot
tbirth = -5:dt:5;           % birth times [fs]
tbirth_cycles = 1;          % birth times only in specifiv cycles, if 0: use tbirth from above
cep_list = -0.5:0.01:2.5;   % CEPs [pi]
in.gamma_0 = 7;             % maximal nearfield enhancement
in.lambda_nf = 100;         % decay length [Å]
in.U_0   = 0;               % constant DC field (here not considered)
in.r_tip = in.lambda_nf;    % constant DC field (here not considered)
in.W = 6.6;                 % work function of metal [eV]
pulse.wavelength = 800;     % wavelength incident laser pulse [nm]
pulse.intensity = 1e12;     % intensity incident laser pulse [W/cm^2]
pulse.width = -1;           % FWHM incident laser pulse [fs], if -1: pi/omega (single-cycle)
pulse.t0 = 0;               % time where envelope function is maximal [fs]

% Create incident laser pulse
Pulse = get_laserfield("intensity", pulse.intensity, "phase", cep_list(1), "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

% Set matrices
Ekin_list = nan(numel(t_list), numel(cep_list));
rate_list = nan(numel(t_list), numel(cep_list));
S_list = nan(numel(t_list), numel(cep_list));
roi_direct_list = nan(numel(t_list), numel(cep_list));
roi_recoll_list = nan(numel(t_list), numel(cep_list));
Ebirth_list = nan(numel(t_list), numel(cep_list));
E_cutoff = nan(numel(cep_list));
E_cutoff_dir = nan(numel(cep_list));
E_cutoff_recoll = nan(numel(cep_list));

% CEP loop
for phi_i = 1:numel(cep_list)
    pulse.phi = cep_list(phi_i);

    % birth times from specific cycles
    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 routine
    [traj, Pulse] = calculate_smm_propagation(in, pulse, t_list, tbirth);
    traj_mat(phi_i) = traj;

    if isempty(traj.Ebirth)
        continue
    end
    % cutoff energy
    if sum(traj.roi_direct) == 0
        E_cutoff_dir(phi_i) = nan;
    else
        E_cutoff_dir(phi_i) = max(traj.Ekin(traj.roi_direct));
    end
    if sum(traj.roi_recoll) == 0
        E_cutoff_recoll(phi_i) = nan;
    else
        E_cutoff_recoll(phi_i) = max(traj.Ekin(traj.roi_recoll));
    end
    
    min_rate = 1e-10*max(traj.rate); % minimal ionization rate considered relevant
    E_cutoff(phi_i) = max(traj.Ekin(traj.rate > min_rate));
    
    Ekin_list(traj.roi_tbirth, phi_i) = traj.Ekin;
    rate_list(traj.roi_tbirth, phi_i) = traj.rate;
    S_list(traj.roi_tbirth, phi_i) = traj.S.tot;
    roi_direct_list(traj.roi_tbirth, phi_i) = traj.roi_direct;
    roi_recoll_list(traj.roi_tbirth, phi_i) = traj.roi_recoll;
    Ebirth_list(traj.roi_tbirth, phi_i) = traj.Ebirth;
end
mean_cutoff = mean(E_cutoff);

%% Extrapolate phases and calculate energy spectrum 

Espec = 0:0.005*Up:30*Up;
in.Up = Up; in.Ephot = Pulse.Ephot;

% evaluation and extrapolation of phase
fit_results = fit_EffPhase_over_energy_cep(Espec, cep_list, traj_mat, weight_scale = 1);

% calculation of energy spectra
spectrum = calc_spec_cep(Espec, cep_list, traj_mat, fit_results.Theta_Fit_list_dir, fit_results.Theta_Fit_list_rec, in, scale=scale);
Espec = Espec + traj.U_0;

%% Create full spectrum from one cycle 
[~, idxmat.Zpi] = min(abs(cep_list));
[~, idxmat.Tpi] = min(abs(cep_list-2));

cep_full = cep_list(cep_list >= 0 & cep_list <= 2);

[~, idxmat.ub_1] = min(abs(cep_list - wrapTo2Pi(cep_list(end)*pi)/pi));
[~, idxmat.lb_1] = min(abs(cep_list - wrapTo2Pi(cep_list(1)*pi)/pi));

[~, idxmat.ub] = min(abs(cep_full - wrapTo2Pi(cep_list(end)*pi)/pi));
[~, idxmat.lb] = min(abs(cep_full - wrapTo2Pi(cep_list(1)*pi)/pi));

[spec_dir_full, dir_fits] = puzzle_from_one_cycle(Espec, cep_full, spectrum.direct, fit_results.Theta_Fit_list_dir, idxmat);
[spec_rec_full, rec_fits] = puzzle_from_one_cycle(Espec, cep_full, spectrum.recoll, fit_results.Theta_Fit_list_rec, idxmat);
spec_full = spec_dir_full + spec_rec_full;

%% Calculate everything needed for the figures
load("Final_Scripts\TDSE_ref.mat") % TDSE results

% find indices of cep_plot for CEP = 0,pi
cep_plot = [0, 1];
[~, ind_c] = ismembertol(cep_plot, cep_full);
[~, ind_c_tdse] = ismembertol(cep_plot, cep_tdse);

% calculate cutoff for direct and recoll 
E_co_dir = nan(size(cep_list));
E_co_rec = nan(size(cep_list));
for i = 1:numel(cep_list)
    roi_birth = ~isnan(Ekin_list(:,i));
    roi_dir = logical(roi_direct_list(roi_birth, i));
    roi_rec = logical(roi_recoll_list(roi_birth, i));
    rate = rate_list(roi_birth, i);
    roi_rate = rate > 0; %1e-10;
    Ekin = Ekin_list(roi_birth, i);

    if ~isempty(Ekin(roi_dir&roi_rate))
        E_co_dir(i) = max(Ekin(roi_dir&roi_rate));
    end
    if ~isempty(Ekin(roi_rec&roi_rate))
        E_co_rec(i) = max(Ekin(roi_rec&roi_rate));
    end
end
E_co_dir_full = nan(size(cep_full));
E_co_rec_full = nan(size(cep_full));
E_co_dir_full(1:idxmat.ub) = max(E_co_dir(idxmat.Zpi:idxmat.ub_1), E_co_dir(idxmat.Tpi:end)); E_co_dir_full(idxmat.ub+1:idxmat.lb-1) = E_co_dir(idxmat.ub_1+1:idxmat.lb_1-1); E_co_dir_full(idxmat.lb:end) = max(E_co_dir(idxmat.lb_1:idxmat.Tpi), E_co_dir(1:idxmat.Zpi));
E_co_rec_full(1:idxmat.ub) = max(E_co_rec(idxmat.Zpi:idxmat.ub_1), E_co_rec(idxmat.Tpi:end)); E_co_rec_full(idxmat.ub+1:idxmat.lb-1) = E_co_rec(idxmat.ub_1+1:idxmat.lb_1-1); E_co_rec_full(idxmat.lb:end) = max(E_co_rec(idxmat.lb_1:idxmat.Tpi), E_co_rec(1:idxmat.Zpi));

% find indices and relevant data (TDSE) for CEP = 1.1pi
cep_test = 1.1;
[~, cep_index] = min(abs(cep_test-cep_tdse));
Espec_tdse = E_tdse;
roi = Espec_tdse < 4.5;
spec_all_tdse = log10(spec_tdse_cep_scan(cep_index,:));

% calculate max/min for TDSE and relative spectral phase
help_lin =  -4.3/5*Espec_tdse-2.5;
max_ind = islocalmax(spec_all_tdse(roi) - help_lin(roi), 'MinSeparation', 0.1, 'MinProminence', 0.05);
min_ind = islocalmin(spec_all_tdse(roi) - help_lin(roi), 'MinSeparation', 0.1, 'MinProminence', 0.05);
n_max = numel(Espec_tdse(max_ind));
phase_max_list = 2*(1:n_max)-2;
n_min = numel(Espec_tdse(min_ind));
phase_min_list = 2*(1:n_min)-1;

% find data of eSMM for CEP = 1.1pi
[~, cep_index] = min(abs(cep_test-cep_full));
spec_plot_all = spec_full(:,cep_index);
spec_plot_dir = spec_dir_full(:,cep_index);
spec_plot_rec = spec_rec_full(:,cep_index);
total_relative_phase = unwrap(angle(spec_plot_rec)-angle(spec_plot_dir));
[~, cep_index] = min(abs(cep_test-cep_list));
traj = traj_mat(cep_index);
direct_list = find(traj.roi_direct);
recoll_list = find(traj.roi_recoll);

% define sample list (dir and rec) for energies up to lower cutoff
co_dir = max(traj.Ekin(direct_list)/traj.Up);
co_rec = max(traj.Ekin(recoll_list)/traj.Up);
max_co = min(co_dir,co_rec);
E_sample_list = linspace(0.0, max_co, 100);
for E_sample_index = 1:numel(E_sample_list)
    E_sample = E_sample_list(E_sample_index);

    index_list = 1:numel(traj.Ekin);
    index_list_direct = index_list(direct_list);
    index_list_recoll = index_list(recoll_list);

    [~, idx_direct] = max(abs(traj.Ekin(traj.roi_direct)/traj.Up)); idx_direct = index_list_direct(idx_direct);
    [~, idx_recoll] = max(abs(traj.Ekin(traj.roi_recoll)/traj.Up)); idx_recoll = index_list_recoll(idx_recoll);

    p_cut_direct = traj.rate(idx_direct);
    p_cut_recoll = traj.rate(idx_recoll);

    index_list_direct = index_list(traj.roi_direct & traj.rate>p_cut_direct);
    index_list_recoll = index_list(traj.roi_recoll & traj.rate>p_cut_recoll);

    [~, idx_direct] = min(abs(traj.Ekin(traj.roi_direct & traj.rate>p_cut_direct)/traj.Up - E_sample)); idx_direct = index_list_direct(idx_direct);
    [~, idx_recoll] = min(abs(traj.Ekin(traj.roi_recoll & traj.rate>p_cut_recoll)/traj.Up - E_sample)); idx_recoll = index_list_recoll(idx_recoll);

    direct_sample_list(E_sample_index) = idx_direct;
    recoll_sample_list(E_sample_index) = idx_recoll;
end

% calc relative spectral phase contributions (theta0, thetaT, thetaV) and
% prepare fill plot (Y1,Y2,Y3)
diff_W = (-traj.S.W(recoll_sample_list) + traj.S.W(direct_sample_list)) / (const.hbar * pi);
diff_T = (traj.S.T(recoll_sample_list) - traj.S.T(direct_sample_list)) / (const.hbar * pi);
diff_V = (traj.S.V(recoll_sample_list) - traj.S.V(direct_sample_list)) / (const.hbar * pi);

X = [E_sample_list, fliplr(E_sample_list)];
Y1 = [zeros(size(E_sample_list)), fliplr(diff_W)];
Y2 = [diff_W, fliplr(diff_W+diff_T)];
Y3 = [diff_W+diff_T, fliplr(diff_W+diff_T+diff_V)];

% 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, 100);
roi_dir = Espec_hist_dir < co_dir*Up;
Espec_hist_rec = linspace(0, 1.1*co_rec*Up, 100);
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);

% calc contrast and spectral phase evolution contour lines for
% CEP-energy-map
contrast = (2*abs(spec_dir_full.').*abs(spec_rec_full.'))./(abs(spec_dir_full.').^2 + abs(spec_rec_full.').^2);
M_C = contourc(Espec/Up, cep_full, 100*contrast, [10, 1e10]);
dir_fits_plot = dir_fits.';
rec_fits_plot = rec_fits.';
M_pd = contourc(Espec/Up, cep_full, (rec_fits_plot-dir_fits_plot)/pi, -19:2:40);

%% Plot spectra for CEP of 0 and pi and compare to TDSE, and plot CEP dependent spectra

figure('Color', 'w')
subplot(3,1,1)
plot(E_tdse, log10(spec_tdse_cep_scan(ind_c_tdse(1),:)), 'Color', my_Lila, 'LineWidth', 1.5)
hold on
plot(E_tdse, log10(spec_tdse_cep_scan(ind_c_tdse(2),:)), 'Color', my_Green, 'LineWidth', 1.5)
plot(Espec/Up, log10(abs(spec_full(:, ind_c(1))).^2)-2, '--', 'Color', my_Lila, 'LineWidth', 1.5)
plot(Espec/Up, log10(abs(spec_full(:, ind_c(2))).^2)-2, '--', 'Color', my_Green, 'LineWidth', 1.5)
xlim([0 15])
ylim([-11 -1])
xlabel 'Energy (U_p)'
ylabel 'log(Yield) (arb.u.)'
set(gca, 'FontSize', 14)

subplot(3,1,2)
imagesc(E_tdse, cep_tdse, log10(spec_tdse_cep_scan))
colormap(cm.vlasov)
clim([-11 -1])
xlim([0 18])
xlabel('E [U_p]')
ylabel('CEP [pi]')
set(gca, 'FontSize', 14)

subplot(3,1,3)
imagesc(Espec/Up, cep_full, log10(abs(spec_full.').^2))
xlim([0 18])
clim([-9 1])
colormap(cm.vlasov)
xlabel 'Energy (U_p)'
ylabel 'CEP (\pi)'
set(gca, 'FontSize', 14)

hold on
plot(E_co_dir_full/traj.Up, cep_full, 'Color', my_B, 'LineWidth', 1)
plot(E_co_rec_full/traj.Up, cep_full, 'Color', my_R, 'LineWidth', 1)

%% Plot for CEP = 1.1pi spectra, phase extrapolation and phase analysis

figure('Color', 'w')
subplot(3,1,1)
hold on
plot([co_dir co_dir], [-8 0], 'Color', my_B)
plot([co_rec co_rec], [-8 0], 'Color', my_R)
plot(Espec_tdse(max_ind), spec_all_tdse(max_ind), 'o', 'MarkerFaceColor', my_Green, 'MarkerEdgeColor', 'none')
plot(Espec_tdse(min_ind), spec_all_tdse(min_ind), 'o', 'MarkerEdgeColor', my_Green)
plot(Espec_tdse, spec_all_tdse, 'LineWidth', 1.5, 'Color', my_Green)
plot(Espec(Espec<max_co*Up)/traj.Up, log10(conj(spec_plot_all(Espec<max_co*Up)).*spec_plot_all(Espec<max_co*Up))-2, 'k', 'LineWidth', 1.5);
plot(Espec(Espec>max_co*Up)/traj.Up, log10(conj(spec_plot_all(Espec>max_co*Up)).*spec_plot_all(Espec>max_co*Up))-2, 'k--', 'LineWidth', 1.5);
xlim([0 2.5])
ylim([-6.5 -1.8])
xlabel 'Energy (Up)'
ylabel 'log10(Yield) (arb.u.)'
box on

% Phase extrapolation
color_line = my_Greylight;
subplot(3,1,2)
hold on
plot([co_dir co_dir], [-1 25], 'Color', my_B)
plot([co_rec co_rec], [-1 25], 'Color', my_R)
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/Up, fit_results.Theta_Fit_list_rec(:, cep_index)/pi + fit_results.offset_rec(cep_index), '--', 'Color', color_line, 'LineWidth', 1.5)
plot(Espec_hist_rec(roi_rec)/Up, fit_results.theta_eff_rec(roi_rec,cep_index)/pi + fit_results.offset_rec(cep_index), 'Color', color_line, 'LineWidth', 1.5)
plot(Espec/Up, fit_results.Theta_Fit_list_dir(:, cep_index)/pi + fit_results.offset_dir(cep_index), '--', 'Color', color_line, 'LineWidth', 1.5)
plot(Espec_hist_dir(roi_dir)/Up, fit_results.theta_eff_dir(roi_dir,cep_index)/pi + fit_results.offset_dir(cep_index), 'Color', color_line, 'LineWidth', 1.5)
xlim([0 2.5])
ylim([2 16])
box on
xlabel 'Energy [Up]'
ylabel 'Phase [pi]'

% Phase analysis
subplot(3,1,3)
hold on
fill(X, Y1, [0.8 1 0.8], 'EdgeColor', 'none')
fill(X, Y2, [0.8 0.8 1], 'EdgeColor', 'none')
fill(X, Y3, [1 0.9 0.8], 'EdgeColor', 'none')
plot([co_dir co_dir], [-1 15], 'Color', my_B)
plot([co_rec co_rec], [-1 15], 'Color', my_R)
plot(Espec(Espec<max_co*Up)/traj.Up, total_relative_phase(Espec<max_co*Up) / pi + 2, 'k', 'LineWidth', 1.5);
plot(Espec(Espec>max_co*Up)/traj.Up, total_relative_phase(Espec>max_co*Up) / pi + 2, 'k--', 'LineWidth', 1.5);
plot(Espec_tdse(max_ind), phase_max_list + 2, 'o', 'MarkerFaceColor', my_Green, 'MarkerEdgeColor', 'none')
plot(Espec_tdse(min_ind), phase_min_list + 2, 'o', 'MarkerEdgeColor', my_Green)
xlabel 'Energy (Up)'
ylabel 'Relative spectral Phase (\pi)'
box on
xlim([0 2.5])
ylim([0 14])


%% Selective spectra and contour lines

figure('Color', 'w')
subplot(2,3,1)
imagesc(Espec/Up, cep_full, log10(abs(spec_dir_full.').^2))
xlim([0 15])
clim([-9 1])
colormap(cm.vlasov)
hold on
plot(E_co_dir_full/traj.Up, cep_full, 'Color', my_B, 'LineWidth', 1)
xlabel 'Energy (U_p)]'
ylabel 'CEP (\pi)'
set(gca, 'FontSize', 14)

subplot(2,3,4)
imagesc(Espec/Up, cep_full, log10(abs(spec_rec_full.').^2))
xlim([0 15])
clim([-9 1])

colormap(cm.vlasov)
hold on
plot(E_co_rec_full/traj.Up, cep_full, 'Color', my_R, 'LineWidth', 1)
xlabel 'Energy (U_p)]'
ylabel 'CEP (\pi)'
set(gca, 'FontSize', 14)

subplot(2,3,[2,3,5,6])
imagesc(Espec/Up, cep_full, log10(abs(spec_full.').^2))
xlim([0 15])
clim([-9 1])
colormap(cm.vlasov)
hold on
plotContourLines(M_C, 'Color', [3 41 129]/255);
plotContourLines(M_pd, 'k');
xlabel 'Energy (U_p)]'
ylabel 'CEP (\pi)'
set(gca, 'FontSize', 14)
