%%%% AnalyzeRatData
%%%%
%%%% 9.290 Spring 2004 Midterm Project
%%%% Code by Mathew Tantama
%%%% 
%%%% Analyzes rat data contained in a subdirectory named 'rat'.
%%%% Provides stay duration histograms, cume-cume plots, and fits data
%%%% to a matching-maximizing model

%% Begin Code

pause on;

% Records of experimental pA, pB, lambdaA, lambdaB leaving and reward rates
rates=struct('dataset','','block1',zeros(4,1),'block2',zeros(4,1));

% Retrieve filenamse from a rat data subdirectory /rat
files=dir('rat');

% Initialize Rich Side data matrix for holding average dwell time 1/pA,
% reward ratio Q=pA/pB>=1, and reward rate lambdaA for each block of each
% data set
exprichdata=zeros(2*length(files),3);
theorrewratios=zeros(2*length(files),1);

% Loop through all data
for k=3:length(files),
    currentfile=files(k).name;
    load(['rat/',currentfile]);

    % Initialize some global variables
    maxtime=max(etimes);
    blockend=max(find(etimes<=blockends(1)));
    startiter=1;
    maxiter=blockend;
    ratio1=rewFracs(1)/(1-rewFracs(1));
    ratio2=rewFracs(2)/(1-rewFracs(2));
    leverA=0;
    leverB=0;
    tstart=0;
    tend=0;
    staydur=0;
    dwellA=[];
    dwellB=[];
    cumstayA=0;
    cumstayB=0;
    cumstay=zeros(1,2);
    cumrewA=0;
    cumrewB=0;
    cumrew=zeros(1,2);
    lasttime=0;
    thistimeperiod=0;
    
    pA1=0;          % Experimental pA and pB leaving rates
    pA2=0;
    pB1=0;
    pB2=0;
    lambdaA1=0;     % Experimental lambdaA and lambdaB reward rates
    lambdaA2=0;
    lambdaB1=0;
    lambdaB2=0;
    

    % Block 1
    for i=startiter:maxiter,
        thistimeperiod=etimes(i+1)-etimes(i);
        
        % Does the rat start pressing on A?
        if and(leverA==0, events(i,1)==1)            
            leverA=1;
            leverB=0;
            tstart=etimes(i);
            
            % Did the rat receive a reward from switching to A?
            if events(i,3)==1
                lambdaA1=lambdaA1-1;    % Don't count an immediate reward from switching in rate
            end;
            
        end;
        
        % Does the rat start pressing on B?
        if and(leverB==0, events(i,2)==1)
            leverA=0;
            leverB=1;
            tstart=etimes(i);
            
            % Did the rat receive a reward from switching to B?
            if events(i,4)==1
                lambdaB1=lambdaB1-1;    % Don't count an immediate reward from switching in rate
            end;
            
        end;
        
        % Did the rat release lever A?
        if and(leverA==1, events(i,1)==-1)
            pA1=pA1+1;
            leverA=0;
            tend=etimes(i);     % The end of a stay on A
            staydur=tend-tstart;
            dwellA=[dwellA;staydur];
        end;
        
        % Did the rat release lever B?
        if and(leverB==1, events(i,2)==-1)
            pB1=pB1+1;
            leverB=0;
            tend=etimes(i);     % The end of a stay on B
            staydur=tend-tstart;
            dwellB=[dwellB;staydur];
        end;
               
        % Update cumulative stay time on A
        if (leverA==1)
            cumstayA=cumstayA+thistimeperiod;
        end;
        
        % Update cumulative stay time on B
        if (leverB==1)
            cumstayB=cumstayB+thistimeperiod;
        end;
        
        % Update records
        cumstay(i,:)=[cumstayA cumstayB];
        cumrewA=cumrewA+events(i,3);
        cumrewB=cumrewB+events(i,4);
        cumrew(i,:)=[cumrewA cumrewB];
        lasttime=etimes(i);
        
    end; % for i Block 1
           
    % Plot histograms for Block 1
    figure(1);
    subplot(2,3,2);
    hist(dwellA, 50);
    title('Block 1 Lever A');
    xlabel('Stay Time');
    ylabel('# Stays');
    subplot(2,3,3);
    hist(dwellB,50);
    title('Block 1 Lever B');
    xlabel('Stay Time');
    ylabel('# Stays');
    
    % Reset for Block 2
    startiter=blockend;
    maxiter=length(etimes);
    thistimeperiod=0;
    
    % Block 2
    for i=startiter:maxiter,
        if (i~=maxiter)
            thistimeperiod=etimes(i+1)-etimes(i);
        else
            thistimeperiod=0;
        end;
        
        % Does the rat start pressing on A?
        if and(leverA==0, events(i,1)==1)
            leverA=1;
            leverB=0;
            tstart=etimes(i);
            
            % Did the rat receive a reward from switching to A?
            if events(i,3)==1
                lambdaA2=lambdaA2-1;    % Don't count an immediate reward from switching in rate
            end;
            
        end;
        
        % Does the rat start pressing on B?
        if and(leverB==0, events(i,2)==1)
            leverA=0;
            leverB=1;
            tstart=etimes(i);
                      
            % Did the rat receive a reward from switching to B?
            if events(i,4)==1
                lambdaB2=lambdaB2-1;    % Don't count an immediate reward from switching in rate
            end;
            
        end;
        
        % Did the rat release lever A?
        if and(leverA==1, events(i,1)~=1)
            pA2=pA2+1;
            leverA=0;
            tend=etimes(i);     % This might be the end of a stay on A
            staydur=tend-tstart;
            dwellA=[dwellA;staydur];
        end;
        
        % Did the rat release lever B?
        if and(leverB==1, events(i,2)~=1)
            pB2=pB2+1;
            leverB=0;
            tend=etimes(i);     % This might be the end of a stay on B
            staydur=tend-tstart;
            dwellB=[dwellB;staydur];
        end;
                
        % Update cumulative stay time on A
        if (leverA==1)
            cumstayA=cumstayA+thistimeperiod;
        end;
        
        % Update cumulative stay time on B
        if (leverB==1)
            cumstayB=cumstayB+thistimeperiod;
        end;
        
        % Update records
        cumstay(i,:)=[cumstayA cumstayB];
        cumrewA=cumrewA+events(i,3);
        cumrewB=cumrewB+events(i,4);
        cumrew(i,:)=[cumrewA cumrewB];
        lasttime=etimes(i);
        
    end; % for i Block 2
    
    % Calculate experimental leaving and reward rates
    pA1=pA1/cumstay(startiter,1);
    pB1=pB1/cumstay(startiter,2);
    pA2=pA2/(cumstay(maxiter,1)-cumstay(startiter,1));
    pB2=pB2/(cumstay(maxiter,2)-cumstay(startiter,2));
    sA1=(-lambdaA1)/cumstay(startiter,1);
    sB1=(-lambdaB1)/cumstay(startiter,2);
    sA2=(-lambdaA2)/(cumstay(maxiter,1)-cumstay(startiter,1));
    sB2=(-lambdaB2)/(cumstay(maxiter,2)-cumstay(startiter,2));
    lambdaA1=(lambdaA1+cumrew(startiter,1))/cumstay(startiter,1);
    lambdaB1=(lambdaB1+cumrew(startiter,2))/cumstay(startiter,2);
    lambdaA2=(lambdaA2+cumrew(maxiter,1)-cumrew(startiter,1))/(cumstay(maxiter,1)-cumstay(startiter,1));
    lambdaB2=(lambdaB2+cumrew(maxiter,2)-cumrew(startiter,2))/(cumstay(maxiter,2)-cumstay(startiter,2));
    rates(k-2).dataset=currentfile(1:5);
    rates(k-2).block1=[pA1;pB1;lambdaA1;lambdaA2];
    rates(k-2).block2=[pA2;pB2;lambdaA2;lambdaB2];
    disp(currentfile(1:5));
    disp(['Block1: Set Rew. Ratio = ',num2str(ratio1),', Obs. Rew. Ratio = ',num2str(lambdaA1/lambdaB1)]);
    disp(['Block2: Set Rew. Ratio = ',num2str(ratio2),', Obs. Rew. Ratio = ',num2str(lambdaA2/lambdaB2)]);
    disp(['Obs. Leaving Ratio pA/pB: Block 1 = ',num2str(pA1/pB1),', Block2 = ',num2str(pA2/pB2)]);
    disp(['s/lambda - 1A=',num2str(sA1/lambdaA1),' 1B=',num2str(sB1/lambdaB1),' 2A=',num2str(sA2/lambdaA2),' 2B=',num2str(sB2/lambdaB2)]);
    
    % Update richdata
    if (lambdaA1>=lambdaB1),
        exprichdata(2*k,:)=[1/pA1, lambdaA1/lambdaB1, lambdaA1];
    else
        exprichdata(2*k,:)=[1/pB1, lambdaB1/lambdaA1, lambdaB1];
    end;
    theorrewratios(2*k)=max(ratio1, 1/ratio1);
    if (lambdaA2>=lambdaB2),
        exprichdata(2*k+1,:)=[1/pA2, lambdaA2/lambdaB2, lambdaA2];
    else
        exprichdata(2*k+1,:)=[1/pB2, lambdaB2/lambdaA2, lambdaB2];
    end;
    theorrewratios(2*k+1)=max(ratio2, 1/ratio2);
    
    % Plot histograms for Block 2
    figure(1);
    subplot(2,3,5);
    hist(dwellA, 50);
    title('Block 2 Lever A');
    xlabel('Stay Time');
    ylabel('# Stays');
    subplot(2,3,6);
    hist(dwellB,50);
    title('Block 2 Lever B');
    xlabel('Stay Time');
    ylabel('# Stays');
    
    % Plot cum-cum for stays and rewards
    figure(1);
    subplot(2,3,1);
    plot(cumstay(:,1),cumstay(:,2),'b');
    title([currentfile(1:5),' Cume Stays']);
    xlabel('Cume Stay on A');
    ylabel('Cume Stay on B');
    hold on;
    plot(cumstay(1:blockend-1,1),cumstay(1:blockend-1,1)/ratio1,'r--');
    plot(cumstay(blockend:maxiter,1),cumstay(blockend,2)+(cumstay(blockend:maxiter,1)-cumstay(blockend,1))/ratio2,'r--');
    axis tight;
    hold off;
    
    figure(1);
    subplot(2,3,4);
    plot(cumrew(:,1),cumrew(:,2),'b');
    title([currentfile(1:5),' Cume Rewards']);
    xlabel('Cume Reward on A');
    ylabel('Cume Reward on B');
    hold on;
    plot(cumrew(1:blockend-1,1),cumrew(1:blockend-1,1)/ratio1,'r--');
    plot(cumrew(blockend:maxiter,1),cumrew(blockend,2)+(cumrew(blockend:maxiter,1)-cumrew(blockend,1))/ratio2,'r--');
    axis tight;
    hold off;
    
    w=waitforbuttonpress;

