%% ------------------------------------------------------------------------
%  Analyse ipRGC sampling structure on spatial image reconstruction limits
%  ------------------------------------------------------------------------
% From "What can the eye see with melanopsin?" by Nugent, T.W. & Zele, A.J.
% Software S2.
%
% Instructions to run:
%   1. This script relies on the data stored in Software S3. Ensure this
%      excel file is named "SoftwareS3_ipRGCSamplingStructureAnalysis.xlsx"
%   2. Both this script and the excel file needs to be in the same folder.
%   3. The variables below can be changed, but changes aren't guaranteed to
%      be valid.
%
%% Clear workspace and set variables
clearvars; close all; clc;

dataFileName = "SoftwareS3_ipRGCSamplingStructureAnalysis.xlsx";
dataMatName = "dataCellTemplates.mat";

% This scales the template cells to a certain image size
cellSizePixels = 500;   % Pixels
                        % Highest resolution: 500px = cellSize um
                        % Fastest execution:  200px = cellSize um

% This is the average dendritic field diameter in retinal region simulated
cellSize = 0.375;       % mm
                        % From Fig. 5A in Nasir-Ahmad 2019
                        % doi: https://doi.org/10.1002/cne.24176
                        % Current range: 375um at 6-12 degrees temporal

% This is the retinal area (a square) simulated
fieldSize = 6;          % Degrees
                        % Dendritic field size changes with eccentricity
                        % and so if field size is too large, cellSize can
                        % no longer be assumed as a constant.

% Define bottom left corner of field simulated
fieldLocation = [6 -3]; % Degrees
                        % fieldLocation = [left bottom]

% Video parameters
fps = 4;
numFrames = 50;
fs_start = 0.16;


fs_end = 2;

%% Pre-script variables
% Seed random number generator for reproducible results
rng(1)

% Define some variables to convert units
DegreesToDistance = 0.291;                      % Degrees to mm
DistanceToPixels = cellSizePixels/cellSize;     % mm to pixels

% What is the retinal area simulated in pixels
fieldSizeDistance = fieldSize * DegreesToDistance;
fieldSizePixels = ceil(fieldSizeDistance * DistanceToPixels);

%% Load in data
% Read in and store the cell templates
if ~exist(dataMatName, 'file')
    [cellTemplates, cellLayoutOuter, cellLayoutInner] = loadDataFromExcel(dataFileName);
    save(dataMatName, 'cellTemplates', 'cellLayoutOuter', 'cellLayoutInner');
else
    load(dataMatName, 'cellTemplates', 'cellLayoutOuter', 'cellLayoutInner');
end

% Resize the cell templates
for ii = 1:size(cellTemplates, 2)
    cellTemplates{ii} = imresize(cellTemplates{ii}, [cellSizePixels, cellSizePixels]);
end

% Plot the cell templates
figure('Units', 'normalized', 'Position', [0.5 0 0.5 0.92])
tiledlayout("flow", 'TileSpacing', 'tight')
for ii = 1:size(cellTemplates, 2)
    nexttile
    imshow(cellTemplates{ii})
    title(['Cell_' num2str(ii)], 'Interpreter', 'none')
end  
sgtitle('Cell Templates')

%% Read in the soma locations
scaleBarDistanceToPixels = 300/2;       % 300 pixels = 2mm
                                        % Fixed value, based on soma  
                                        % location image

% Binarize image of soma locations
thresh = 0.1;                           % Fixed - based on soma image
cellLayoutOuter = double(cellLayoutOuter < thresh);
cellLayoutInner = double(cellLayoutInner < thresh);

% Restrict soma locations to only the region of retina of the test field
FAVCenter = [318, 476];                 % Pixels - fixed


