%% Main tracking code for analyzing the position of cockroaches marked with BEEtags within the circular experimental arena

%Get filename from user
[filename pathname] = uigetfile('*.avi');
mov = VideoReader(strcat(pathname,filename));
nframes = mov.NumberOfFrames;

% Where to save output data?
%savepath = 'C:\Users\OEB131-E\Dropbox\Data\';


load roachCohorts %% Need a matlab file in the path with list of roaches belonging to each experimental cohort
[startFrame cohort] = roachAnalyzeMovie(filename); %Extract cohort and start frame data from movie

%Figure out what codes belong in current cohort
codeind = zeros(size(cohorts,1),1);

for i = 1:length(codeind)
    try
        c = mean(char(cohorts(i,2)) == char(cohort));
    catch
        continue
    end
    
    if  c == 1;
        codeind(i) = 1;
    end
end

%Generate list of valid codes to look for
codelist = [cohorts{codeind == 1,1}];

%Create empty output structure
movData = struct();

r = 46; %Radius of black region trackers
angles = 2:2:360;
thr = 0.3;
ft = 0; %Apply gaussian high pass filter?
%% Create empty data output matrix
load('roachMask.mat');
anglesOutput = nan(nframes, numel(angles));
%

%Pre-allocate data frames for output
%codelist = [movData(105).R.number]; %To be replaced with list of numbers in this cohort
roachAngles = nan(nframes,numel(codelist));
roachRadius = nan(nframes,numel(codelist));
roachOrientation = nan(nframes,numel(codelist));

kymograph=zeros(nframes,size(angles,1));

%% Insert code for finding best threshold here
ff = startFrame;
numCalibFrames=10;
threshCalibVals=linspace(0.5,0.7,20);
testFrames=ceil(linspace(ff,nframes,numCalibFrames));
calibrationResults=zeros(numel(threshCalibVals),numCalibFrames);


%% Get frame center
calFrame = ff + 1000;

xl = [1100 1500];
yl = [650 1100];
im = read(mov,calFrame);
imshow(im);
title('Click on center of projection (where the black cones meet)')
xlim(xl);
ylim(yl);
[xc yc button] = ginput(1);
%If this frame has a cockroach covering it, skip forward 500 frames and
%overwrite the value
if char(button) == 's'
    calFrame = calFrame+500;
    im = read(mov,calFrame);
    imshow(im);
    title('Click on center of projection (where the black cones meet)')
    xlim(xl);
    ylim(yl);
    [xc yc button] = ginput(1);
end
hold off
%%
for i=1:length(threshCalibVals)
    for j=1:length(testFrames)
        
        
        %         if i == 1 && j==1
        %             im = read(mov, ff + 1000);
        %             title('Click on center of projection (where the black cones meet)')
        %             xlim([1000 1600]);
        %             ylim([600 1200]);
        %             [xc yc] = ginput(1);
        %             hold off
        %             imshow(uint8(imtmp));
        %         end
        
        im = read(mov,testFrames(j));
        im = im(:,:,1); %Get only a single frame
        brightnessFactor=160;
        im=uint8(double(im)/mean(mean(im))*brightnessFactor);
        imtmp=double(im).*(128./double(mask));
        imshow(uint8(imtmp));
        %% Track codes
        imt = imtmp;
        imt(imt > 255) = 255;
        try
            R = locateCodes(uint8(imt),0,threshCalibVals(i),0,1,400); %Threshold value is the third value here
        catch
            continue
        end
        calibrationResults(i,j)=length(R);
        
    end
end
calibSums=sum(calibrationResults');
whichThresh=find(calibSums==max(calibSums));
whichThresh=round(median(whichThresh));
bestThresh=threshCalibVals(whichThresh);


%%
for i = startFrame:nframes
    %% load and plot frame
    disp(i);
    im = read(mov,i);
    im = im(:,:,1); %Get only a single frame
    brightnessFactor=160;
    im=uint8(double(im)/mean(mean(im))*brightnessFactor);
    imtmp=double(im).*(128./double(mask));
    %     im(im==255)=0;
    imshow(uint8(imtmp));
    
    %% Track stimulus regions
    [x y] = pol2cart(deg2rad(angles),repmat(r,1,numel(angles)));
    x = round(x + xc);
    y = round(y + yc);
    
    pixVals = [];
    for aa = 1:numel(x)
        pixVals(aa) = double(im(y(aa),x(aa)))/255;
    end
    ind = pixVals < thr;
    
    % Output data
    anglesOutput(i,:) = ind;
    %% Look for codes within image
    imt = imtmp;
    imt(imt > 255) = 255;

    try
        R = locateCodes(uint8(imt),0,bestThresh,0,1,400); %Threshold value is the third value here
    catch
        continue
    end
    %end
    
    %% Extract theta for each code
    for zz = 1:numel(R)
        xd = R(zz).Centroid(1) - R(zz).frontX;
        yd = R(zz).Centroid(2) - R(zz).frontY;
        [th rad] = cart2pol(xc, yd);
        R(zz).orientationTh = th;
    end
    %% PLot tracking results
    hold on;
    plot(x(ind), y(ind),'r.');
    
    movData(i).R = R;
    
    kymograph(i,ind)=1;
    
    %Extract and reformat tracking data
    for j=1:size(R,1)
        [theta radialDist]=cart2pol(R(j).Centroid(1)-xc,R(j).Centroid(2)-yc);
        roachAngles(i,find(R(j).number == codelist))= theta;
        roachRadius(i,find(R(j).number == codelist))= radialDist;
        roachOrientation(i,find(R(j).number == codelist))= R(zz).orientationTh;
        
        
    end
    
    hold off;
    
    %Every few frames, save matlab data format
    if mod(i, 5) == 0;
        save(strcat(pathname,filename, '.mat'), 'roachAngles', 'roachRadius', 'roachOrientation', 'kymograph', 'xc', 'yc', 'codelist', 'movData');
        sprintf(strcat('saving frame ', num2str(i), ' of ', num2str(nframes)))
    end
    
end