end; % for k

%% Fit the data to a matching-maximizing model

% Remove non-finite or zero elements
bump=length(theorrewratios(find(theorrewratios==0)));
theorrewratios=theorrewratios(bump+1:length(theorrewratios));
uniqueratios=unique(theorrewratios);

% Separate groups of data with different reward ratios
for i=1:(length(uniqueratios)-1),
    if (i==(length(uniqueratios)-1)),
        setsofinterest=sort([find(theorrewratios==uniqueratios(i));find(theorrewratios==uniqueratios(i+1))]);
        disp('yep');
    else
        setsofinterest=find(theorrewratios==uniqueratios(i));
    end; % if else

% Display rich data to see if 1/Ta ~ (1/f(c))*1/Q*lambdaA
figure(i+1);
subplot(1,2,1);
t=1./exprichdata(setsofinterest,1);
bump=length(t(find(t==inf)));               % Get rid of initial 0/infinite values
t=t(bump+1:length(t));
y=(exprichdata(setsofinterest,3)./exprichdata(setsofinterest,2));
y=y(bump+1:length(y));
xl=t;                                       % Fit to a line
al=xl\y;                                    % Solve for coefficients
T=min(t(find(t~=inf)))+(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):max(t(find(t~=inf)));
Yl=T'*al;
plot(y,t,'bo',Yl,T,'r-');
title(['Q = ',num2str(uniqueratios(i))]);
ylabel('1/T');
xlabel('{\lambda}/Q');
subplot(1,2,2);
loglog(y,t,'bo',Yl,T,'r-');
title('Log Log Plot');

