%% Fourier Prony Recovery Algorithm
% This script showcases USF Fourier Prony Recovery Algorithm.
% First, define a functions for zero-centered modulo operation and Dirichlet kernel.

function folded_samples = cent_mod(samples, lambda)
    tmp = double(samples) ./ (2 * lambda) + 1/2;
    fractional_part = tmp - floor(tmp);
    folded_samples = 2 * lambda * (fractional_part - 1/2);
end

function dir_vals = gen_dir_ker(order, t)
    dir_vals = sin((2*order + 1) * t / 2) ./ sin(t / 2);

    dir_vals(t == 0) = 2*order + 1;

    dir_vals = dir_vals / (2*order + 1);
end

%% Signal Acquisition
% Define parameters for the script and time variable.

N = 2000;
LAMBDA = 0.75;
tau = 2 * pi;
T = tau / N;
M = 6;

n = (0:1:(N - 1))';
t = T * n;

% Generate input signal to USF acquisition pipeline, apply modulo and plot input signal and modulo samples.

gn = 2.5 * gen_dir_ker(M, t - tau/2);
yn = cent_mod(gn, LAMBDA);

figure;
hold on;
scatter(t, gn, "filled", "Color", "blue");
stem(t, yn, "Color", "red");
xlim([min(t) max(t)]);
xlabel("Time (s)");
legend("g[n]", "y[n]");
legend boxoff;

% Using the residue signal calculate the number of signal folds.

resn = gn - yn;
resn_diff = diff(resn);

n_k_gt = [];
a_k_gt = [];

for n_idx = 1:1:length(resn_diff)
    if abs(resn_diff(n_idx)) > LAMBDA
        n_k_gt = [n_k_gt; n_idx];
        a_k_gt = [a_k_gt; resn_diff(n_idx)];
    end
end

K = length(n_k_gt);

%% Signal Reconstruction
% Calculate first-order differences of y[n], obtain out of the band part of the spectra and plot.

M_OFFSET = 3; % Offset to mitigate potential effect of spectral leakage
spect_offset = M + M_OFFSET;

yn_diff = diff(yn);
yn_d_dft = fft(yn_diff);
res_dft = yn_d_dft(spect_offset + 1: end - spect_offset);
m = (spect_offset:1:(N - spect_offset - 2))';

figure;
hold on;
plot(m, abs(res_dft), "Color", "blue", "LineWidth", 2);
xlabel("Spectral index");
xlim([min(m) max(m)]);
title("Out band part of the spectra");

% Use annihilation filter to estimate folding locations.

toep_mtx = toeplitz(-flip(res_dft));
A = toep_mtx(1:end - K, end - K + 1:end);
b = res_dft(K+1:end);
hn = A\b;
hn = [1;hn];

z_roots = roots(hn);
log_mod = mod(-imag(log(z_roots)) + 2 * pi, 2 * pi);
n_k_hat = log_mod * (N - 1) / (2 * pi);
exps = exp(-1j * 2 * pi * m * transpose(n_k_hat) / (N - 1));
a_k_hat = exps\(-res_dft);

% Reconstruct the residue, log MSE on parameter estimates.

res_diff_rec = zeros(size(yn_diff));
res_diff_rec(round(n_k_hat) + 1) = a_k_hat;
res_rec = cumsum([0; res_diff_rec]);

[n_k_hat_srt, srt_idx] = sort(n_k_hat);
a_k_hat_srt = a_k_hat(srt_idx);

n_k_MSE = mean((n_k_hat_srt - n_k_gt) .^ 2);
a_k_MSE = mean(abs(a_k_hat_srt - a_k_gt) .^ 2);

fprintf("n_k MSE: %d; a_k MSE: %d\n", n_k_MSE, a_k_MSE);

% Reconstruct the signal using the residue and modulo signal.

gn_rec = res_rec + yn;

gn_MSE = mean(abs(gn_rec - gn) .^ 2);

fprintf("g_n MSE: %d\n", gn_MSE);

figure;
hold on;
plot(t, gn, "Color", "blue");
plot(t, real(gn_rec), "--", "Color", "red");
xlim([min(t) max(t)]);
legend("Original", "Reconstruction");
legend boxoff;
xlabel("Time (s)");