clear
close all
hmfolder = pwd;
datadir = '..\datasets\videos\registered\'; %Directory containing all of the data. Assumed to end in \
load(strcat(datadir, '20220824_all35_focuslabels.mat'));
load(strcat(datadir,'20220826_registeredFOVs.mat'))
cd(datadir)
bigNaN = cell(2, 3); %Stores the dFF with out of focus frames removed for genotypexstimulation sorted by coefficient of first PC
bigPCA = cell(2, 3); %Stores the PCA output for genotypexstimulation 
bigNorm = cell(2, 3); %Stores the normalized dFF with out of focus frames removed sorted by the coefficient of the first PC
gdFF = {}; %Stores the dFF with out of focus frames removed as mouse and order by which matching was performed (1-2, 1-3, 2-3)
gnsdFF = {}; %Stores the normalized dFF with out of focus frames removed as mouse and order by which matching was performed (1-2, 1-3, 2-3)
gmag = {}; %Stores the average magnitude of the movement as mouse and order by which matching was performed (1-2, 1-3, 2-3)
count = 0;
window = 1000; %Window for the normalization
calcdF = @(F, B) (F-B)./medfilt1(F, window, [], 2, 'truncate', 'omitnan');
subdF = @(F, B) calcdF(F, B)-median(calcdF(F, B), 2); %Currently not used in favor of subtracting the baseline from the first 500 frames

