%% v17: 
% edited on 04/30/2017
% this is the 6th version for new recording system in new recording room
% basically same code as in v16
% but only change the paradigm in sample trials to the same as that in test trials
% i.e. report the incorrect stop location when the rat stops at a later location
% set incorrect location +-3 locations around the correct location

% The experiment paradihm is setting as:
% 
% 1. All reward locations are marked by white door bumpers (0.7inch) at the inner edge of
%    the track.
% 2. In sample trials, a striking marker in a different color (orange) was
%    at reward location besides the corresponding white bumper; but the orange
%    sticker will be removed in the following test trials.
% 3. Sample1(marker)-test1(no marker)-Sample2(marker)-test2(no marker)-
%    Sample3(marker)-test3(no marker)
%    Sample1-20s-test1-40s-sample2-20s-test2........
% 4. 3 mins rest in the animal box in between sessions, and before the
%    first session and after the last session
% 5. A correct/incorrect tone goes off when he stops at correct/incorrect spot.
% 6. In sample trial, rats can stop at other locations besides reward
%    locaiton, and food will be delivered as long as they stop at reward
%    location. However, in test trials, rats are not allowed to stop and
%    turn to the center at the other locations. That means, if rats stop at
%    a wrong location and turn to the center, a incorrect tone will go off
%    and then no food reward will be delivered.
% 7. However, if the rat turned back, criteria are different. If the rat
%    turned back without already passing by the reward location, he still
%    has chance to get the reward; otherwise, he can't get the reward
%    anymore even if he ran back and force around the reward location.
% 8. At the end of each lap, the rat needs to be back to the animal box.
%    This lap (sample or test) will end if he's angle is around 0 degrees,
%    which is the center of the animal box.

% set the missing data [0,0] as the same as the last visited data
locationArray_end=[0,0];
fig = figure('Units','normalized','Position',[0 0 0.4 0.5]);

%% initialize variables for this session
ts_sample_start=nan(n_sample_test,1);
ts_sample_reward=nan(n_sample_test,1);
ts_sample_end=nan(n_sample_test,1);
ts_test_start=nan(n_sample_test,1);
ts_test_reward=nan(n_sample_test,1);
ts_test_end=nan(n_sample_test,1);
sign_correct_sample=nan(n_sample_test,1);
sign_correct_test=nan(n_sample_test,1);
ind_rewardloc_sample=nan(n_sample_test,1);
ind_rewardloc_test=nan(n_sample_test,1);
ang_sample_reward_ontrack=nan(n_sample_test,1);
ang_test_reward_ontrack=nan(n_sample_test,1);

%% Samples and tests
for n_trial = 1:n_sample_test
    %% 40s before each sample
    figure(fig);
    cla
    hold on
    text(0.5,0.8,strcat('Location .', num2str(ind_rewardloc)),'color','b','fontsize',70,'HorizontalAlignment','center');
    text(0.5,0.4,strcat('Sample . ', num2str(n_trial)),'color','b','fontsize',70,'HorizontalAlignment','center');
    xlim([0,1]);ylim([0,1])
    hold off
    axis off
    
    h=waitbar(0, strcat('Sample:',num2str(n_trial)));
    for min0=1:time_test_sample
        pause(1);
        waitbar(min0/time_test_sample);
    end
    close(h)
    
    %% Sample trials
%     sound(sound_beep_start,sound_beep_fs);  % make a sound when a sample trial starts
    sound(beep0,fs_sin);
    
    pause(sample_time);
    get_current_loc_v1  % get the current location and angle from on-line video  
    ts_sample_start(n_trial,1)=time0;
    sign_enter=0;  % 1 = the rat was in reward zone
    sign_leave=0; % 1 = has stopped at one location ; 0 = not stopped at any location
    sign_correct0=NaN;
    ind_loc_max=0;
    time_stay_reward=0;
    ang_ontrack=0;
    ang_ontrack_end=0;
    locationArray_end=[0,0];
    while 1
        pause(sample_time);
        get_current_loc_v1  % get the current location and angle from on-line video
        ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
        vel_ontrack=(ang_ontrack-ang_ontrack_end)/sample_time;
        if (sign_leave==1 || ind_loc_max > N_loc/2)...
                && (ang_ontrack>Ang_StartZone_arrive_ontrack || ang_ontrack<Ang_StartZone_depart_ontrack)
            % if the rat is back entering the animal box, sample tests end
            % sound(sound_beep_start,sound_beep_fs);
            % use this "if" to make sure the rat has finished this
            % lap (pass the middle of track), but not turned around at the
            % starting position
            break
        end
        if sign_leave==0
            % this rat hasn't stopped at either locations yet
            
            ind_loc_temp=find(ang_ontrack>Ang_RewardLoc_ontrack_zone_narrow(:,1) & ang_ontrack<Ang_RewardLoc_ontrack_zone_narrow(:,2));
            if ~isempty(ind_loc_temp)
                % rat is walking in location zone, and reaching this loc
                sign_enter=1;
                if ind_loc_temp > ind_loc_max & ind_loc_temp <= ind_loc_max+N_loc/4
                    % exclude the situation when the rat run backwards from the
                    % box to the last reward location.
                    ind_loc_max=ind_loc_temp;
                    time_stay_reward=0;
                elseif ind_loc_temp >= ind_loc_max-1 && dist0 <= R_inner_RewardLoc(ind_loc_temp)
                    time_stay_reward=time_stay_reward+sample_time;
                else
                    time_stay_reward=0;
                    sign_enter=0;
                end
            end
            
            if time_stay_reward >= time_stay_sample
                % if the rat has stopped at reward location for "time_stay_sample"
                if ind_loc_max==ind_rewardloc
                    % correct reward location
                    sound(tone_correct,fs_tone_correct);
                    sign_correct0=1;
                elseif ind_loc_max~=ind_rewardloc
                    % the rat passes by the location
                    if ind_loc_max > ind_rewardloc
                        % sound(tone_incorrect,fs_tone_incorrect);
                        sign_correct0=0;
                    else
                        % not report the incorrect reward location in Samples
                        time_stay_reward=0;
                        ang_ontrack_end=ang_ontrack;
                        continue
                    end
                end
                ind_rewardloc_sample(n_trial,1)=ind_loc_max;
                ang_sample_reward_ontrack(n_trial,1)=Ang_RewardLoc_ontrack(ind_loc_max);
                ts_sample_reward(n_trial,1)=time0;
                sign_correct_sample(n_trial,1)=sign_correct0;
                sign_enter=0;
                sign_leave=1;
            end
        end
        ang_ontrack_end=ang_ontrack;
    end
    if sign_leave==0  % if didn't stopped at any location
        ind_rewardloc_sample(n_trial,1)=nan;
        ang_sample_reward_ontrack(n_trial,1)=nan;
        ts_sample_reward(n_trial,1)=nan;
        sign_correct_sample(n_trial,1)=nan;
    end
    ts_sample_end(n_trial,1)=time0;
    
    %% 20s between sample and test
    figure(fig);
    cla
    hold on
    text(0.5,0.8,strcat('Location .', num2str(ind_rewardloc)),'color','b','fontsize',70,'HorizontalAlignment','center');
    text(0.5,0.4,strcat('Test . ', num2str(n_trial)),'color','b','fontsize',70,'HorizontalAlignment','center');
    xlim([0,1]);ylim([0,1])
    hold off
    axis off
    
    h=waitbar(0, strcat('Test:',num2str(n_trial)));
    for min0=1:time_sample_test
        pause(1);
        waitbar(min0/time_sample_test);
    end
    close(h)
    
    %% Test trials
