clear all
close all

%% loading, splitting data and metadata, data sorting, and definition of constants 
load('../data/Neuron22.mat', 'data_unsorted')
sample_rate = 1000;
conditions = data_unsorted(:, 1);
spike_data = data_unsorted(:, 2:end);
data_sorted = sortrows(data_unsorted, 1, 'descend');
sorted_conditions = data_sorted(:,1);
sorted_spike_data = data_sorted(:, 2:end);

condition_colors = ["red", "green", "blue", "k"];
segments = ["Fixation", "Stimulus ON", "Delay"];
segment_times = [-0.5, 0; 0.001, 0.8; 0.801, 1.8];

%% Scatterplots of the spike responses  
fig = scatterplot(spike_data, conditions, sample_rate, segments, segment_times, condition_colors);
fig = format_figure(fig, [15, 10]);
saveas(fig, "scatterplot_unsorted", "pdf")

fig = scatterplot(sorted_spike_data, sorted_conditions, sample_rate, segments, segment_times, condition_colors);
fig = format_figure(fig, [15, 10]);
saveas(fig, "scatterplot_sorted", "pdf");

%% Firing rate analysis
bin_width = 0.1;
[rates, bins] = get_firing_rates(spike_data, sample_rate, bin_width);
fig = plot_firing_rates(conditions, rates, bins, segments, segment_times, condition_colors);
fig = format_figure(fig, [15, 10]);
saveas(fig, "firing_rates", "pdf")

fig = plot_average_rates(rates, bins, conditions, segments, segment_times, condition_colors);
fig = format_figure(fig, [10,10]);
saveas(fig, "rate_comparison", "pdf")

%% function definitions
function f = plot_average_rates(rates, bins, conditions, segments, segment_times, condition_colors)
    f = figure();
    unique_conditions = sort(unique(conditions));
    bins = bins + segment_times(1,1); 
    max_rate = floor(max(rates(:))/5) * 5;
    for i = 1:length(segments)
        subplot(1, length(segments), i)
        hold on
        box_data = [];
        grouping_data = [];
        for j = 1:length(unique_conditions)
            condition_rates = mean(rates(conditions == unique_conditions(j), ...
                bins >= segment_times(i,1) & bins < segment_times(i, 2)));
            box_data = cat(1, box_data, condition_rates);
            grouping_data = cat(1, grouping_data, ... 
                ones(length(condition_rates),1) * unique_conditions(j));
        end
        boxplot(box_data, grouping_data, "Widths", 0.3, "ColorGroup", condition_colors)
        xticklabels(1:length(unique_conditions))
        xlabel("Number of dots")
        if i == 1
            ylabel("Firing rate [Hz]")
        end
        ylim([0, max_rate])
        box("off")
        title(segments(i))
    end
end

function f = format_figure(f, size)
    f.Color = "White";
    f.PaperSize = size;
    f.PaperPosition = [0, 0, size(1), size(2)];
end

function f = plot_firing_rates(conditions, rates, bins, segments, segment_times, condition_colors)
    f = figure();
    unique_conditions = sort(unique(conditions));

    for i = 1:length(unique_conditions)
        condition = unique_conditions(i);
        subplot(length(unique_conditions), 1, i)
        hold on
        e = errorbar(bins + segment_times(1,1), mean(rates(conditions == condition, :)), std(rates), ...
            "Color", condition_colors(condition-1), "DisplayName", ...
            sprintf("Number of points: %i", condition-1));
        xlim([segment_times(1,1)-0.1, segment_times(end, end) + 0.1])
        ylim([0, 5]) 
        ylabel("Firing rate [Hz]")
        if i < length(unique_conditions)
            xticklabels([])
        else
            xlabel("Time [s]")
        end
        
        for j = 1:length(segments)
            if i == 1 
                label_x_pos = segment_times(j,2) - (segment_times(j, 2) - segment_times(j, 1))/2;
                t = text(label_x_pos, 5.5, segments(j));
                t.HorizontalAlignment = "center";
                t.FontSize = 10;
            end
            if j < length(segments)
                plot([segment_times(j, 2), segment_times(j, 2)], ... 
                    [0, 5],'k--', 'linewidth', 1.5)
            end
        end
        legend(e);
        box('off')
    end  
end

function [rates, bins] = get_firing_rates(data, sample_rate, bin_width)
    
    time = (0:size(data, 2)-1) ./ sample_rate;
    edges = 0:bin_width:time(end)+bin_width; 
    rates = zeros(size(data, 1), length(edges)-1);
    for i = 1:size(data, 1)
        rates(i, :) = histcounts(time(data(i,:) == 1), edges);
    end
    bins = edges(1:end-1);
end

function f = scatterplot(data, conditions, sample_rate, segments, segment_times, colors)
     
    f = figure();
    hold on
    time = (0:size(data, 2)-1) ./ sample_rate + segment_times(1,1);
    for i = 1:size(data, 1)
        c = colors(conditions(i) - 1);
        spike_times = time(data(i,:) == 1);
        scatter(spike_times, ones(length(spike_times),1) * i, c)
    end
    xlabel('Time [s]')
    ylabel('Trials')
    for i = 1:length(segments)
        label_x_pos = segment_times(i,2) - (segment_times(i, 2) - segment_times(i, 1))/2;
        t = text(label_x_pos, 275, segments(i));
        t.HorizontalAlignment = "center";
        t.FontSize = 10;
        if i < length(segments)
            plot([segment_times(i, 2), segment_times(i, 2)],[0, size(data,1)],'k--', 'linewidth', 1.5)
        end
    end
    ylim([0, size(data, 1) + 50])
    yticks(0:50:250)
    box off
    hold off
end