%% v8: 
% edited on 20150203
% 1. in v7, we don't need to track the reward location before the recording,
%    but rather, we can track the reward location only based on where the rat
%    stops in sample trials. However, this doesn't always work because
%    sometimes it's hard to detect if the rat is heading out to get food.
%    So in v8, we have to track the reward locations first. And then in sample
%    trials, it's easier to figure out where the rat tends to stop.
% 2. in v8, we don't need to judge if the rat is heading out
% 3. the defalt running direction is 'clockwise'. If the rat runs in
%    counterclockwise direction, we rotate the data to clockwise, i.e. 
%    the angle will be 2pi-angle
% 
% the experiment is setting as:
%
% In sample trial, a piece of tape at reward location. The rat has to run 6 
% laps, and stop at the tape in each lap to get the reward. 
% At the beginning of the sample trial, the rat was sitting in starting-zone for
% 30s, to let me sit down and close the curtain. At the end of this trial,
% he needs to be back to starting-zone, and then sitting for another 30s
% right before the test trial.
%
% In the test trial, no tape was on the track. If he stops at correct
% spot, i.e. ~10 degress close to the spot, the correct tone will go off. 
% Otherwise, if he stops outside ~10 degree but inside 90 degree around
% the spot, or run out of the spot without stopping, 
% the incorrect tone will go off. 
% At the end of test trial, the rat needs to be back to starting-zone. When
% he leave the track, this session ends. 


% ******* Importantly, we get the reward location from his behavior.
% That is, to record the location where he stops to get food

% set the missing date [0,0] as the same as the last visited data

%%
ang_RewardLoc_ontrack=Ang_RewardLoc_ontrack(ind_rewardloc);

