%explanation figure for MOT paper
rng(12)

%% default Parameters
Vdist=1.75; % viewing distance
L1=0.9; %ne-cd model
L2=0;
Lapse=0;
c=0.0082; % frontal plane noise 
d=0.0202; % depth Noise scaling
NP=2; %only two combinatons for 2 targets
%% parameters of simulation
sigV=0.2; %velocity sigma
sigX=2; %position Sigma

sigV=visAngle2m(Vdist,sigV);
sigX=visAngle2m(Vdist,sigX);

T = 5; %length of simualtion /s
hz=60; %simulation rate /hz
dt=1; %time step (currently unused)
time=T*hz; %total number of time steps
a=0.0; % slope of velocity noise, fitted
b=0.05; %weber fraction velocity noise, from ()
nDim=6; % number of state dimensions
NT=2; %Number of tracked targets
runs=2; % number of monte carlo repeats


%% Initalize variables
%variables to be stored if parfor is used
probtarget=nan(NT,runs,length(sigV)); %probability of target given sample

    %% generate trajectories
    idx=randperm(NT); %randomise measuremet order
    
    L=0.9;
    k=((L+1)*sigV^2)/(2*sigX^2);
    sig=0.5*sqrt(sigV^2*(4*sigX^2-4*L^2*sigX^2-sigV^2+L^2*sigV^2)/sigX^2);
    A=([[1-k,0,0, L,0,0] ;[0,1-k, 0,0,L,0];[0,0, 1-k,0,0,L];[-k,0,0 L,0,0];[0,-k, 0,0,L,0];[0,0, -k,0,0,L]]);
    Q=([[sig^2,0 0, sig^2,0,0] ;[0,sig^2,0 0,sig^2,0];[0,0,sig^2 0,0,sig^2]; [sig^2,0 0, sig^2,0,0] ;[0,sig^2,0 0,sig^2,0];[0,0,sig^2 0,0,sig^2]]);
    
    prepv=zeros(NP,runs);
    logprop=zeros(NP,runs,time);
    EST=zeros(nDim,time,NT,runs);
    sampleNoise=zeros(nDim,nDim,runs);
    estcov=zeros(nDim,nDim,time,NT,runs);
    mnoise=randn(nDim,time,NT,runs);
    proNoise=mvnrnd(zeros(nDim,1),Q,time*NT*runs)';
    proNoise=reshape(proNoise,nDim,time,NT,runs);
    S=zeros(length(A),time,NT,runs);
    m=zeros(length(A),time,NT,runs);
    S(1:nDim/2,1,:,:,:)=normrnd(0,sigX,[nDim/2,1,NT,runs]);
    S((nDim/2)+1:nDim,1,:,:,:)=normrnd(0,sigV,[nDim/2,NT,runs]);
    M=zeros(nDim,nDim,time,NT,runs);
    
    %generate trajectories
        for r =1:runs
            %can probably calc across targerts
            for j=1:NT
                M(:,:,1,j,r)=covM(S(:,1,j,r),a,b,c,d);
                m(:,1,j,r)=S(:,1,j,r)+mnoise(:,1,j,r).*diag(M(:,:,1,j,r));
                
                for i=1:time-1
                    S(:,i+1,j,r)=A*S(:,i,j,r)+proNoise(:,i,j,r);
                    M(:,:,i+1,j,r)=covM(S(:,i+1,j,r),a,b,c,d);
                    m(:,i+1,j,r)=S(:,i+1,j,r)+mnoise(:,i+1,j,r).*diag(M(:,:,i+1,j,r));
                    M(:,:,i+1,j,r)=covM(m(:,i+1,j,r),a,b,c,d); %subject
                    %does not know true variance add bck in 
                end
            end
        end

    
    %% filter parameters
    L=L1;
    k=((L+1)*sigV^2)/(2*sigX^2);
    sig=0.5*sqrt(sigV^2*(4*sigX^2-4*L^2*sigX^2-sigV^2+L^2*sigV^2)/sigX^2);
    L=L2;
    A=([[1-k,0,0, L,0,0] ;[0,1-k, 0,0,L,0];[0,0, 1-k,0,0,L];[-k,0,0 L,0,0];[0,-k, 0,0,L,0];[0,0, -k,0,0,L]]);
    Q=([[sig^2,0 0, sig^2,0,0] ;[0,sig^2,0 0,sig^2,0];[0,0,sig^2 0,0,sig^2]; [sig^2,0 0, sig^2,0,0] ;[0,sig^2,0 0,sig^2,0];[0,0,sig^2 0,0,sig^2]]);
    
    %% intialize variables for filter
    %dimension state,targets,runs,particles
    %define variables needed for kalman filter
    startEST=S(:,1,:,:);
    ESTstate=repmat(squeeze(startEST),1,1,1,NP); %set initial velocity state to zero
    ESTstate=reshape(ESTstate,nDim,1,NT,runs,NP);
    ESTstate=permute(ESTstate,[1,2,3,5,4]);
    EST(:,1,:,:)=ESTstate(:,1,:,1,:);
    
    %covariance of prediciton
    startCov=M(:,:,1,:,:);
    redo=[4,5,6];
    for j=1:length(redo)
        startCov(redo(j),redo(j),1,:,:)=repmat(a,size(startCov(redo(j),redo(j),1,:,:)));
    end
    
    ESTcov=repmat(squeeze(startCov.^2),1,1,1,1,NP); %eliminate velocity set to small value, 0.0001
    ESTcov=permute(ESTcov,[1,2,3,5,4]);
    estcov(:,:,1,:,:)=ESTcov(:,:,:,1,:);
    
    prepv=zeros(NP,runs);
    logprop=zeros(NP,runs,time);
    allgamma=perms(1:NT);
    gammaLen=length(allgamma); %length of gamma vector
    gammaMap=zeros(gammaLen,NT*NT);
    map=0:NT:NT*(NT-1);
    gammaMap(sub2ind(size(gammaMap),(1:gammaLen)'*ones(1,NT),allgamma+ones(gammaLen,1)*map))=1; %compute sum of log likelihood my mapping it
    gammaMap=sparse(gammaMap);
    LH=nan(NT,NP,runs,NT);
    s=zeros(gammaLen,NP,runs);
    i=nan(NP,runs);
    weights=ones(NP,runs)/NP;
    particleProb=ones(NP,runs)/NP;
    
    datassoc=nan(NP,runs,10);
    idx=[1 2 3 4 5 6] ;
    
    for k=2:time%from second time step (first used for intialization
        %prediction step for all targets
        ESTref=ESTstate; ESTcovref=ESTcov;
        [StateEstimate,CovEstimate]=kalPredMD(ESTstate,ESTcov,A,Q);
        StateRef=StateEstimate; StatecovRef=CovEstimate;
        
        
        %%
        for j=1:NT %go through each measurment
            
            LH(:,:,:,j)=kalMeasLik(m(:,k,idx(j),:),StateEstimate,CovEstimate,M(:,:,k,idx(j),:).^2,NT,NP,nDim);
           
        end
        squeeze(ESTcov(:,:,:,2,2));
        
        
        for r=1:runs
            for e=1:NP
                lik=squeeze(LH(:,e,r,:));
                loglik=lik;
                loglik=loglik(:);
                if k==2
                    s(:,e,r)=gammaMap*loglik; %probability for each partile
                else
                    s(:,e,r)=logprop(e,r,k-1)+gammaMap*loglik; %likelihood of each particle (L(i,e)=L(i-1,e)+log(p(data/part))
                    %s(:,e,r)=gammaMap*loglik; %probability for each partile
                end
                
            end
            
            if k==2
                i(:,r)= maxInd(s(:,1,r),NP); %if first time step only use one particle
            else
                i(:,r)=maxInd(s(:,:,r),NP); %find NP highest probaility assignment from all particles
            end
            
            assign=i;
            for e=1:NP
                prepv(e,r)=div(assign(e,r)-1,gammaLen)+1;
                i(e,r)=rem(i(e,r)-1,gammaLen)+1; % convert data association back.
                %               inverseGamma(:,e,r)=invgamma(allgamma(i(e,r),:)');
                ESTstate(:,:,:,e,r)=ESTref(:,:,:,prepv(e,r),r);
                ESTcov(:,:,:,e,r)=ESTcovref(:,:,:,prepv(e,r),r);
                
                StateEstimate(:,:,:,e,r)=StateRef(:,:,:,prepv(e,r),r);
                CovEstimate(:,:,:,e,r)=StatecovRef(:,:,:,prepv(e,r),r);
                
                %s(i(e,r),e,r)=s(i(e,r),prepv(e,r),r);
                logprop(e,r,k)=s(i(e,r),prepv(e,r),r);
                assoc(:,k,e,r)=allgamma(:,i(e,r));
            end
        end
        
        
        [ESTstate, ESTcov]= kalMeasUpdate(m(:,k,:,:),M(:,:,k,:,:).^2,ESTstate,ESTcov,StateEstimate,CovEstimate,i,NT,NP,nDim,runs,allgamma,idx);
        
        Estate(:,k,:,:,:)=ESTstate;
        Ecov(:,:,k,:,:,:)=ESTcov;
    end
    for k=2:time
    for r=1:runs
        for e=1:NP
            %particleProb(e,r) = particleProb(e,r)* exp(s(i(e,r),e,r));
            particleProb(e,r,k) = exp(logprop(e,r,k)-logprop(1,r,k));
        end
        particleProb(:,r,k)=particleProb(:,r)./sum(particleProb(:,r));
    end
    end
    
    for r=1:runs
    for k=2:time
        for j=1:NP
        EST(:,k,:,r)=EST(:,k,:,r)+Estate(:,k,:,j,r)*particleProb(j,r,k);
        estcov(:,:,k,:,r)=estcov(:,:,k,:,r)+Ecov(:,:,k,:,j,r)*particleProb(j,r,k);
        end
    end
    end
    
    testsample=squeeze(S(:,end,1,:))+mvnrnd(zeros(nDim,1),squeeze(M(:,:,end,1,:)).^2)';
    postobs=zeros(NT,runs);
    for r=1:runs
        for u=1:NT
            for e=1:NP
                mu=ESTstate(:,:,u,e,r);
                cm=ESTcov(:,:,u,e,r)+covM(testsample(:,r),a,b,c,d).^2;
                postobs(u,r)=postobs(u,r)+particleProb(e,r)*mvnpdf(testsample(1:3,r),mu(1:3),cm(1:3,1:3));
            end
        end
        postobs(:,r)=postobs(:,r)./sum(postobs(:,r));
    end
    
    probtarget=postobs; %probability of final sample from each object

%% graphics parameters
msize=7; %marker size
lsize=2; %line size
axnLat=12; %size od non latex labels
axFS=10;

%%
figure()
subplot(2,3,[1,2,3])
hold on
t=T*hz; %number of timesteps
tstart=150; 
dt=1/60; %dt
Time=linspace(0,T,t); %time in seconds
obj=1;%object to compare
incor=find(probtarget(1,:)-probtarget(2,:)<0); %we drew sample from object 1 so should always be higher than object 2
[row,col]=find(assoc(1,2:end,1,incor(1))-allgamma(1,2)~=0&assoc(1,2:end,2,incor(1))-allgamma(1,2)~=0); %find when inferred data association was incorrect
Cpoint=col(1)-1; %the posterior prior to change point
trial=incor(1);

plot(Time(tstart:t),squeeze(S(1,tstart:t,1,trial)),'b--','LineWidth',lsize) %plot estimate object 1
plot(Time(tstart:t),squeeze(S(1,tstart:t,2,trial)),'r--','LineWidth',lsize) %plot estimate object 2
plot(Time(tstart:t),squeeze(m(1,tstart:t,1,trial)),'b.','MarkerSize',msize,'LineWidth',lsize) %plot measurment object 1
plot(Time(tstart:t),squeeze(m(1,tstart:t,2,trial)),'r.','MarkerSize',msize,'LineWidth',lsize) %plot measurment object 2
plot(Time(tstart:t),squeeze(EST(1,tstart:t,obj,trial)),'g-','LineWidth',lsize) %plot estimate of object 1


set(gca, 'Ticklength', [0 0])
set(gca, 'LineWidth', 1.5,'FontName','Arial')
box off
set(gca,'FontSize',axFS,'FontName','Arial')
xlim([2.5,4.5])
ylim([-0.15,0.15])
setTicks(5,5,2,2)
ylabel('Position (m)','Fontsize',axnLat,'FontName','Arial')
xlabel('Time(s)','Fontsize',axnLat,'FontName','Arial')
drawline(Time(Cpoint),'dir','vert','linewidth',1.5)
   set(gca,'TickDir','in')
            set(gca, 'Ticklength', [0.02 0])

h=legend('True Object 1','True Object 2','Measurement 1','Measurement 2','Estimate Object 1','Orientation','horizontal');
legend boxoff
set(h,'Location','southeast','FontSize',axFS,'FontName','Arial')
pos=get(h,'Position');
pos=pos+[0.2 0.31 0 0];
set(h,'Position',pos)
clear h


x = -0.5:.001:0.5; %// x axis
subplot(2,3,4)
hold on
predm1=A(1,1)*EST(1,Cpoint,obj,trial); %x_t+1 = A*x %compute prediction distribtuons
predvar1=A(1,1)^2*estcov(1,1,Cpoint,obj,trial)+Q(1,1)+(M(1,1,Cpoint+1,obj,trial))^2; %A*x_cov*A^T+Q+M % measurment probability variance 

predm2=A(1,1)*EST(1,Cpoint,obj+1,trial); %x_t+1 = A*x %compute prediction distribtuons
predvar2=A(1,1)^2*estcov(1,1,Cpoint,obj+1,trial)+Q(1,1)+(M(1,1,Cpoint+1,obj+1,trial))^2; %A*x_cov*A^T+Q+M % measurment probability variance 

plot(x,normpdf(x,predm1,sqrt(predvar1)),'r','LineWidth',lsize) % plot measurment likelihood distribution 
plot(x,normpdf(x,predm2,sqrt(predvar2)),'b','LineWidth',lsize) % plot measurment likelihood distribution 
plot(m(1,Cpoint+1,obj,trial),0,'or','MarkerFaceColor','r','MarkerSize',msize) %plot measurment of confusion
plot(m(1,Cpoint+1,obj+1,trial),0,'ob','MarkerFaceColor','b','MarkerSize',msize)
set(gca, 'Ticklength', [0 0])
set(gca, 'LineWidth', 1.5)
box off
set(gca,'FontSize',axFS,'FontSize',axFS,'FontName','Arial')


ylabel('Likelihood','Fontsize',axnLat)
xlabel('Position (m)','Fontsize',axnLat)
h=legend('Object 1','Object 2','Measurement 1','Measurement 2');
legend boxoff
set(h,'Location','northeast','FontSize',axFS,'FontName','Arial')
xlim([-0.05 0.07])
   set(gca,'TickDir','in')
            set(gca, 'Ticklength', [0.02 0])
setTicks(5,3,2,2)

subplot(2,3,5)
hold all
% take posterior and comput prediction to calc measurement probability.
mu = [A(1,1)*EST(1,Cpoint,obj,trial),0]; %A*x %compute prediction distribtuons
sigma = [A(1,1)^2*estcov(1,1,Cpoint,obj,trial)+Q(1,1)+M(1,1,Cpoint+1,obj,trial)^2 0;...
    0 A(2,2)^2*estcov(1,1,Cpoint,obj,trial)+Q(2,2)+M(1,1,Cpoint+1,obj,trial)^2]; %// data
[X, Y] = meshgrid(x,x); %// all combinations of x
Z = mvnpdf([X(:) Y(:)],mu,sigma); %// compute Gaussian pdf
Z = reshape(Z,size(X)); %// put into same size as X, Y
contour(X,Y,Z,4,'r','LineWidth',lsize)%// contour plot; set same scale for x and y...
%imagesc(x,y,Z)
%colormap(gray)
plot(m(1,Cpoint+1,obj,trial),0,'or','MarkerFaceColor','r','MarkerSize',msize)
%xlim([-0.05 0.1])
%ylim([-0.05 0.1])



offset=0.05;
mu = [A(1,1)*EST(1,Cpoint,obj+1,trial),offset]; %// data
sigma = [A(1,1)^2*estcov(1,1,Cpoint,obj+1,trial)+Q(1,1)+M(1,1,Cpoint+1,obj+1,trial)^2 0; 0 A(2,2)^2*estcov(1,1,Cpoint+1,obj+1,trial)+Q(2,2)+M(1,1,Cpoint+1,obj+1,trial)^2]; %// data
[X, Y] = meshgrid(x,x); %// all combinations of x, y
Z = mvnpdf([X(:) Y(:)],mu,sigma); %// compute Gaussian pdf
Z = reshape(Z,size(X)); %// put into same size as X, Y
contour(X,Y,Z,4,'b','LineWidth',lsize) %// contour plot; set same scale for x and y...
%imagesc(x,y,Z)
%colorbar
plot(m(1,Cpoint+1,obj+1,trial),offset,'ob','MarkerFaceColor','b','MarkerSize',msize,'LineWidth',lsize)
xlim([-0.05 0.07])
ylim([-0.04  0.08])
set(gca, 'Ticklength', [0 0])
    set(gca, 'LineWidth', 1.5)
    box off
     set(gca,'FontSize',axFS,'FontSize',axFS,'FontName','Arial')
xlabel('Position (m)','Fontsize',axnLat)
ylabel('Position (m)','Fontsize',axnLat)
   set(gca,'TickDir','in')
            set(gca, 'Ticklength', [0.02 0])
setTicks(5,3,2,2)
%%
L=0.9;
k=((L+1)*sigV^2)/(2*sigX^2);
sig=0.5*sqrt(sigV^2*(4*sigX^2-4*L^2*sigX^2-sigV^2+L^2*sigV^2)/sigX^2);
Ao=([[1-k,0,0, L,0,0] ;[0,1-k, 0,0,L,0];[0,0, 1-k,0,0,L];[-k,0,0 L,0,0];[0,-k, 0,0,L,0];[0,0, -k,0,0,L]]); %A matrix FE
Qo=sig^2; %Q matrix FE

L=L1; %dynamics are correct
k=((L+1)*sigV^2)/(2*sigX^2);
sig=0.5*sqrt(sigV^2*(4*sigX^2-4*L^2*sigX^2-sigV^2+L^2*sigV^2)/sigX^2);
L=L2; %but velocity is zero
Asub=([[1-k,0,0, L,0,0] ;[0,1-k, 0,0,L,0];[0,0, 1-k,0,0,L];[-k,0,0 L,0,0];[0,-k, 0,0,L,0];[0,0, -k,0,0,L]]); %A matrix NE
Qsub=sig^2; %Q matrix NE-cd

Aopt=[Ao(1,1) Ao(1,4); Ao(4,1), Ao(4,4)];
As=[Asub(1,1) Asub(1,4); Asub(4,1), Asub(4,4)];
Qv=[Q(1,1) Q(1,4); Q(4,1), Q(4,4)];

subplot(2,3,6)
hold on


pos=[0.02;0.02];%set example position and velocityEST([1 2],Cpoint,obj,trial);
MnoiseEx=[stdPos(pos(1),c),0; 0 0.05*pos(2)];
obs=pos+mvnrnd([0 ;0],Qv)'+mvnrnd([0 ;0],MnoiseEx.^2)';

x = -0.2:.001:0.25; %// x axis
y = -0.1:.001:0.1; %// y axis
muopt=Aopt*pos; %predicted mean
musub=As*pos;


%covest=estcov([1 4],[1 4],Cpoint,obj,trial)
covest=Qv+MnoiseEx.^2;
sigopt=Aopt*covest*Aopt' +Qo+M([1,4],[1,4],Cpoint+1,obj,trial)^2;
sigsub=As*covest*As' +Qsub+M([1,4],[1,4],Cpoint+1,obj,trial)^2;
[X, Y] = meshgrid(x,y); %// all combinations of x, y


Z = mvnpdf([X(:) Y(:)],muopt',sigopt); %// compute Gaussian pdf
Z = reshape(Z,size(X)); %// put into same size as X, Y
contour(X,Y,Z,4,'m','LineWidth',lsize)



Z = mvnpdf([X(:) Y(:)],musub',sigsub); %// compute Gaussian pdf
Z = reshape(Z,size(X)); %// put into same size as X, Y
contour(X,Y,Z,4,'c','LineWidth',lsize)


Z = mvnpdf([X(:) Y(:)],pos',covest); %// compute Gaussian pdf
Z = reshape(Z,size(X)); %// put into same size as X, Y
contour(X,Y,Z,4,'k','LineWidth',lsize)


plot(obs(1),obs(2),'og','MarkerFaceColor','g','MarkerSize',msize,'LineWidth',lsize)
h=legend('NF prediction','NE-cd prediction','Posterior','Next measurement');
set(h,'FontSize',axFS,'FontName','Arial');
xlim([-0.05 0.07])
ylim([-0.01 0.03])
set(gca,'TickDir','in')
set(gca, 'Ticklength', [0.02 0])
set(gca, 'LineWidth', 1.5)
  legend boxoff
     set(gca,'FontSize',axFS)
xlabel('Position (m)','FontSize',axnLat)
ylabel('Velocity (m/frame)','FontSize',axnLat)
   set(gca,'TickDir','in')
            set(gca, 'Ticklength', [0.02 0])
setTicks(5,3,2,2)
%%
fig=gcf;
%# centimeters units
%a4=21 by 29.7
X = 18;                  %# A3 paper size
Y = 12;                  %# A3 paper size
xMargin = 1;               %# left/right margins from page borders
yMargin = 1;               %# bottom/top margins from page borders
xSize = X - 2*xMargin;     %# figure size on paper (widht & hieght)
ySize = Y - 2*yMargin;     %# figure size on paper (widht & hieght)

%# figure size displayed on screen (50% scaled, but same aspect ratio)
set(fig, 'Units','centimeters', 'Position',[0 0 xSize ySize])
movegui(fig, 'center')

%# figure size printed on paper
set(fig, 'PaperUnits','centimeters')
set(fig, 'PaperSize',[X Y])
set(fig, 'PaperPosition',[xMargin yMargin xSize ySize])
set(fig, 'PaperOrientation','portrait')
parentpath = cd(cd('..'));
%print(fullfile(parentpath,'figures/explanationFigure'), '-dsvg', '-r400')
print(fullfile(parentpath,'figures/explanationFigure'), '-dpdf', '-r800')


