% Author: Otar Akanyeti
% Date: 21 Sep 2021
% This program fits a second degree polynomial to amplitude envelopes of
% fishes and plots
%   1) histogram of r-square values
%   2) polynomial coefficient 0 versus polynomial coefficient 1
%   3) polynomial coefficient 0 versus polynomial coefficient 2
%   4) polynomial coefficient 1 versus polynomial coefficient 2
%   5) all amplitude envelopes (fish)
%   6) all amplitude envelopes (polynomial models)
% It also saves polynomial coefficients and their r-square values into a
% csv file.

% Input data
% dat.mat

% Output data
% outdat.csv


% variables
sfactor = 20;   % moving average window for smoothing

r2th = 0.90;    % r-square threshold

% load input data
load('dat.mat');

datnum = size(dat.ampenv,1);    % number of datasets
posnum = size(dat.ampenv,2);    % number of positions along the body

% create a position vector
x = (1:posnum)/posnum; x = x'; 

% smooth data
for k=1:datnum
    dat.ampenv(k,:) = smooth(dat.ampenv(k,:),sfactor);
end

% initialise output variables
regcoeff = nan(datnum,3);
r2val = nan(datnum,1);
ypredict = nan(datnum,posnum);
ypredict_nooffset = nan(datnum,posnum);

for n=1:datnum
   
    % output
    y = dat.ampenv(n,:)';
    
    % input
    xv = [x.^2, x, ones(length(x),1)];
    
    % fit a linear regression
    [b,~,~,~,stats] = regress(y,xv);
    
    regcoeff(n,1:3) = b;   % polynomial coefficients
    r2val(n) = stats(1);   % r-square value
    
    % calculate predicted output
    ypredict(n,:) = xv(:,1) * b(1) + xv(:,2) * b(2) + xv(:,3) * b(3); 
    
end


% plot r-square values histogram
hedges = 0.5:0.02:1; hbins = 0.51:0.02:0.99;
hn = histcounts(r2val,hedges);
fig = figure;
b = bar(hbins,hn);
b.EdgeColor = 'k';
b.FaceColor = 'k';
xlabel('r-square value'); ylabel('counts');


% select data with r-square value above the threshold
cind = r2val > r2th;

% plot polynomial coefficient 0 versus polynomial coefficient 1
fig = figure;
plot(regcoeff(cind,3),regcoeff(cind,2),'ko');
xlabel('coefficient 0'); ylabel(' coefficient 1')

% plot polynomial coefficient 0 versus polynomial coefficient 2
fig = figure;
plot(regcoeff(cind,3),regcoeff(cind,1),'ko');
xlabel('coefficient 0'); ylabel(' coefficient 2')

% plot polynomial coefficient 1 versus polynomial coefficient 2
fig = figure;
plot(regcoeff(cind,2),regcoeff(cind,1),'ko');
xlabel('coefficient 1'); ylabel(' coefficient 2');

% plot amplitude envelope
fig = figure;
plot(x,dat.ampenv(cind,:)','k');
xlabel('position along the body (L)'); 
ylabel('amplitude envelope (fish) (L)');

% plot amplitude envelope
fig = figure;
plot(x,ypredict','k');
xlabel('position along the body (L)'); 
ylabel('amplitude envelope (predicted) (L)');

% save results in a csv file
clipid = dat.clipid;
r2 = r2val;
c0 = regcoeff(:,3);
c1 = regcoeff(:,2);
c2 = regcoeff(:,1);
out = table(clipid,r2,c0,c1,c2);
writetable(out,'outdat.csv');