clear
clc

%% load data and set some variables
load('p-unit_spike_times.mat');
load('p-unit_stimulus.mat');

sample_rate = 20000;
max_time = max(stimulus(:,1));
kernel_width = 0.015;

%% plot the stimulus
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')

subplot(3, 1, 1)
plot(stimulus(:, 1), stimulus(:, 2), 'LineWidth', 1, 'DisplayName', 'stimulus')
xlim([0, 1])
ylim([-1, 1])
xlabel('time [s]')
ylabel('stimulus [arb. units]')
box('off')
set(gca, 'TickDir', 'out')
set(gca, 'XGrid', 'on')
legend('show')

%% calculate the PSTH and plot it
[average_psth, std_psth] = psth(spike_times, max_time, sample_rate, kernel_width);
x = (1:length(average_psth))/sample_rate;

subplot(3, 1, 2)
errX = [x, fliplr(x)];
errY = [average_psth - std_psth, fliplr(average_psth + std_psth)];
fill(errX, errY, 'b', 'FaceAlpha', 0.3, 'EdgeAlpha', 0.3, 'DisplayName', 'std');
hold on
plot(x, average_psth, 'LineWidth', 1, 'DisplayName', 'PSTH')

xlim([0, 1.0])
xlabel('time [s]')
ylim([0, max(average_psth)*1.2])
ylabel('firing rate [Hz]')
set(gca, 'TickDir', 'out')
box('off')
set(gca, 'XGrid', 'on')
legend('show')

%% calculate the cross-correlation and plot it

[c, lags] = xcorr(stimulus(:,2), average_psth, 1000, 'coeff');
subplot(3, 1, 3)
plot(lags/sample_rate, c, 'LineWidth', 1, 'DisplayName', 'cross-correlation')
xlabel('time [s]')
ylabel('correlation')
set(gca, 'TickDir', 'out')
legend('show')
grid('on')
box('off')