leftEdge = round(FAVCenter(1) + fieldLocation(1) * DegreesToDistance * scaleBarDistanceToPixels);
rightEdge = round(FAVCenter(1) + (fieldLocation(1) + fieldSize) * DegreesToDistance * scaleBarDistanceToPixels);
upperEdge = round(FAVCenter(2) + (fieldLocation(2) + fieldSize) * DegreesToDistance * scaleBarDistanceToPixels);
lowerEdge = round(FAVCenter(2) + fieldLocation(2) * DegreesToDistance * scaleBarDistanceToPixels);
cellLayoutOuter = cellLayoutOuter(lowerEdge:upperEdge, leftEdge:rightEdge);
cellLayoutInner = cellLayoutInner(lowerEdge:upperEdge, leftEdge:rightEdge);

% Resize soma locations image based off cellSizePixels choice and re-binarize
cellLayoutOuter = imresize(cellLayoutOuter, [fieldSizePixels, fieldSizePixels]);
cellLayoutInner = imresize(cellLayoutInner, [fieldSizePixels, fieldSizePixels]);
cellLayoutOuter = double(cellLayoutOuter > 1-thresh);
cellLayoutInner = double(cellLayoutInner > 1-thresh);

% Image processing to separate overlapping markers
SE = strel('disk', round(6 * (cellSizePixels/500)), 0);
cellLayoutOuter = imerode(cellLayoutOuter, SE);
cellLayoutInner = imerode(cellLayoutInner, SE);

% Pad the soma locations to fit cell templates near the edge
padval = cellSizePixels;
cellLayoutOuter = padarray(cellLayoutOuter, [padval, padval], 0, 'both');
cellLayoutInner = padarray(cellLayoutInner, [padval, padval], 0, 'both');

% Identify pixel positions of each soma (separately for inner and outer)
positionsOuter = regionprops(bwconncomp(logical(cellLayoutOuter),4), 'Centroid');
positionsOuter = round(vertcat(positionsOuter.Centroid));
positionsInner = regionprops(bwconncomp(logical(cellLayoutInner),4), 'Centroid');
positionsInner = round(vertcat(positionsInner.Centroid));
positions = [positionsInner; positionsOuter];

%% Create the simulated sampling structure from cellTemplates and positions

field = zeros(size(cellLayoutInner,1), size(cellLayoutInner,2), size(positions,1));

% These positions were manually found for each cell template image we use
cellOffset = [283 317; 259 244; 237 281; 283 280; 162 250; 218 181; 195 257;...
              247 285; 244 267; 278 316; 251 222; 266 244; 213 276; 260 154;...
              253 268; 315 289];        % Based on the 500 x 500 cellTemplates
cellOffset = round(cellSizePixels * cellOffset / 500);

% Place randomly selected cell templates at each soma location
for ii = 1:size(positions,1)
    cellChoice = randi(size(cellTemplates,2));
    [row, col] = find(logical(cellTemplates{cellChoice}));
    offset = cellOffset(cellChoice,:);
    for jj = 1:size(row,1)
        field(positions(ii,2) + row(jj) - offset(2), positions(ii,1) + col(jj) - offset(1),ii) = 1;
    end
end
% Remove padding from the field to constrain back to field area
field = field(padval+1:end-padval,padval+1:end-padval,:);

% Plot the simulated sampling structure of the field
receptiveFieldFig = figure();
imInner = rescale(sum(field(:,:,1:size(positionsInner,1)),3),0,1);
imOuter = rescale(sum(field(:,:,size(positionsInner,1) + (1:size(positionsOuter,1))),3),0,1);
imshow(cat(3, imInner, imOuter, imOuter))
hold on
plot(positions(:,1)-padval, positions(:,2)-padval, 'g.')

%% Parameters for the grating stimuli
% Set up for the grating stimuli
x = -size(field,2)/2:size(field,2)/2-1;
y = -size(field,1)/2:size(field,1)/2-1;
[X, Y] = meshgrid(x,y);
X = X / DistanceToPixels / DegreesToDistance;
Y = Y / DistanceToPixels / DegreesToDistance;

