clear 
clc

%% load data
load p-unit_spike_times.mat
load p-unit_stimulus.mat
sample_rate = 20000;

%% calculate STA
all_times = [];
for i = 1:length(spike_times)
    all_times = cat(1, all_times, spike_times{i});
end

[st_average, sta_sd, num] = sta(stimulus(:,2), all_times, 1000, sample_rate);
fig = figure();
set(fig, 'PaperUnits', 'centimeters');
set(fig, 'PaperSize', [11.7 9.0]);
set(fig, 'PaperPosition',[0.0 0.0 11.7 9.0]);
set(fig,'Color', 'white')
plot(((1:length(st_average))-1000)/sample_rate, st_average)
xlabel('time [s]')
ylabel('stimulus')

%% reverse reconstruction

% make binary representation of the spike times
estimated_stims = zeros(size(stimulus, 1), length(spike_times));
for i = 1:length(spike_times)
    binary_spikes = zeros(size(stimulus, 1));
    binary_spikes(round(spike_times{i}*sample_rate)) = 1;
    estimated_stims(:,i) = conv(binary_spikes, st_average, 'same');
end

fig = figure();
set(fig, 'PaperUnits', 'centimeters');
set(fig, 'PaperSize', [11.7 9.0]);
set(fig, 'PaperPosition',[0.0 0.0 11.7 9.0]);
set(fig,'Color', 'white')
plot(stimulus(:,1), stimulus(:,2), 'displayname','original')
hold on
plot(stimulus(:,1), mean(estimated_stims,2), 'r', 'displayname', 'reconstruction')
xlabel('time [s]')
ylabel('stimulus')
legend show