% In this tutorial, we will predict data from Koerding and Wolpert 2004, then
% fit the actual data.

% In the experiments, they have subjects move a cursor to a target at 0cm
% with a random lateral shift between the hand position and the cursor. The
% cursor's position is rendered with 2 levels of blur or occluded, increasing 
% the uncertainty about target position, predicting an increased reliance on the prior.

% On each trial, there is a true shift, given by the cursor position,
% "xtrue"
%  subjects form an estimate of xtrue from data and prior knowledge that I 
%  call hp below
%  
% There is also a noisy estimate of the cursor position, 
%  "xsensed"
% 
%  assuming that reach endpoints reflect the best estimate of cp, we can
%  compare predictions and reach endpoints.

% cp = argmax P(xtrue|xsensed), which for Gaussians is the mean of P(xtrue|xsensed)
%  By Bayes P(xtrue|xsensed) = P(xsensed| xtrue)P(xtrue)/P(xsensed)
% where P(xsensed) = Int_{xtrue} P(xsensed| xtrue)P(xtrue) dxtrue

% We need to form the prior P(xtrue) and the likelihood P(xsensed| xtrue)

% PRIOR
x = (-.03:.0001:.03)';  % in meters -2cm through 2cm
mu_prior = .01;  sd_prior = 0.5/100;
z_p = (x-mu_prior).*(x-mu_prior)/sd_prior^2;  % general quadratic form

p_p = exp(-0.5*z_p); % compute unnormalized density
p_p = p_p/sum(p_p);  % normalize into a probability over the discretized x

% Likelihood for a "small blur"

xtrue = [-0.015 -0.01 -0.005 0 0.005 0.01 0.015];
sd_sensed = 0.1/100;

for k=1:length(xtrue)
    z_sensed = (x-xtrue(k)).*(x-xtrue(k))/sd_sensed^2;  % general quadratic form
    p_s(:,k) = exp(-0.5*z_sensed); % compute unnormalized density
    p_s(:,k) = p_s(:,k)/sum(p_s(:,k));
    post(:,k) = p_s(:,k).*p_p; 
    post(:,k) = post(:,k)/sum(post(:,k));
    cp_pred(k) = sum(post(:,k).*x);
end

% Likelihood for a "larger blur"

xtrue = [-0.015 -0.01 -0.005 0 0.005 0.01 0.015];
sd_sensed_l = 1/100;

for k=1:length(xtrue)
    z_sensed = (x-xtrue(k)).*(x-xtrue(k))/sd_sensed_l^2;  % general quadratic form
    p_l(:,k) = exp(-0.5*z_sensed); % compute unnormalized density
    p_l(:,k) = p_l(:,k)/sum(p_l(:,k));
    post_l(:,k) = p_l(:,k).*p_p; 
    post_l(:,k) = post_l(:,k)/sum(post_l(:,k));
    cp_pred_l(k) = sum(post_l(:,k).*x);
end

% now we plot the predictions, generated by formula, and by computing the
% probabilities

weights_s = 1./[sd_prior^2 sd_sensed^2]; 
weights_s = weights_s/sum(weights_s);
mean_pred_s = mu_prior*weights_s(1) + xtrue*weights_s(2);

weights_l = 1./[sd_prior^2 sd_sensed_l^2]; 
weights_l = weights_l/sum(weights_l);
mean_pred_l = mu_prior*weights_l(1) + xtrue*weights_l(2);

subplot(2,3,1);
plot(x,p_p,x,p_s)
subplot(2,3,2);
plot(x,post)
subplot(2,3,3)
plot(xtrue,cp_pred,xtrue,mean_pred_s+0.00001)
axis([-0.02 0.02 -0.02 0.02])

subplot(2,3,4);
plot(x,p_p,x,p_l)
subplot(2,3,5);
plot(x,post_l)
subplot(2,3,6)
plot(xtrue,cp_pred_l,xtrue,mean_pred_l+0.00001)
axis([-0.02 0.02 -0.02 0.02])

% SIMULATE THE OBSERVER FOR ONE CONDITION
% we will assume the observer has about a thousand trials of feedback to
% estimate the prior.
feedback = randn(1,1010)*sd_prior+mu_prior;
mu_prior_est = mean(feedback);
sd_prior_est = std(feedback);
motor_sd = 0.2/100;
sim_sd = .15/100;

weights_sim = 1./[sd_prior_est^2 sim_sd^2];
weights_sim = weights_sim/sum(weights_sim);

for j=1:300,  % number of simulated trials
   cursor(j) = randn*sd_prior+mu_prior;
   sensedpos(j) = cursor(j)+randn*sim_sd;
   hp_sim(j) = weights_sim(1)*mu_prior_est+weights_sim(2)*sensedpos(j) + randn*motor_sd;
end
figure(2);
plot(cursor,hp_sim,'.')

%  Now we analyze this data.  

% Note that the simulated data has the distribution P(hp_sim,cursor)
% P(hp_sim|cursor) = 
% Normal(mu_sim, sd_sim) 
%  mu_sim = mu_prior_est*weights_sim(1) + cursor *weights_sim(2), 
% sd_sim^2 = motor_sd^2 + (sd_prior_est^2*sim_sd^2)/(sd_prior_est^2+sim_sd^2) )
%  NOW P(hp_sim,cursor) =  P(hp_sim|cursor)P(cursor) 
%  This is a joint Gaussian: = Normal(mu_sim, sd_sim^2)*N(mu_prior,sd_prior^2)
% we can estimate a mean and covariance for the data and back-infer our
% parameters using the conditional gaussian formula:
%  hp|cursor ? Normal( mu(1) + C/B*(cursor ? mu(2)), A ? CB?1CT );
mu_sim = mean([hp_sim(:) cursor(:)]);
C_sim = cov([hp_sim(:) cursor(:)]);
A = C_sim(1,1); B = C_sim(2,2); C = C_sim(1,2);