%% Simulation loop

% Create the video figure
figure()
set(gcf, 'Units', 'normalized', 'Position', [0 0 1 0.4*16/9])
buffX = 0.01;
buffY = buffX/2*5;

% Video file parameters
movName = 'SampledIm';
frame = 1;
v = VideoWriter(movName, "MPEG-4");
v.FrameRate = fps;
open(v);
msg = sprintf('Simulating different fs: 0/%.0f done.', numFrames);
fprintf(msg);
reverseStr = repmat(sprintf('\b'), 1, length(msg));
time = nan(1, numFrames);
fs_vec = logspace(log10(fs_start), log10(fs_end), numFrames);       % cpd
for ii = 1:numFrames
    fs = fs_vec(ii);
    tic
    clf
    ax = axes();
    set(ax, 'Units', 'normalized', 'Position', [1/5 1/2 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    RI = imref2d(size(imOuter));
    RI.XWorldLimits = fieldLocation(1) + [0 fieldSize];
    RI.YWorldLimits = fieldLocation(2) + [0 fieldSize];
    imshow(1.3*cat(3, imOuter, imInner, imInner), RI)
    yticks([])
    set(ax, 'Box', 'off')
    xlabel('Retinal eccentricity [degrees]')
    title('Sampling structure of human ipRGCs')
    text(-0.05, 1, 'B', 'FontSize', 18, 'Units', 'normalized', 'HorizontalAlignment', 'left', 'VerticalAlignment', 'top')
    
    % Create Stimuli with certain spatial frequency
    theta = 0;
    t = 0.25;
    stimulus = 0.5 * sin(2*pi*t) .* cos(2 * pi * fs * (X * cos(theta) + Y * sin(theta)));
    radial = sqrt(X.^2 + Y.^2);
    
    % Apply spatial mask to stimuli field
    beta = 0.25;
    mask = 1/2 * (1 + cos(pi/beta*(abs(radial)/fieldSize - (1-beta)/2)));
    mask(abs(radial) <= fieldSize*(1-beta)/2) = 1;
    mask(abs(radial) >= fieldSize*(1+beta)/2) = 0;
    stimulus = mask.*stimulus + 0.5;
    
    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [0 1/2 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    RI = imref2d(size(stimulus));
    RI.XWorldLimits = fieldLocation(1) + [0 fieldSize];
    RI.YWorldLimits = fieldLocation(2) + [0 fieldSize];
    imshow(stimulus, RI), hold on
    plot(positions(:,1) - padval, positions(:,2) - padval, 'r.')
    yticks([])
    set(ax, 'Box', 'off')
    xlabel('Retinal eccentricity [degrees]')
    title('Spatial grating presented')
    text(-0.05, 1, 'A', 'FontSize', 18, 'Units', 'normalized', 'HorizontalAlignment', 'left', 'VerticalAlignment', 'top')
    
    % Each ipRGC collects light across soma and dendrites
    samples = nan(size(positions,1),1);
    for jj = 1:size(positions,1)
        samples(jj) = mean(stimulus(logical(field(:,:,jj))),'all');
    end


    
    % Make new coordinate grid
    xx = 1:size(field,2)-1;
    yy = 1:size(field,1)-1;
    [XX, YY] = meshgrid(xx,yy);
    
    % Sample coordinates (including edges of image)
    sampX = [positions(:,1) - padval; 0; 0; 2328; 2328];
    sampInnerX = [positions(1:size(positionsInner,1),1) - padval; 0; 0; 2328; 2328];
    sampOuterX = [positions(size(positionsInner,1) + (1:size(positionsOuter)),1) - padval; 0; 0; 2328; 2328];

    sampY = [positions(:,2) - padval; 0; 2328; 0; 2328];
    sampInnerY = [positions(1:size(positionsInner,1),2) - padval; 0; 2328; 0; 2328];
    sampOuterY = [positions(size(positionsInner,1) + (1:size(positionsOuter)),2) - padval; 0; 2328; 0; 2328];
    
    % Sample values (including 50% for the edge samples)
    samples = [samples; 0.5; 0.5; 0.5; 0.5]; %#ok<*AGROW>
    samplesInner = [samples(1:size(positionsInner,1)); 0.5; 0.5; 0.5; 0.5]; %#ok<*AGROW>
    samplesOuter = [samples(size(positionsInner,1) + (1:size(positionsOuter))); 0.5; 0.5; 0.5; 0.5]; %#ok<*AGROW>
    sampledIm = griddata(sampX, sampY, samples, XX, YY, 'natural');
    sampledInnerIm = griddata(sampInnerX, sampInnerY, samplesInner, XX, YY, 'natural');
    sampledOuterIm = griddata(sampOuterX, sampOuterY, samplesOuter, XX, YY, 'natural');
    
    ax = axes;
    ax.Visible = false;
    set(ax, 'Units', 'normalized', 'Position', [0 0 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    text(0.5, 0.5, ['fs = ' num2str(round(fs, 2)) ' cpd'], 'FontSize', 18, 'Color', 'k', 'HorizontalAlignment','center', 'VerticalAlignment','middle')
    text(0, 0, 'Nugent, T.W. & Zele, A.J.', 'FontSize', 12, 'Color', [0.5 0.5 0.5], 'HorizontalAlignment','left', 'VerticalAlignment','middle')
    xlim([0 1]), ylim([0 1])

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [2/5 1/2 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    RI = imref2d(size(sampledIm));
    RI.XWorldLimits = fieldLocation(1) + [0 fieldSize];
    RI.YWorldLimits = fieldLocation(2) + [0 fieldSize];
    imshow(ones(size(sampledIm))*0.5, RI), hold on
    imshow(sampledIm, RI)
    title('All ipRGCs')
    yticks([])
    xticks([6 7 8 9 10 11 12])


    xticklabels([])
    set(ax, 'Box', 'off')
    yl = ylabel('Reconstructed image');
    text(-0.05, 1, 'C', 'FontSize', 18, 'Units', 'normalized', 'HorizontalAlignment', 'left', 'VerticalAlignment', 'top')

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [2/5 0 1/5 1/2] + [buffX 2.5*buffY -2*buffX -2*buffY])
    imFreq = abs(fftshift(fft2(sampledIm-mean(sampledIm, 'all'))));
    samplesPerDegree = fieldSizePixels./fieldSize;
    RI = imref2d(size(imFreq));
    RI.XWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2];
    RI.YWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2];
    if fs == 0.16
        maxPower = max(imFreq,[],'all');
    end

    imshow(imFreq, RI, [0 maxPower])
    axis([-2 2 -2 2])
    yticks([])
    xlabel('Spatial frequency [cpd]')
    yl = ylabel('Reconstructed spectrum');
    set(ax, 'Box', 'off')
    yl.Position = yl.Position;
    text(-0.05, 1, 'D', 'FontSize', 18, 'Units', 'normalized', 'HorizontalAlignment', 'left', 'VerticalAlignment', 'top')

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [3/5 1/2 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    RI = imref2d(size(sampledInnerIm));
    RI.XWorldLimits = fieldLocation(1) + [0 fieldSize];
    RI.YWorldLimits = fieldLocation(2) + [0 fieldSize];
    imshow(ones(size(sampledInnerIm))*0.5, RI), hold on
    imshow(sampledInnerIm, RI)
    yticks([])
    xticks([6 7 8 9 10 11 12])
    xticklabels([])
    title('Inner-stratifying ipRGCs')

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [3/5 0 1/5 1/2] + [buffX 2.5*buffY -2*buffX -2*buffY])
    imFreq = abs(fftshift(fft2(sampledInnerIm-mean(sampledInnerIm, 'all'))));
    samplesPerDegree = fieldSizePixels./fieldSize;
    RI = imref2d(size(imFreq));
    RI.XWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2];
    RI.YWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2]; 
    imshow(imFreq, RI, [0 maxPower])
    yticks([])


    xlabel('Spatial frequency [cpd]')
    set(ax, 'Box', 'off')
    axis([-2 2 -2 2])

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [4/5 1/2 1/5 1/2] + [buffX buffY -2*buffX -2*buffY])
    RI = imref2d(size(sampledOuterIm));
    RI.XWorldLimits = fieldLocation(1) + [0 fieldSize];
    RI.YWorldLimits = fieldLocation(2) + [0 fieldSize];
    imshow(ones(size(sampledOuterIm))*0.5, RI), hold on
    imshow(sampledOuterIm, RI)
    yticks([])
    xticks([6 7 8 9 10 11 12])
    xticklabels([])
    set(ax, 'Box', 'off')
    title('Outer-stratifying ipRGCs')

    ax = axes;
    set(ax, 'Units', 'normalized', 'Position', [4/5 0 1/5 1/2] + [buffX 2.5*buffY -2*buffX -2*buffY])
    imFreq = abs(fftshift(fft2(sampledOuterIm-mean(sampledOuterIm, 'all'))));
    samplesPerDegree = fieldSizePixels./fieldSize;
    RI = imref2d(size(imFreq));
    RI.XWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2];
    RI.YWorldLimits = [-samplesPerDegree./2 samplesPerDegree./2];    
    imshow(imFreq, RI, [0 maxPower])
    yticks([])
    set(ax, 'Box', 'off')
    xlabel('Spatial frequency [cpd]')
    axis([-2 2 -2 2])
    
    imageFrame = getframe(gcf);
    writeVideo(v, frame2im(imageFrame));
    
    time(ii) = toc;
    msg = sprintf(['Simulating different fs: %.0f/%.0f done\n'...
                   '                         %.1f seconds left.\n'],...
                   ii, numFrames, mean(time, 'omitnan')*(numFrames-ii));
    fprintf([reverseStr, msg]);
    reverseStr = repmat(sprintf('\b'), 1, length(msg));
