% Postprocessing
% read tracked data and plot the traces, and distances

% Input
% calibration.mat : read the chamber position
%        *-bg.mat : read background image
%      *-track.mat: read position data of each flies

clear
close all
%% Parameter setting
%The logic to set parameters:
% 1. filter_win---first smooth traces to reduce small movement
% 2. t_gap---the fly cannot go to another (far) place and then return in 3 time units.
%    By setting this parameter, we remove the movement due to the small body vibration
% 3. static_thres_pixel---if the fly does not even move 0.5 pixel in 3 time
%    units, then we think it does not move at all
filter_win=3;          % [pixel] must be an odd number. Length for trace smoothing,larger then trace is smoother
t_gap =1;              % [dt] time interval to calculate averaging speed, larger then less speed spikes but not accurate
static_thres_pixel =1;% [pixel] if displacement is smaller than 3 pixel then we think it does not move at all.

%% open file
[file_cali, path_cali] = uigetfile(fullfile('./','*calibration.mat'),'Select the calibration data','MultiSelect','off');
[file_bg, path_bg] = uigetfile(fullfile(path_cali,'*-bg.mat'),'Select the background data','MultiSelect','off');
[file_trk, path_trk] = uigetfile(fullfile(path_bg,'*-track.mat'),'Select the track data','MultiSelect','off');
file_segs= split(file_trk,'-');%extract file name for saving
fileprefix = file_segs{1};
for m=2:numel(file_segs)-1
    fileprefix = join([string(fileprefix), string(file_segs{m})],'-');
end

% load data
load(fullfile(path_cali,file_cali));
load(fullfile(path_bg,file_bg));
load(fullfile(path_trk,file_trk));

fly_index = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18]; %remove the index to exclude 

%% Clean data and calculate distance
dt = 1/calib.FPS;
pixel2mm = 1/calib.PPM;                           %[mm/pixel]

t = dt*(0: size(trk.data,2)-1);
npts = numel(t);
data_smooth = zeros(size(trk.data,1),numel(t), 2);% fly-t-[px,py]
data_stat   = 0;                                  %[t, displacement, speed]
trk_backup = trk;

%clean data
for mm=1:numel(fly_index)%number of flies
    kk=fly_index(mm);           %target fly
    if (isnan(trk.data(kk,1,1)))%check the first x position     ,if lost
        trk.data(kk,1,1) = mean(trk.data(kk,1:20,1),'omitnan');%take the average of the first 10 frames as the lost data
    end
    if (isnan(trk.data(kk,1,2)))%check the first y position
        trk.data(kk,1,2) = mean(trk.data(kk,1:20,2),'omitnan');
    end
    for itt=2:numel(t)
        if isnan(trk.data(kk,itt,1));%x
            trk.data(kk,itt,1)=trk.data(kk,itt-1,1);%use last position
        end
        if isnan(trk.data(kk,itt,2));%y
            trk.data(kk,itt,2)=trk.data(kk,itt-1,2);%use last position
        end
        1;
    end
    data_smooth(kk, :,1) = mov_filter(trk.data(kk,:,1),filter_win);% smooth pos_x
    data_smooth(kk, :,2) = mov_filter(trk.data(kk,:,2),filter_win);% smooth pos_y
    mm=1;
    for itt=1:t_gap: numel(t)
        if  itt+t_gap<npts
            %dist1 =sqrt((trace(2).x(kk) - trace(2).x(kk-1))^2 + (trace(2).y(kk) - trace(2).y(kk-1))^2)*pixel2mm;
            dist1 =sqrt((data_smooth(kk,itt,1) - data_smooth(kk,itt+t_gap,1))^2 + (data_smooth(kk,itt,2) - data_smooth(kk,itt+t_gap,2))^2 )*pixel2mm;
            v_dist1=dist1/dt;
            if v_dist1<4
                dist1 = 0;
            end
            data_stat(kk,mm,1) = t(round(itt+t_gap/2));% [s] time point for velocity
            data_stat(kk,mm,2) = dist1;                % [mm] distance
            assert(~isnan(dist1));
            data_stat(kk,mm,3) = dist1/(t_gap*dt);     % [mm/s] average speed in dt*t_gap interval
            mm = mm+1;
        end
    end
end

%only calculate timepoints when flies are MOVING
% !!!!!!! use this only when t_gap = 1;otherwise results might be wrong !!!!!!!!
velMov = zeros(1,size(data_stat,1));
for kk=1:size(data_stat,1)
    iMov = data_stat(kk,:,2)>0;
    sMov = data_stat(kk,:,2)<0.001;
    distMov = sum(data_stat(kk,:,2),'omitnan');
    tMov  = sum(iMov)*dt;
    sMov  = sum(sMov)*dt;
    velMov(kk) = distMov / tMov;
    sMovArray(kk)=sMov;
end


data_dist = sum(data_stat(:,:,2),2);                   % [mm] total displacement
fprintf('Fly                     :\t');fprintf('%5d\t',1:numel(data_dist));fprintf('\n');
fprintf('Dist                    :\t');fprintf('%5.2f\t',data_dist);fprintf('[mm]\n');
fprintf('Ave_speed               :\t');fprintf('%5.2f\t',mean(data_stat(:,:,3),2));fprintf('[mm/s]\n');
fprintf('Moving_speed            :\t');fprintf('%5.2f\t',velMov);fprintf('[mm/s]\n');
fprintf('Immobile_time           :\t');fprintf('%5.2f\t',sMovArray);fprintf('[s]\n');
%% Plot
figure(1)
color_line = {'r','y','b'};
imshow(bg.bg_mean);
hold on;
for mm=1:numel(fly_index)
    kk = fly_index(mm);   %only plot flies in the list
    plot(trk.data(kk,:,1), trk.data(kk,:,2),'Color',color_line{mod(kk+2,3)+1});
end
%exportgraphics(gca,fullfile(path_trk,[fileprefix,'_trak_all.jpg']),'Resolution',600);



%% test



%%
function out = mov_filter(in, win_len)
out = in;
N = length(in);
win_len_half = (win_len-1)/2;
for kk= win_len_half+1: N-win_len_half
    out(kk) = mean(in(kk-win_len_half:kk+win_len_half),'omitnan');
    %out(kk) = 0.1*in(kk-2) + 0.3*in(kk-1) + 0.6*in(kk);
end
end