% detect if the rat arrived at Starting-zone, if so, this session starts
locationArray_end=[];
while 1
    %Pause to let some new data load up in NetCom's buffers.
    pause(sample_time);
    %Request all new VT data that has been acquired since the last pass.
    [~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
    %Get the X and Y location value for this VT frame. The
    %locationArray is in the form [x1, y1, x2, y2, ... xN, yN].
    if ~isempty(locationArray)
        x_loc = locationArray(end - 1);
        y_loc = locationArray(end);
        if ~x_loc && ~y_loc %% if x_loc=0, y_loc=0, means it's a missing data from the video
            if isempty(locationArray_end)
                continue
            end
            x_loc=locationArray_end(1);
            y_loc=locationArray_end(2);
        end
        x0 = double(x_loc).*scale_x-x_center;
        y0 = double((-1)*y_loc).*scale_y-y_center;
        
        if x0>=0
            if x0==0 x0=0.01; end
            ang0=mod(atan(-y0/x0),2*pi);
        else
            ang0=mod(atan(-y0/x0)+pi,2*pi);
        end
        if strcmp(rot,'Counterclockwise');
            ang0=2*pi-ang0;
        end
        ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
        locationArray_end=[x_loc,y_loc];
        distance_center=sqrt(x0^2+y0^2);
        
        if distance_center < R_outer && (ang_ontrack > Ang_StartZone_arrive_ontrack - ang_startzone_border || ang_ontrack < Ang_StartZone_depart_ontrack + ang_startzone_border)
            break
        end
    end
end
sound(beep0,fs_sin);


%% sample trial

[~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
ts_sample_start1=timeStampArray(end);
pause(time_fix);
% give me 30s to sit into the track, close the curtain and prepare the
% treat
[~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
ts_sample_start2=timeStampArray(end);
sound(sound_beep_start,sound_beep_fs);

n_sample;
n_reward_lap=0;
n_running_lap=0;
locationArray_end=[];
ang_ontrack_end=0;
sign_enter=0;  % sign_enter=1 means the rat is ready to enter a reward spot;
% otherwise, sign_enter=0 means the rat is just leaving a reward spot
sign_leave=0;  % sign_leave=1 means the rat is just leaving a reward spot;
% otherwise, sign_enter=0 means the rat is just leaving a reward spot
% Note: when the rat is back to starting zone, both sign_enter and
% sign_leave are set to 0
sign_getreward=0;  % sign_reward=1 means the rat has found the reward location,
% and ready to eat the reward. (the rat headed out toward the center)
% otherwise, sign_reward=0 means the rat is not at reward location yet

time_stay_reward=0;
ang_sample_reward_actual=zeros(n_sample,1);
ang_sample_reward_ontrack=zeros(n_sample,1);
time_sample_reward=zeros(n_sample,1);
while n_reward_lap<n_sample
    % get the current location
    pause(sample_time);
    [~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
    time0=timeStampArray(end);
    x_loc = locationArray(end - 1);
    y_loc = locationArray(end);
    if ~x_loc && ~y_loc %% if x_loc=0, y_loc=0 means it's a missing data from the video
        if isempty(locationArray_end)
            continue
        end
        x_loc=locationArray_end(1);
        y_loc=locationArray_end(2);
    end
    x0 = double(x_loc).*scale_x-x_center;
    y0 = double((-1)*y_loc).*scale_y-y_center;
    
    if x0>=0
        if x0==0 x0=0.01; end
        ang0=mod(atan(-y0/x0),2*pi);
    else
        ang0=mod(atan(-y0/x0)+pi,2*pi);
    end
    if strcmp(rot,'Counterclockwise');
        ang0=2*pi-ang0;
    end
    ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
    vel_ontrack=(ang_ontrack-ang_ontrack_end)/sample_time;
    
    % check the current location
    if vel_ontrack<0 && (ang_ontrack > Ang_StartZone_arrive_ontrack - ang_startzone_border || ang_ontrack < Ang_StartZone_depart_ontrack + ang_startzone_border)
        % rat is walking back to starting zone, almost happens when the
        % rat starts and then walk back a little bit.
        sign_enter=0;
        sign_leave=0;
        sign_getreward=0;
    else
        % the rat is walking forward
        % or the rat is walking back, but not in the starting zone
        if sign_enter==0 && sign_leave==0 && vel_ontrack>0 && ang_ontrack > ang_RewardLoc_ontrack - ang_correct_zone && ang_ontrack < ang_RewardLoc_ontrack + ang_correct_zone
            % as soon as the rat is starting from the starting zone, and
            % enter the reward location zone
            sign_enter=1;
            sign_leave=0;
            time_stay_reward=0; % reset it to 0
        elseif sign_enter==1 && sign_leave==0 && ang_ontrack > ang_RewardLoc_ontrack - ang_correct_zone && ang_ontrack < ang_RewardLoc_ontrack + ang_correct_zone
            % if the rat is already in the reward location zone
            if time_stay_reward >= time_stay_sample
                % if the rat has stopped at reward location for "time_stay_sample"
                % mark here as a reward location
                n_reward_lap=n_reward_lap+1;
                ang_sample_reward_actual(n_reward_lap)=ang0;
                ang_sample_reward_ontrack(n_reward_lap)=ang_ontrack;
                fprintf('Loc: %s \r',num2str(ang_ontrack));
                time_sample_reward(n_reward_lap)=time0;
                sign_enter=0;
                sign_leave=1;
                sign_getreward=1;
            else % time_stay_reward < time_stay_sample
                if abs(ang_ontrack-ang_ontrack_end) <= ang_nomving && abs(ang_ontrack-ang_RewardLoc_ontrack) <= ang_correct_tolerance
                    % if the rat is still close to the reward location, and he is staying here
                    time_stay_reward=time_stay_reward+sample_time;
                else
                    time_stay_reward=0; % reset it to 0
                end
            end
        elseif ang_ontrack > Ang_StartZone_arrive_ontrack - ang_startzone_border || ang_ontrack < Ang_StartZone_depart_ontrack + ang_startzone_border
            % the rat is back into starting-zone
            if sign_enter || sign_leave || sign_getreward
                % if the rat just run back to the starting zone after a lap
                % count the number of running laps
                n_running_lap=n_running_lap+1;
                
                if n_reward_lap < n_running_lap && sign_getreward==0
                    %in this lap, the rat runs without stop
                    n_reward_lap=n_reward_lap+1;
                    ang_sample_reward_actual(n_reward_lap)=NaN;
                    ang_sample_reward_ontrack(n_reward_lap)=NaN;
                    fprintf('Loc: NaN\r');
                    time_sample_reward(n_reward_lap)=NaN;
                end
            end
            sign_enter=0;
            sign_leave=0;
            sign_getreward=0;
            time_stay_reward=0; % reset it to 0
        end
    end
    locationArray_end=[x_loc,y_loc];
    ang_ontrack_end=ang_ontrack;
end

% as soon as the rat is back into starting-zone, this sample trial ends
while ang_ontrack > Ang_StartZone_depart_ontrack && ang_ontrack < Ang_StartZone_arrive_ontrack
    % if rat is still running on the track
    pause(sample_time);
    [~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
    time0=timeStampArray(end);
    x_loc = locationArray(end - 1);
    y_loc = locationArray(end);
    if ~x_loc && ~y_loc %% if x_loc=0, y_loc=1, means it's a missing data from the video
        if isempty(locationArray_end)
            continue
        end
        x_loc=locationArray_end(1);
        y_loc=locationArray_end(2);
    end
    x0 = double(x_loc).*scale_x-x_center;
    y0 = double((-1)*y_loc).*scale_y-y_center;
    
    if x0>=0
        if x0==0 x0=0.01; end
        ang0=mod(atan(-y0/x0),2*pi);
    else
        ang0=mod(atan(-y0/x0)+pi,2*pi);
    end
    if strcmp(rot,'Counterclockwise');
        ang0=2*pi-ang0;
    end
    ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
    locationArray_end=[x_loc,y_loc];
end
ts_sample_stop=time0;
sound(beep0,fs_sin);

%% get the correct reward location from the sample trial
% check if there's a wrong loc in ang_sample_reward_ontrack
ang_sample_reward_ontrack0=ang_sample_reward_ontrack(~isnan(ang_sample_reward_ontrack));
std0=nanstd(ang_sample_reward_ontrack0);
if std0 > 0.2
    % there's at lease one loc is not close to others
    [ang_temp,ind]=sort(ang_sample_reward_ontrack0);
    n=length(ang_sample_reward_ontrack0);
    mu=mean(ang_temp(2:n-1));
    if abs(ang_sample_reward_ontrack0(1)-mu)>abs(ang_sample_reward_ontrack0(n)-mu)
        mu0=mean(ang_sample_reward_ontrack0(ind(2:n)));
    else
        mu0=mean(ang_sample_reward_ontrack0(ind(1:n-1)));
    end
    ang_sample_reward_ontrack_mean=mu0;
else
    ang_sample_reward_ontrack_mean=nanmean(ang_sample_reward_ontrack0);
end
if isempty(ang_sample_reward_ontrack0)
    ang_sample_reward_ontrack_mean=ang_RewardLoc_ontrack;
end
ang_test_reward_correct_left=max(ang_sample_reward_ontrack_mean-ang_correct_tolerance,Ang_StartZone_depart_ontrack);
ang_test_reward_correct_right=min(ang_sample_reward_ontrack_mean+ang_correct_tolerance,Ang_StartZone_arrive_ontrack);
ang_test_reward_incorrect_left=max(ang_sample_reward_ontrack_mean-ang_incorrect_tolerance,Ang_StartZone_depart_ontrack);
ang_test_reward_incorrect_right=min(ang_sample_reward_ontrack_mean+ang_incorrect_tolerance,Ang_StartZone_arrive_ontrack);


%% test trial

ts_test_start1=ts_sample_stop;
pause(time_fix);
% give me 30s to swipe the track and prepare the treat
[~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
ts_test_start2=timeStampArray(end);
sound(sound_beep_start,sound_beep_fs);

n_test;
n_reward_lap=0;
locationArray_end=[];
time_stay_reward=0;
ang_ontrack_end=0;
sign_enter=0;
sign_leave=0;
sign_getreward=0;
ang_test_reward_ontrack=0;
time_test_reward=0;
while n_reward_lap<n_test
    % get the current location
    pause(sample_time);
    [~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
    time0=timeStampArray(end);
    x_loc = locationArray(end - 1);
    y_loc = locationArray(end);
    if ~x_loc && ~y_loc %% if x_loc=0, y_loc=0 means it's a missing data from the video
        if isempty(locationArray_end)
            continue
        end
        x_loc=locationArray_end(1);
        y_loc=locationArray_end(2);
    end
    x0 = double(x_loc).*scale_x-x_center;
    y0 = double((-1)*y_loc).*scale_y-y_center;
    
    if x0>=0
        if x0==0 x0=0.01; end
        ang0=mod(atan(-y0/x0),2*pi);
    else
        ang0=mod(atan(-y0/x0)+pi,2*pi);
    end
    if strcmp(rot,'Counterclockwise');
        ang0=2*pi-ang0;
    end
    ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
    vel_ontrack=(ang_ontrack-ang_ontrack_end)/sample_time;
    
    % check the current location
    if vel_ontrack<0 && (ang_ontrack > Ang_StartZone_arrive_ontrack - ang_startzone_border || ang_ontrack < Ang_StartZone_depart_ontrack + ang_startzone_border)
        % rat is walking back to starting zone, almost happens when the
        % rat starts and then walk back a little bit.
        sign_enter=0;
        sign_leave=0;
        sign_getreward=0;
    elseif sign_leave==0 && ang_ontrack > ang_test_reward_incorrect_left && ang_ontrack < ang_test_reward_incorrect_right
        % the rat is walking in reward zone
        % as soon as the rat is starting from the starting zone
        sign_enter=1;
        if time_stay_reward >= time_stay_test
            % if the rat has stopped at reward location for "time_stay_test"
            % mark the his "turning head"-location as his final stop location
            if ang_ontrack > ang_test_reward_correct_left && ang_ontrack < ang_test_reward_correct_right
                % the rat stops at reward location
                sound(tone_correct,fs_tone_correct);
                sign_correct=1;
            else % the rat make an error infront of or behind the reward location
                sound(tone_incorrect,fs_tone_incorrect);
                sign_correct=0;
            end
            n_reward_lap=n_reward_lap+1;
            ang_test_reward_actual=ang_test_reward_actual_temp;
            ang_test_reward_ontrack=ang_test_reward_ontrack_temp;
            time_test_reward=time_test_reward_temp;
            sign_enter=0;
            sign_leave=1;
            sign_getreward=1;
        else % time_stay_reward < time_stay_test
            if sign_getreward==0
                if abs(ang_ontrack-ang_ontrack_end) <= ang_nomving
                    sign_getreward=1;
                    ang_test_reward_actual_temp=ang0;
                    ang_test_reward_ontrack_temp=ang_ontrack;
                    time_test_reward_temp=time0;
                end
                time_stay_reward=time_stay_reward+sample_time;
            else % sign_getreward==1
                if abs(ang_ontrack-ang_ontrack_end) <= ang_nomving && abs(ang_ontrack-ang_test_reward_ontrack_temp) <= ang_nomving
                    % if the rat is still staying here
                    time_stay_reward=time_stay_reward+sample_time;
                else
                    sign_getreward=0;
                    time_stay_reward=0; % reset it to 0
                end
            end
        end
    elseif sign_enter==1 && sign_leave==0 && (ang_ontrack > Ang_StartZone_arrive_ontrack - ang_startzone_border || ang_ontrack < Ang_StartZone_depart_ontrack + ang_startzone_border)
        % the rat is back into starting-zone and he doesn't stop at all around the reward location
        n_reward_lap=n_reward_lap+1;
        ang_test_reward_actual=NaN;
        ang_test_reward_ontrack=NaN;
        time_test_reward=NaN;
        sign_correct=0;
        sign_enter=0;
        sign_leave=0;
    end
    locationArray_end=[x_loc,y_loc];
    ang_ontrack_end=ang_ontrack;
end
fprintf('TestLoc: %s \r',num2str(ang_test_reward_ontrack));

% as soon as the rat is back into starting-zone, this test trial ends
while ang_ontrack > Ang_StartZone_depart_ontrack && ang_ontrack < Ang_StartZone_arrive_ontrack
    pause(sample_time);
    [~, timeStampArray, locationArray, ~, VTRecsReturned, ~] = NlxGetNewVTData(vt_acq_ent);
    time0=timeStampArray(end);
    x_loc = locationArray(end - 1);
    y_loc = locationArray(end);
    if ~x_loc && ~y_loc %% if x_loc=0, y_loc=1, means it's a missing data from the video
        if isempty(locationArray_end)
            continue
        end
        x_loc=locationArray_end(1);
        y_loc=locationArray_end(2);
    end
    x0 = double(x_loc).*scale_x-x_center;
    y0 = double((-1)*y_loc).*scale_y-y_center;
    
    if x0>=0
        if x0==0 x0=0.01; end
        ang0=mod(atan(-y0/x0),2*pi);
    else
        ang0=mod(atan(-y0/x0)+pi,2*pi);
    end
    if strcmp(rot,'Counterclockwise');
        ang0=2*pi-ang0;
    end
    ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
    locationArray_end=[x_loc,y_loc];
end

ts_test_stop=time0;
sound(beep0,fs_sin);