cursorS = sort(cursor);
mu_est = mu_sim(1) + C/B*(cursorS - mu_sim(2));
plot(cursorS,mu_est,cursorS,weights_sim(1)*mu_prior_est+weights_sim(2)*cursorS);
sig_est = sqrt(A - C*inv(B)*C);
% Now let's try to extract the parameters from sig_est and mu_est
% sig_est = motor_sd^2 +
% (sd_prior_est^2*sim_sd^2)/(sd_prior_est^2+sim_sd^2) );
% mu_est(1) = weights_sim(1)*mu_prior_est+weights_sim(2)*cursor(j);
% collecting terms, we have weights_sim(2) = C/B 
%  and weights_sim(1)*mu_prior_est = mu_sim(1) -C/B*mu_sim(2);
% IS IT POSSIBLE TO SOLVE FOR ALL THE PARAMETERS?
%  WHAT OTHER DATA MIGHT NEED TO BE COLLECTED?



% REAL DATA ANALYSIS
% First we load the data
load Kording_2004.mat;

% There are 10 subjects.  Let's analyze one, say 
subjnum = 1;
% we grab the condition indices
Task.Conditions
% we will analyze 'occluded'    'small blur'    'large blur'
%  the first 1010 trials are training trials.  The last 1000 are test.
cond = [Subject(subjnum).Trial(:).Condition];
Jinf = find(cond(1011:end)==4);
Jsmb = find(cond(1011:end)==2);
Jlgb = find(cond(1011:end)==3);

for j=1:length(Jinf), 
    if ~isempty(Subject(subjnum).Trial(Jinf(j)).CursorPos{end}), 
        cpinf(j) = Subject(subjnum).Trial(Jinf(j)).CursorPos{end}(1,1); 
    end; 
    hpinf(j) = Subject(subjnum).Trial(Jinf(j)).HandPos(end,1);
end
for j=1:length(Jsmb), 
    if ~isempty(Subject(subjnum).Trial(Jsmb(j)).CursorPos{end}), 
        cpsmb(j) = Subject(subjnum).Trial(Jsmb(j)).CursorPos{end}(1,1); 
    end; 
    hpsmb(j) = Subject(subjnum).Trial(Jsmb(j)).HandPos(end,1);
end
for j=1:length(Jlgb),
    if ~isempty(Subject(subjnum).Trial(Jlgb(j)).CursorPos{end}), 
        cplgb(j) = Subject(subjnum).Trial(Jlgb(j)).CursorPos{end}(1,1); 
    end;
    hplgb(j) = Subject(subjnum).Trial(Jlgb(j)).HandPos(end,1);
end

figure(5);
subplot(1,3,1);
plot(cplgb(cplgb~=0),hplgb(cplgb~=0),'.');
title('sigma_{lg}')
subplot(1,3,2);
plot(cpsmb(cpsmb~=0),hpsmb(cpsmb~=0),'.');
title('sigma_{sm}')
subplot(1,3,3);
plot(cpinf(cpinf~=0),hpinf(cpinf~=0),'.');
title('sigma_{\inf}')
% These are the data - compute means and Cov, and decompose the Cov as
% above.  
% for example:
mu_inf = mean([cpinf(cpinf~=0),hpinf(cpinf~=0)]);
C_inf = cov([cpinf(cpinf~=0),hpinf(cpinf~=0)]);

%DO:  Plot mu_est for each condition and compare to Kording and Wolpert's
%published data - remember this is only one subject.  


% PART II  LEARNING THE PRIOR

% is there evidence of learning in the training phase?
% let's visualize
% first analyze the training data

Jinf = find(cond(1:1010)==4);
Jsmb = find(cond(1:1010)==2);
Jlgb = find(cond(1:1010)==3);

for j=1:length(Jinf), 
    if ~isempty(Subject(subjnum).Trial(Jinf(j)).CursorPos{end}), 
        cpinf(j) = Subject(subjnum).Trial(Jinf(j)).CursorPos{end}(1,1); 
    end; 
    hpinf(j) = Subject(subjnum).Trial(Jinf(j)).HandPos(end,1);
end
for j=1:length(Jsmb), 
    if ~isempty(Subject(subjnum).Trial(Jsmb(j)).CursorPos{end}), 
        cpsmb(j) = Subject(subjnum).Trial(Jsmb(j)).CursorPos{end}(1,1); 
    end; 
    hpsmb(j) = Subject(subjnum).Trial(Jsmb(j)).HandPos(end,1);
end
for j=1:length(Jlgb),
    if ~isempty(Subject(subjnum).Trial(Jlgb(j)).CursorPos{end}), 
        cplgb(j) = Subject(subjnum).Trial(Jlgb(j)).CursorPos{end}(1,1); 
    end;
    hplgb(j) = Subject(subjnum).Trial(Jlgb(j)).HandPos(end,1);
end

 plot3(Jinf,hp(Jinf-1),hp(Jinf),'.')
% doing this plot suggests learning is occurring - why?  
% Tomorrow we will see if we can use temporal integration models to account
% for this effect.











