%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% This file was used for two purposes:
%% 1. clean up raw data from the lab
%% 2. embed the omega data to produce phase-space reconstruction
%%
%% Still very dodgy for many reasons:
%% 1. Haven't thought about camera angles (camera projection matrix)
%% 2. There is a manual step in cleaning up the theta time series
%% 3. No false nearest neighbour analysis yet (but could be done
%%    so long as file is not too big.
%%

clear all
close all

%% The embedding should only be done on cleaned data: set embed = 1 to
%% move on to the embedding step, otherwise the script will go no
%% further than calculating omega by numerically differentiating theta
embed     = 1;

%% set time interval between frames (seconds)
dt        = 0.5; 

%% read in raw data
raw_data  = dlmread('0630_12cups_xy.dat.txt');

%% create time vector
raw_time  = [1:size(raw_data,1)]'*dt;

%% pull out raw x and y data, and convert to r and theta
raw_x     = raw_data(:,1)-mean(raw_data(:,1));
raw_y     = raw_data(:,2)-mean(raw_data(:,2));
raw_r     = sqrt(raw_x.^2+raw_y.^2);
raw_theta = atan2(raw_y,raw_x);

%% find the indices for the dropped frames
[nonzero_index,ignore] = find(raw_data(:,1)|raw_data(:,2));

%% remove the zero points
cleaner_data           = raw_data(nonzero_index,:);
cleaner_time           = raw_time(nonzero_index);

%% estimate the wheel centre from mean x and y position
guess_centre           = mean(cleaner_data)';

%% find the centre that minimises the sum of squares difference
%% between r and the mean r
centre                 = fminsearch(@radius_error,guess_centre,[],cleaner_data);

%% recalculate x and y using new centre estimate
first_pass             = cleaner_data-repmat(centre',size(cleaner_data,1),1);

%% calculate r and theta based on new centre estimate
r                      = sqrt(first_pass(:,1).^2+first_pass(:,2).^2);
theta                  = atan2(first_pass(:,2),first_pass(:,1));

%% find the r values that are within 1 std deviation of mean and reject the rest
good_r_index           = find((r<mean(r)+std(r))&(r>mean(r)-std(r)));
even_cleaner_data      = cleaner_data(good_r_index,:);

%% estimate the centre again and do the minimising thing again
better_guess_centre    = mean(even_cleaner_data)';
better_centre          = fminsearch(@radius_error,better_guess_centre,[],even_cleaner_data);

%% calculate final x and y values using new centre estimate
final_data             = even_cleaner_data-repmat(better_centre',size(even_cleaner_data,1),1);
time                   = cleaner_time(good_r_index);

%% calculate r and theta
r                      = sqrt(final_data(:,1).^2+final_data(:,2).^2);
theta                  = atan2(final_data(:,2),final_data(:,1));


%% Find spikes in data to indicate where theta is > pi or < pi
%% The spikes are places where there is a difference exceeding
%% pi (plus a little bit) in the theta signal over one time
%% step.
pi_t                  = pi+pi/4;
spike_spots           = [0;(abs(diff(theta))>pi_t)];
spike_index           = [find(spike_spots);length(spike_spots)];
spike                 = -round([0;(abs(diff(theta))>pi_t).*diff(theta)]/2/pi);

num_spikes            = sum(abs(spike_spots));
spike_free_theta      = theta;

%% I'm too tired to comment the next bit. We're searching through for
%% chunks of the time series that need multiples of 2*pi added to them
%% We can't automate the detection of all such chunks, so there is
%% a need to direct it manually at the end

n = 0;
for ii = 1:num_spikes
    chunk                   = [spike_index(ii):spike_index(ii+1)];
    n                       = n+spike(spike_index(ii));
    spike_free_theta(chunk) = theta(chunk)+2*pi*n*abs(spike(spike_index(ii)));
end

diffspike2            = sign([0;diff(spike_free_theta)]);
diffspike_orig        = diffspike2;

spike_free_theta2     = spike_free_theta;

for ii = 1:length(diffspike2)-3
    if (diffspike2(ii:ii+3) == [-1,1,-1,1]')     % rightside-up N
        %shift up
        spike_free_theta2(ii+2:end) = spike_free_theta2(ii+2:end)+2*pi;
        diffspike2                  = sign([0;diff(spike_free_theta2)]);
    elseif (diffspike2(ii:ii+3) == [1,-1,1,-1]') % upside-down  N
        %shift down
        spike_free_theta2(ii+2:end) = spike_free_theta2(ii+2:end)-2*pi;
        diffspike2                  = sign([0;diff(spike_free_theta2)]);
    end
end

 for ii = 1:length(diffspike2)-2
    if (diffspike2(ii:ii+2) == [1,-1,1]')     % rightside-up N
        %shift up
        spike_free_theta2(ii+1:end) = spike_free_theta2(ii+1:end)+2*pi;
        diffspike2                  = sign([0;diff(spike_free_theta2)]);
    elseif (diffspike2(ii:ii+2) == [-1,1,-1]') % upside-down  N
        %shift down
        spike_free_theta2(ii+1:end) = spike_free_theta2(ii+1:end)-2*pi;
        diffspike2                  = sign([0;diff(spike_free_theta2)]);
    end
end
           
    
%% manual spots for this file
%% THESE NEED TO BE FOUND AND SPECIFIED MANUALLY FOR EACH NEW DATA FILE

start    = [505 602 610 611 815 816 1277 length(spike_free_theta2)+1];
updown   = [2*pi 4*pi 6*pi 8*pi 6*pi 4*pi 6*pi];

for ii=1:length(start)-1
    spike_free_theta2(start(ii):start(ii+1)-1) = spike_free_theta2(start(ii):start(ii+1)-1)+updown(ii);
end


if embed~=1 
    figure(1)
    plot(final_data(:,1),final_data(:,2),'o',raw_data(:,1)-guess_centre(1),raw_data(:,2)-guess_centre(2),'+')
    legend('clean data','raw data')
    xlabel('x')
    ylabel('y')
    
    figure(2)
    plot(time,theta,'-o',time,spike,'+',[time(1),time(end)],[pi,pi],[time(1),time(end)],[-pi,-pi])
    title('\theta time series before correction.')
    xlabel('t')
    ylabel('\theta')
    legend('\theta time seris','spike location and direction')
    
    figure(3)
    plot(time, spike_free_theta,'-o',time, spike_free_theta2,'-o',time,diffspike_orig,'-o',time,diffspike2,'-o')
    xlabel('t')
    ylabel('\theta')
    legend('1st spike removal pass','2nd spike removal pass','original diff','final diff')
    
    figure(4)
    plot(spike_free_theta2,'-o')
    title('final \theta time series')
    xlabel('t')
    ylabel('\theta')
    return
end

%% interpolate theta to a finer time step to get
%% derivative estimates
dt_fine       = dt/10;
t_fine        =[5:dt_fine:max(raw_time)]';
interp_theta  = interp1(time,spike_free_theta2,t_fine,'spline');

%% calculate omega by numerically differentiating
%% use centred differencing: dydt~(y(t+dt)-y(t-dt))/(2*dt)
%% this method is O(dt^2)
%% An alternative centred differencing technique O(dt^4) is:
%% dydt~-1/(12*dt)*(y(t+2*dt)-y(t-2*dt))+2/(3*dt)*(y(t)-y(t-dt))
omega         = (interp_theta(3:end)-interp_theta(1:end-2))/(2*dt_fine);
t             = t_fine(3:end);

%% embed omega data
numbins       = 100;
tau           = modified_estimateTauAMI(numbins, omega, 1);
%d_e          = FNN(omega,tau);  <----use this only if omega is not very big (ask Aaron how big?)
d_e           = 3;

n        = floor(length(omega)/tau)*tau;
embedded = zeros(n-(d_e-1)*tau,d_e);
for k=1:d_e
    embedded(:,k) = omega((k-1)*tau+1:n-(d_e-k)*tau);
end

figure (5)
plot(t,omega)
ylabel('\omega')
xlabel('t')

figure(6)
plot3(embedded(:,1),embedded(:,2),embedded(:,3))
title(strcat('Delay Coordinate Embedding; d_{E} = ', int2str(d_e), '; \tau = ', num2str(tau*dt_fine),' s'));
xlabel('\omega_{t}');
ylabel(strcat('\omega_{t+\tau}'));
zlabel(strcat('\omega_{t+2\tau}'));