subplot(1,2,1);
hold on;
xe=[ones(size(t)),exp(-t)];                 % Fit to an exponential
ae=xe\y;                                    % Solve for coefficients
T=min(t(find(t~=inf)))+(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):max(t(find(t~=inf)));
Ye=cat(1,ones(size(T)),exp(-T))'*ae;
plot(Ye,T,'g-');
subplot(1,2,2);
hold on;
loglog(Ye,T,'g-');

end; % for i

% Display rich data to see if 1/Ta ~ (1/f(c))*1/Q*lambdaA
figure(length(uniqueratios)+1);
subplot(1,2,1);
t=1./exprichdata(:,1);
bump=length(t(find(t==inf)));                   % Get rid of initial 0/infinite values
t=t(bump+1:length(t));
y=(exprichdata(:,3)./exprichdata(:,2));
y=y(bump+1:length(y));
xl=t;                                           % Fit to a line
al=xl\y;                                        % Solve for coefficients
T=min(t(find(t~=inf)))+(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):max(t(find(t~=inf)));
Yl=T'*al;
plot(y,t,'bo',Yl,T,'r-');
title('All Data');
ylabel('1/T');
xlabel('{\lambda}/Q');
subplot(1,2,2);
loglog(y,t,'bo',Yl,T,'r-');
title('Log Log Plot');

subplot(1,2,1);
hold on;
xe=[ones(size(t)),exp(-t)];                     % Fit to an exponential
ae=xe\y;                                        % Solve for coefficients
T=min(t(find(t~=inf)))+(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):(max(t(find(t~=inf)))-min(t(find(t~=inf))))/length(t):max(t(find(t~=inf)));
Ye=cat(1,ones(size(T)),exp(-T))'*ae;
plot(Ye,T,'g-');
subplot(1,2,2);
hold on;
loglog(Ye,T,'g-');

%% END AnalyzeRatData