for fn = 1:length(focus_labels)
    vidname = focus_labels(fn).name; %Video name associated with the focus labels
    ovname = vidname;
    indmat = []; %Acts as the indexes for gdFF, gnsdFF, and gmag
    for g = 1:length(group) %Identify which mouse and number video that vidname comes from
        gnames = group{g};
        for v = 1:length(gnames)
            tname = gnames{v};
            tgind = strfind(tname, 'Green');
            if contains(vidname, tname(1:tgind))
                indmat = [g, v]; %[Group/mouse, video number for mask matching] 
            end
        end
    end
    dind = strfind(ovname, '-'); %Remove dashes from file name
    tind = strfind(ovname, '30'); %Identify where the 30 is that frequently breaks up the stimulation type
    vname = ovname;
    if ~isempty(tind)
     vname((0:1)+tind(end)) = [];
    end
    Dind = strfind(vname, 'D');
    if isempty(Dind)
    else
    zind = strfind(vname(Dind(end):end), '0');
    stim = lower(vname(((1):(zind-2))+Dind)); %Get the stimulation type
    end
    
    greenind = strfind(vidname, 'Green'); %Find the corresponding Fluorescence matrix
    fluomat = dir(strcat(vidname(1:greenind), '*Fluorescence.mat'));
    load(fluomat(1).name)
    dFF = calcdF(F, B); %Calculate the dFF
    dFF = dFF-repmat(median(dFF(:, 1:500), 2), 1, size(dFF, 2)); %Subtract the baseline of the first 500 frames
    sdFF = dFF; 
    nsdFF = normalize(sdFF', 'center', 'median', 'scale', 'mad')';
    raster = zeros(size(nsdFF)); %Create and populate the raster of frames above 5 MAD
    raster(nsdFF > 5) = 1;
    labels = focus_labels(fn).labels;
    idx = focus_labels(fn).idx; %Remove out of focus frames for the relevant clusters of neurons
    for ngroup = 1:size(labels, 1)
        cind = find(idx == ngroup);
        remind = find(labels(ngroup, :) == 1);
        for neuron = cind
            dFF(neuron, remind) = NaN;
            nsdFF(neuron, remind) = NaN;
        end
    end
    sdFF = dFF;
    numNaN = sum(isnan(dFF), 1); %Remove neurons that go out of focus before frame 500. This only impacts the pooled PCA, bigPCA and pooled dFF
    nanind = find(numNaN);
    nanind(nanind < 500) = [];
    capnan = find(numNaN(nanind) > size(dFF, 1)*0.2, 1);
    diffnanind = diff(nanind);
    split = find(diffnanind(1:capnan) > 1, 1, 'last');
    if isempty(split)
        split = capnan;
        [row col] = find(isnan(dFF(:, nanind(1:split))));
        dFF(unique(row), :) = []; 
    else
        [row col] = find(isnan(dFF(:, nanind(1:split))));
        dFF(unique(row), :) = []; 
    end
    
    if contains(lower(vidname), 'chat') %Determine genotype and corresponding index for pooled cells
        big1 = 1;
    else
        big1 = 2;
    end
    if contains(stim, 'astim') %Determine stimulation and corresponding index for pooled cells
        big2 = 1;
    elseif contains(stim, 'ostim')
        big2 = 2;
    elseif contains(stim, 'efs')
        big2 = 3;
    end
    
    count = size(bigNaN{big1, big2}, 1);
   
    %Calculate movement
    movemag = sqrt((xTraj-xTraj(:, 1)).^2+(yTraj-yTraj(:, 1)).^2);
    meanmag = mean(movemag);

    %Perform PCA
    [coeff,score,latent,tsquared,explained,mu] = pca((dFF')/max(dFF(:)),'Centered', 'off');
    [B, I] = sort(coeff(:, 1));

    %Store the data
    gdFF{indmat(1)}{indmat(2)} = sdFF;
    gnsdFF{indmat(1)}{indmat(2)} = nsdFF;
    gmag{indmat(1)}{indmat(2)} = meanmag;
    bigPCA{big1, big2}((1:3)+size(bigPCA{big1, big2}, 1), 1:size(dFF, 2)) = score(:, 1:3)';
    bigNaN{big1, big2}((1:size(dFF, 1))+count, 1:size(dFF, 2)) = dFF(I, :)/max(dFF(:));
    bigNorm{big1, big2}((1:size(dFF, 1))+count, 1:size(dFF, 2)) = nsdFF(I, :);
end

%% Plot the pooled dFF by genotype and stimulation with neurons that go out of focus prior to 500 frames removed
c = 1;
tl = {'Chat Astim', 'NOS Astim', 'Chat Ostim', 'NOS Ostim', 'Chat EFS',  'NOS EFS'};
for fn = 1:3
    for nn = 1:2
        subplot(3, 2, c)
        imagesc(bigNaN{nn, fn})
        title(tl{c})
        xlabel('Frame')
        ylabel('Cell number')
        c = c+1;
    end
end

%% Plot the percent of frames that are out of focus when neurons that go out of focus prior to frame 500 are removed
c = 1;
tl = {'Chat Astim', 'NOS Astim', 'Chat Ostim', 'NOS Ostim', 'Chat EFS',  'NOS EFS'};
for fn = 1:3
    for nn = 1:2
        subplot(3, 2, c)
        plot(sum(isnan(bigNaN{nn, fn}), 1)/size(bigNaN{nn, fn}, 1)*100)
        title(tl{c})
        hold on
        yline(20)
        xlabel('Frame')
        ylabel('Percent out of focus')
        c = c+1;
    end
end


%% Show the raster of the neural activity above threshold and plot the percent of neurons above threshold
c = 1;
tl = {'Chat Astim', 'NOS Astim', 'Chat Ostim', 'NOS Ostim', 'Chat EFS',  'NOS EFS'};
f1 = figure;
f2 = figure;
for fn = 1:3
    for nn = 1:2
        figure(f1)
        subplot(3, 2, c)
        block = bigNorm{nn, fn};
        raster = zeros(size(block));
        raster(block > 5) = 1; 
        imagesc(raster)
        title(tl{c})
        xlabel('Frame')
        ylabel('Cell number')
        figure(f2)
        subplot(3, 2, c)   
        plot((1:length(sum(raster, 1)))*0.05, sum(raster, 1)/size(raster, 1)*100)
        hold on
        xline(30)
        title(tl{c})
        xlabel('Time (s)')
        ylabel('Percentage cells active')
        ylim([0 50])
        c = c+1;
    end
end

%% Plot the first 2 PCs for each video within a given genotype/stimulation and the 3d plot of all 3 PCs. Figure by stimulation type
tl = {'Chat Astim', 'NOS Astim'; 'Chat Ostim', 'NOS Ostim'; 'Chat EFS',  'NOS EFS'};
f1 = figure;
f2 = figure;
f3 = figure;
for fn = 1:3
    if fn == 1
        figure(f1);
    elseif fn == 2
        figure(f2)
    elseif fn == 3
        figure(f3)
    end
    for nn = 1:2
        subplot(3, 2, 1+1*(nn-1))
        num = size(bigPCA{nn, fn}, 1);
        t = (1:size(bigPCA{nn, fn}, 2))*0.05;
        ntrace = num/3;
        plot(t, bigPCA{nn, fn}(1:3:num, :)')
        hold on
        xline(30)
        xlabel('Time (s)')
        ylabel('PC 1')
        title(tl{fn, nn})
        subplot(3, 2, 3+1*(nn-1))
        plot(t, bigPCA{nn, fn}(2:3:num, :)')
        hold on
        xline(30)
        xlabel('Time (s)')
        ylabel('PC 2')
        subplot(3, 2, 5+1*(nn-1))
        xc = [];
        yc = [];
        zc = [];
        for vid = 1:num/3
        x = bigPCA{nn, fn}(1+3*(vid-1), :);
        y = bigPCA{nn, fn}(2+3*(vid-1), :);
        z = bigPCA{nn, fn}(3+3*(vid-1), :);
        plot3(x, y, z)
        hold on
        xlabel('PC1')
        ylabel('PC 2')
        zlabel('PC 3')
        title('Registered PC evolution')
        end
        view([-37.5 30])     
    end
    hold off
end

%% Separate figure for each genotype and stimulation type plotting PC1 vs. PC2
c = 1;
gl = {'Chat', 'Astim'};
sl = {'Astim', 'Ostim', 'EFS'};
col = {'k', 'r'};
for fn = 1:3
    for nn = 1:2
        figure
        num = size(bigPCA{nn, fn}, 1);
        t = (1:size(bigPCA{nn, fn}, 2))*0.05;
        ntrace = num/3;
        for vid = 1:num/3
        x = bigPCA{nn, fn}(1+3*(vid-1), :);
        y = bigPCA{nn, fn}(2+3*(vid-1), :);
        plot(x, y)
        hold on
        xlabel('PC1')
        ylabel('PC 2')        
        end     
       sgtitle(strcat(gl{nn}, 32, sl{fn}))
       c = c+1;
    end
    hold off
end

%% Calculate and plot the PCA for matched neurons without removing neurons that go out of focus prior to frame 500
for fn = 1:length(group)
    dFF = gdFF{fn};
    nsdFF = gnsdFF{fn};
    m = matches{fn};
    if length(m) == 1
        figure
        M = m{1};
        i1 = 1; 
        i2 = 2;
        r = 4; 
        c = 1; 
        ind = 1:4;
        matchedPCAplot(dFF, M, sttype{fn}, i1, i2, r, c, ind)
        sgtitle(group{fn}{1})
    elseif length(m) == 3 && sum(cellfun(@(x) size(x,1), matches{fn})>2)==3
        figure
        M = m{1};
        i1 = 1; 
        i2 = 2;
        r = 4; 
        c = 3; 
        ind = 1:3:12;
        matchedPCAplot(dFF, M, sttype{fn}, i1, i2, r, c, ind)
        M = m{2};
        i1 = 1; 
        i2 = 3;
        r = 4; 
        c = 3; 
        ind = 2:3:12;
        matchedPCAplot(dFF, M, sttype{fn}, i1, i2, r, c, ind)
        M = m{3};
        i1 = 2; 
        i2 = 3;
        r = 4; 
        c = 3; 
        ind = 3:3:12;
        matchedPCAplot(dFF, M, sttype{fn}, i1, i2, r, c, ind)
        sgtitle(group{fn}{1})
    end
end

%% Create the table for Astim vs Ostim

percthresh1 = 0.1; %Threshold for the movement to identify MCs
avgframes = 10; %Window for a moving median to smooth the movement. Could be removed
mt = [1 2; 1 3; 2 3]; %The matching order
totaltable = table;
for fn = 1:length(group)
    stimulation = sttype{fn};
    stimind = [];
    m = matches{fn};
    if ~isempty(m)
        movement = gmag{fn};
        dFF = gdFF{fn};
        nsdFF = gnsdFF{fn};
        mouse = strcat(num2str(fn),'-',group{fn}{1}(1:6));
        if contains(lower(group{fn}{1}), 'nos') %Identify Genotype
            gtype = 'nos';
        else
            gtype = 'chat';
        end
        for t = 1:length(stimulation) %Determine which of the 3 videos are astim and ostim
            if length(stimulation{t}) == 5
                if all(stimulation{t} == 'astim')
                    stimind(1) = t;
                elseif all(stimulation{t} == 'ostim')
                    stimind(2) = t;
                end
            end
        end
        if length(stimind) == 2 && isempty(find(stimind == 0))
            adFF = nsdFF{stimind(1)}; %Astim dFF
            arast = zeros(size(adFF)); %Astim raster
            arast(adFF > 5) = 1;
            odFF = nsdFF{stimind(2)}; %ostim dFF
            orast = zeros(size(odFF)); %ostim raster
            orast(odFF > 5) = 1;
            amove = movement{stimind(1)}; %ID start of motor complex for astim
            amove = amove - mean(amove((end-20):end));
            abin = zeros(size(amove));
            abin(amove >= max(movmean(amove, avgframes))*percthresh1) = 1;
            astart = strfind([0 abin], [0 1]); %Merges blocks that have a small number of frames between them
            aend = strfind([abin 0], [1 0]);
            mergind = find(astart(2:end)-aend(1:end-1) < 10);
            astart(mergind+1) = [];
            aend(mergind) = [];
            omove = movement{stimind(2)}; %ID start of motor complex for ostim. Should probably just be its own separate function now
            omove = omove - mean(omove((end-20):end));
            obin = zeros(size(omove));
            obin(omove >= max(omove)*percthresh1) = 1;
            ostart = strfind([0 obin], [0 1]);
            oend = strfind([obin 0], [1 0]);
            mergind = find(ostart(2:end)-oend(1:end-1) < 10);
            ostart(mergind+1) = [];
            oend(mergind) = [];
            MCsa = astart(find(astart > 600, 1, 'first')); %MC start for astim
            MCea = aend(end); %MC end for astim
            MCso = ostart(find(ostart > 600, 1, 'first')); %MC start for ostim
            MCeo = oend(end); %End for ostim
            sortstimind = sort(stimind); %Determine which of the cells in the matched neurons corresponds to the data
            comparind = all(mt == sortstimind, 2);
            mc = m{comparind};
            [mval maxind] = max(stimind);
            aonly = 1:size(adFF, 1);
            oonly = 1:size(odFF, 1);
            if maxind == 2 %Get the correct order for the matching matrix, m, in case ostim occurs before astim in the list
                omatch = mc(:, 2);
                amatch = mc(:, 1);
            else
                omatch = mc(:, 1);
                amatch = mc(:, 2);
            end
            aonly(amatch) = []; %Make the list of neurons that occur only in astim or ostim
            oonly(omatch) = [];
            astart = sum(arast(:, 1:600), 2); %Sum the number of frames above threshold for astim or ostim for the 3 different time periods
            amid = sum(arast(:, 601:MCsa), 2);
            aend = sum(arast(:, MCea:end), 2);
            ostart = sum(orast(:, 1:600), 2);
            omid = sum(orast(:, 601:MCsa), 2);
            oend = sum(orast(:, MCea:end), 2);
            totalcell = length(aonly)+length(oonly)+length(omatch);
            aactivity = NaN(totalcell, 3);
            oactivity = NaN(totalcell, 3);

            %Make the list of matching neurons between astim and ostim.
            %Ordered so astim only goes first, matched between astim and
            %ostim second, and ostim only last in the list of neurons
            aactivity(1:length(aonly), :) = [astart(aonly) amid(aonly) aend(aonly)];
            aactivity((1:length(amatch))+length(aonly), :) = [astart(amatch) amid(amatch) aend(amatch)];
            oactivity((length(astart)+1):end, :) = [ostart(oonly) omid(oonly) oend(oonly)];
            oactivity((1:length(amatch))+length(aonly), :) = [ostart(omatch) omid(omatch) oend(omatch)];
            astartframes = repmat(600, totalcell, 1);
            amidframes = repmat(length(601:MCsa), totalcell, 1);
            aendframes = repmat(length(orast(1, MCea:end)),totalcell, 1);
            ostartframes = repmat(600, totalcell, 1);
            omidframes = repmat(length(601:MCso), totalcell, 1);
            oendframes = repmat(length(orast(1, MCeo:end)), totalcell, 1);
            repmouse = repmat({mouse}, totalcell, 1);
            repgtype = repmat({gtype}, totalcell, 1);
            cellnum = (1:totalcell)';

            %Make the table and set up the names for the mouse
            mousetable = table(repmouse, repgtype, cellnum, aactivity(:, 1),...
                astartframes, aactivity(:, 2), amidframes, aactivity(:, 3),...
                aendframes, oactivity(:, 1), ostartframes, oactivity(:, 2),...
                omidframes, oactivity(:, 3), oendframes);
            mousetable.Properties.VariableNames = {'Mouse', 'Genotype', 'Cell',...
                'astim baseline # frames active', 'astim baseline # frames', ...
                'astim active # frames active', 'astim active # frames',...
                'astim end # frames active', 'astim end # frames', ...
                'ostim baseline # frames active', 'ostim baseline # frames', ...
                'ostim active # frames active', 'ostim active # frames',...
                'ostim end # frames active', 'ostim end # frames'};

            %Concatenate to make the pooled table
            totaltable = [totaltable; mousetable];
        end
    end
end
        
%% All 3 matched
percthresh1 = 0.1;
avgframes = 10;
mt = [1 2; 1 3; 2 3];
c = 1;
numcell = []; %Number of cells that match
for fn = 1:length(group)
    stimulation = sttype{fn};
    stimind = [];
    m = matches{fn};
    names = group{fn};
    if length(m) == 3
        movement = gmag{fn};
        dFF = gdFF{fn};
        nsdFF = gnsdFF{fn};
        mouse = strcat(num2str(fn),'-',group{fn}{1}(1:6));
        if contains(lower(group{fn}{1}), 'nos')
            gtype = 'nos';
        else
            gtype = 'chat';
        end
        for t = 1:length(stimulation)
            if length(stimulation{t}) == 5
                if all(stimulation{t} == 'astim')
                    stimind(1) = t;
                elseif all(stimulation{t} == 'ostim')
                    stimind(2) = t;
                end
            else
                stimind(3) = t;
            end
        end
        mdFF = {}; 
        
        if length(stimind) == 3 && isempty(find(stimind == 0))
           m1 = m{1};
           m2 = m{2};
           mall1 = find(ismember(m1(:, 1), m2(:, 1)));
           mall2 = find(ismember(m2(:, 1), m1(:, 1)));
           ind1 = m1(mall1, 1);
           ind2 = m1(mall1, 2);
           ind3 = m2(mall2, 2);
           dFF1 = dFF{1};
           dFF1 = dFF1(ind1, :);
           dFF2 = dFF{2}; 
           dFF2 = dFF2(ind2, :);
           dFF3 = dFF{3};
           dFF3 = dFF3(ind3, :);
           mdFF{1} = dFF1;
           mdFF{2} = dFF2; 
           mdFF{3} = dFF3;
           catdFF = cat(2, mdFF{stimind(1)},mdFF{stimind(2)},mdFF{stimind(3)}); %Order Astim, Ostim, EFS
           l1 = length(mdFF{stimind(1)});
           l2 = length(mdFF{stimind(2)});
           subplot(2, 1, 1)
           imagesc(catdFF);
           colorbar
           xlabel('frame')
           ylabel('Cell number')
           title('Astim to Ostim to EFS')
           [coeff,score,latent,tsquared,explained,mu] = pca(catdFF','Centered', 'off');
           subplot(2, 1, 2)
           if size(catdFF, 1) >= 3
               x = score(:, 1);
               y = score(:, 2); 
               z = score(:, 3);
               plot(x(1:l1), y(1:l1))
               hold on
               plot(x((1:l2)+l1), y((1:l2)+l1))
               plot(x(l1+l2+1:end), y(l1+l2+1:end))
               legend({'Astim', 'Ostim', 'EFS'})
               hold off
           end
           numcell(c) = size(catdFF, 1);
           c = c+1;
           
        end
    end
end
cd(hmfolder)