%     sound(sound_beep_start,sound_beep_fs);  % make a sound when a test trial starts
    sound(beep0,fs_sin);
    
    pause(sample_time);
    get_current_loc_v1  % get the current location and angle from on-line video  
    ts_test_start(n_trial,1)=time0;
    sign_enter=0;  % 1 = the rat was in reward zone
    sign_leave=0; % 1 = has stopped at one location ; 0 = not stopped at any location
    sign_correct0=NaN;
    ind_loc_max=0;
    time_stay_reward=0;
    ang_ontrack=0;
    ang_ontrack_end=0;
    locationArray_end=[0,0];
    while 1
        pause(sample_time);
        get_current_loc_v1  % get the current location and angle from on-line video
        ang_ontrack=mod(ang0-Ang_StartZone_center,2*pi);
        vel_ontrack=(ang_ontrack-ang_ontrack_end)/sample_time;
        if (sign_leave==1 || ind_loc_max > N_loc/2)...
                && (ang_ontrack>Ang_StartZone_arrive_ontrack || ang_ontrack<Ang_StartZone_depart_ontrack)
            % if the rat is back entering the animal box, sample tests end
            % sound(sound_beep_start,sound_beep_fs);
            break
        end
        if sign_leave==0
            % this rat hasn't stopped at either locations yet
            
            ind_loc_temp=find(ang_ontrack>Ang_RewardLoc_ontrack_zone_narrow(:,1) & ang_ontrack<Ang_RewardLoc_ontrack_zone_narrow(:,2));
            if ~isempty(ind_loc_temp)
                % rat is walking in location zone, and reaching this loc
                sign_enter=1;
                if ind_loc_temp > ind_loc_max & ind_loc_temp <= ind_loc_max+N_loc/4
                    % exclude the situation when the rat run backwards from the
                    % box to the last reward location.
                    ind_loc_max=ind_loc_temp;
                    time_stay_reward=0;
                elseif ind_loc_temp >= ind_loc_max-1 && dist0 <= R_inner_RewardLoc(ind_loc_temp)
                    time_stay_reward=time_stay_reward+sample_time;
                else
                    time_stay_reward=0;
                    sign_enter=0;
                end
            end
            
            if time_stay_reward >= time_stay_test
                % if the rat has stopped at reward location for "time_stay_test"
                if ind_loc_max==ind_rewardloc
                    % correct reward location
                    sound(tone_correct,fs_tone_correct);
                    sign_correct0=1;
                elseif ind_loc_max~=ind_rewardloc
                    % first time get to incorrect reward location
                    if abs(ind_loc_max-ind_rewardloc) <= 3
                        sound(tone_incorrect,fs_tone_incorrect);
                        sign_correct0=0;
                    else
                        % if the rat stopped at a location very far away from the reward location,
                        % we don't mark here as incorrect location
                        time_stay_reward=0;
                        ang_ontrack_end=ang_ontrack;
                        continue
                    end
                end
                ind_rewardloc_test(n_trial,1)=ind_loc_max;
                ang_test_reward_ontrack(n_trial,1)=Ang_RewardLoc_ontrack(ind_loc_max);
                ts_test_reward(n_trial,1)=time0;
                sign_correct_test(n_trial,1)=sign_correct0;
                sign_enter=0;
                sign_leave=1;
            end
        end
        ang_ontrack_end=ang_ontrack;
    end
    if sign_leave==0  % if didn't stopped at any location
        ind_rewardloc_test(n_trial,1)=nan;
        ang_test_reward_ontrack(n_trial,1)=nan;
        ts_test_reward(n_trial,1)=nan;
        sign_correct_test(n_trial,1)=nan;
    end
    ts_test_end(n_trial,1)=time0;
    
end
close(fig)