end
close(v);
figure()
imshow(sampledIm), hold on
plot(positions(:,1) - padval, positions(:,2) - padval, 'r.')
title(['Spatial Frequency = ' num2str(fs) ' cpd'])
exportgraphics(receptiveFieldFig, 'ReceptiveFieldFigure.png', 'Resolution', '100')

%% Local Functions
function [cellTemplates, cellLayoutOuter, cellLayoutInner] = loadDataFromExcel(dataFileName)


    numCellTemplates = numel(sheetnames(dataFileName)) - 2;
    time = nan(1, numCellTemplates);
    cellTemplates = cell(1, numCellTemplates);
    msg = sprintf('Reading in cell templates: 0/%.0f done.', numCellTemplates);
    fprintf(msg);
    reverseStr = repmat(sprintf('\b'), 1, length(msg));
    for ii = 1:numCellTemplates
        tic;
        cellTemplates{ii} = readmatrix(dataFileName, 'Sheet', ['cell_' num2str(ii)]);
        time(ii) = toc;
        msg = sprintf(['Reading in cell templates: %.0f/%.0f done\n'...
            '                           %.1f seconds left.\n'],...
            ii, numCellTemplates, mean(time, 'omitnan')*(numCellTemplates-ii));
        fprintf([reverseStr, msg]);
        reverseStr = repmat(sprintf('\b'), 1, length(msg));
    end
    msg = sprintf('Reading in cell locations.\n');
    fprintf(msg)
    cellLayoutOuter = readmatrix(dataFileName, 'Sheet', 'cellLayoutOuter');
    cellLayoutInner = readmatrix(dataFileName, 'Sheet', 'cellLayoutInner');
end