function p = runKalman(se,sh,s,R,varargin)
%simulatate kalman filter
%se: observation variance
%sh: innovation variance
%s: stimulus sequence, trial x cue
%R: outcome sequence
%varargin suppresses plots

global Sigma

n = size(s,1); %number of trials
k = size(s,2); %number of cues
mu = zeros(n+1,k); %mean of prior on weights, trial x cue
Sigma = zeros(k,k,n+1); %variance matrix for prior, trial x trial x cue
Sigma(:,:,1) = sh*eye(k); %initialize prior variance (somewhat arbitrary)
p = zeros(n,1); %sequence of predictions

%run model
for t=1:n %loop through trials
    p(t) = s(t,:)*mu(t,:)'; %expected outcome
    mu(t+1,:) = mu(t,:) + s(t,:)*(Sigma(:,:,t)^(-1)*se+s(t,:)'*s(t,:))^(-1)*(R(t)-p(t)); %update mean
    Sigma(:,:,t+1) = se*Sigma(:,:,t)*(se*eye(k)+s(t,:)'*s(t,:)*Sigma(:,:,t))^(-1)+sh*eye(k); %udpate variance
end

%plots
if nargin<5
    figure(2),clf
    co = get(gca,'ColorOrder');
    for c=1:k %loop through cues
        subplot(k+7,1,c)
        plot(s(:,c),'.','Color',co(c,:)); %plot values of cue number c
        set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.5 1.5],'YTick',0:1,'YTickLabel','','FontSize',16); ylabel('Cue');
    end
    subplot(k+7,1,k+1)
    plot(R,'k.'); %plot outcome sequence
    set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.5 1.5],'YTick',0:1,'YTickLabel','','FontSize',16); ylabel('Outcome');
    subplot(k+7,1,k+2:k+3)
    plot(mu) %plot weights
    set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.05 1.05],'FontSize',16); ylabel('Means');
    subplot(k+7,1,k+4:k+5);hold on;box
    for c=1:k, for d=1:c-1 %loop through cue pairs
        plot(squeeze(Sigma(c,d,:)./sqrt(Sigma(c,c,:))./sqrt(Sigma(d,d,:)))); %plot cue correlations
    end,end
    set(gca,'XLim',[0 n+1],'YLim',[-1.01,1.01],'FontSize',16); xlabel('Trial'); ylabel('Correlations');
    subplot(k+7,1,k+6:k+7)
    plot(p,'k') %plot predictions
    set(gca,'XLim',[0 n+1],'YLim',[-.05 1.05],'FontSize',16); xlabel('Trial'); ylabel('Predictions');
end