%[H,M] = trk(S,X,OP,US,<M>)
%
% Synopsis: Tracks the feature at location X of frame 1 of sequence S
%    with user selected optimisation algorithm OP and update strategy US.
% Eg: H = trk(mov/binSz,[x y xsz ysz],[2 3 1 1 1 2 binSz],[1 1 0 1 1]);
% Outputs:
%    H - Tracking History
%   *M - Model storing state of current track. (* Also an output)
% Inputs:
%    S  - Sequence of Images (3D cube of pixels: x*y*time)
%    X  - Initial position of feature (2,3,4,6 values for translation,
%         euclidean, similarity or affine coordinates) + 2 for rect size
%    OP - Optimisation Parameters (5-7 value vector, type "warthog" for
%         detailed help)
%    US - Verbosity (0/1/2 for quiet/verbose/movie) and...
%         Presearch params (window size and density) and...
%         Update strategy: 1-6 for:
%         none, naive, strategic*, naive+real, strategic+real* and
%         forward-back respectively. (The * indicates that a second element
%         specifying a tolerence parameter is also required.)

% Author:  Nicholas Dowson
% Revised: 6 Jul 2006 - 11.22
% (c) CVSSP, University of Surrey
%---------------------------------------------------------------------

function [oHstry, gM] = trk(iSeq, iX, iOP, iUpd, iM)
  
  if numel(findstr(path,'zPers/dev/zuseful'))==0
    addpath ~/zPers/dev/zuseful;
  end
  
  %% Error checking
  if nargin < 4
    disp('trk error: Less than 4 parameters given. Exiting.');
    disp('Usage: [H] = trk(S,X,OP,Upd). Type "help trk" for more info.');
    return;
  end

  clear global gOP gUpd gM gBin;
  global gOP gUpd gM gBin;
  
  if exist('iM','var'), gM = iM; end
  gOP = iOP(:); gUpd = iUpd;
  clear iUpd iOP;

  fX = iX(1:gOP(1));
  kShp = iX( gOP(1)+(1:2) ); %Auto-Error: Shape of the template
  global ix iy;
  ix=1; iy=2;
  if gOP(1)==6, ix=5; iy=6; end
  
  % Check number of update parameters
  if numel(gUpd)<4
    disp(['trk error: gUpd needs >=4 params: [Verb, PreSrchSz, ' ...
          'PreSrchRes, Strategy, UpdateParam] ']);
    return;
  end
  % Check for update parameter
  if numel(gUpd)<5
    if (gUpd(4)==3)||(gUpd(4)==5)
      gUpd(5) = 1;
      disp(['trk warning: Strategic Update needs Param. Using 1.']);
    end
  end
  
  % Check for 7th param in gOP
  if numel(gOP)<6,
    disp(['trk error: gOP needs 6 parameters: [wrp, msr, ord, intrp, '...
          'rate, optMeth]']);
    return;
  end
  gBin = 256 / 2^ceil(log(max(iSeq(:)))/log(2));
  %gOP = gOP(1:6);
  
  %% Main part of algorithm
    
  %Put model together
  if nargin<5
    if gUpd(1)~=0 disp('trk msg: Creating model.'); end
    gM.T0 = wartr( iSeq(:,:,1), fX, kShp );
    gM.TN = gM.T0;
    gM.TNo = gM.T0;
    gM.XtX0 = zeros(gOP(1),1);
  end

  % Initialisation
  t = gOP; t(1) = 2;
  t = warts(gM.T0,gM.T0,(size(gM.T0)-1)/2,t);
  oHstry  = [fX(:)' t 1]';

  %Load the kernel from the first frame
  fI  = iSeq(:,:,1);
  fXo = zeros(gOP(1),1); %Get pixel offset

  gM.updFrame=1;  
  trkStatus( fI, oHstry ); %Display

  %For all the frames from frame 2
  for lFrm = [2:size(iSeq,3)]
    %Find position of best mutual information
    fIp = fI;
    fI = iSeq(:,:,lFrm);

    fXp = preSearch( fI, fX ); %Presearch
    gM.TNo = gM.TN; %Store old template
    
    switch gUpd(4),
     case 1, %No update
      [fY,fX,fH,gM.TN ] = warthog( gM.T0, fI, fXp, gOP ); %Reg
      
     case 2, %Naive update
      [fY,fX,fH,gM.TN ] = warthog( gM.TN, fI, fXp, gOP ); %Reg
     
     case 3, %Strategic update
      [fY,fX          ] = warthog( gM.TN, fI, fX, gOP );
      gM.fX1 = fX;
      [fY,fX,fH,gM.TNt] = warthog( gM.T0, fI, fX, gOP );
      %Update if correction is small enough
      t = gM.fX1(1:2)-fX(1:2);
      if t'*t < gUpd(5)^2, gM.TN = gM.TNt; gM.updFrame=lFrm; end;
     
     case 4, %Naive + ZI update------------------------------
      if ~isfield(gM,'XtX0') gM.XtX0 = zeros(gOP(1),1); end
      [fY,fXtI      ] = warthog( gM.TN, fI, fXp, gOP );%Find best matching pos
      fX0I = fXtI - gM.XtX0; %From Xt:I and Xt:X0, get X0:I
      fXniI= zeros(gOP(1),1); %Get zero interp position
      fXniI(ix:iy) = round(fX0I(ix:iy)); 
      gM.TN = wartr( fI, fXniI(ix:iy), size(gM.T0) ); %Update template
      gM.XtX0 = fXniI - fX0I; %Get new Xt:X0
      fX = fXniI; %Back into main function variables names
      clear fX0I fXtI fXniI; %Clear extraneous variables
      
     case 5, %Strategic + ZI update-------------------------
      if ~isfield(gM,'XtX0') gM.XtX0 = zeros(gOP(1),1); end
      [fY,fXtI      ] = warthog( gM.TN, fI, fXp,  gOP );
      eX0I = fXtI - gM.XtX0; %From Xt:I and Xt:X0, estimate X0:I
      [fY,fX0I      ] = warthog( gM.T0, fI, eX0I, gOP ); %Correct fX0I
      %Block update if correction is large
      if sqrt(sum((eX0I(ix:iy)-fX0I(ix:iy)).^2)) < gUpd(5)
        fXniI= zeros(gOP(1),1); %Get zero interp position
        fXniI(ix:iy) = round(fX0I(ix:iy)); 
        gM.TN = wartr( fI, fXniI(ix:iy), size(gM.T0) ); %Update template
        gM.XtX0 = fXniI - fX0I; %Get new Xt:X0
        fX = fXniI; %Back into main function variables names
        clear fXniI; %Clear extraneous variables
        gM.updFrame=lFrm;
      else
        fX = fX0I;
      end
      clear eX0I fX0I fXtI;
% $$$       
% $$$      case 6, %Naive Forwards-Back track---------------------
% $$$       fXp = fX;
% $$$       fTp = gM.TN; %Store previous template
% $$$       [fY,fX,fH,gM.TN ] = warthog( gM.TN, fI , fX, gOP ); %Forwards track
% $$$       [fYb,fXb,fHb,fTb ] =warthog( gM.TN, fIp, fXp, gOP ); %Backwards track
% $$$       fX=fX-1.0*(fXp-fXb);
% $$$       gM.TN = wartr(fI,fX,size(gM.T0));
% $$$       clear fYb fXb fHb fTb;
    end
      
    oHstry = [oHstry, [fX(:)' fY lFrm]' ]; %History
    trkStatus( fI, oHstry ); %Display
  end
  
%---------------------------------------------------------------------
function oX = preSearch(iI,iX)
  
  global gOP gM gUpd;
  
  [gx gy] = ndgrid( -gUpd(2):gUpd(3):gUpd(2) );
  if numel(gx)==0
    oX = iX;
  else    
    fX = repmat( iX(:), [1 numel(gx)] );
    
    if gOP(1) == 6, %Affine warps
      fX(5,:) = iX(5)+gx(:)';
      fX(6,:) = iX(6)+gy(:)';
    else %Other warps
      fX(1,:) = iX(1)+gx(:)';
      fX(2,:) = iX(2)+gy(:)';
    end
    
    [fY] = warts( gM.TN, iI, fX, gOP ); %Measure values on grid

    [t idx] = min(fY);
    
    oX = fX(:,idx(1));
  end

%---------------------------------------------------------
% Shows the status of a tracked region in an image
function oFig = trkStatus( iI, iHstry )

  global gM gUpd gBin gOP ix iy;
  
  % Are we displaying
  if gUpd(1)==0
    oFig = 0;
    return;
  end
    
  %% Draw the image
  oFig = figure(2);
  clf;
  imagesc( iI*gBin );
  colormap(gray(256));
  t = get(gcf,'Position');
  set(gcf,'Position',[t(1:2) fliplr(size(iI))*3]);
  set(gca,'Position',[0 0 1 1]);

  hold on;
  axis off;
  
  % What measurement is being used?
  switch gOP(2),
   case 1, tMsr = 'SSD';
   case 2, tMsr = 'NC';
   case 3, tMsr = 'MI-std';
   case 4, tMsr = 'MI-pve';
   case 5, tMsr = 'MI-ipz';
   case 6, tMsr = 'MI-npw';
  end
  th = text( 2,10, ['Msr: ' tMsr]);
  set( th,'HorizontalAlignment','Left','VerticalAlignment','Top', ...
          'BackgroundColor',[1 1 1]);   
  
  %% Write Frame number
  fFrm = size(iHstry,2);
  th = text( 2,2, ['m = ' num2str(fFrm)] );
  set( th,'HorizontalAlignment','Left','VerticalAlignment','Top', ...
          'BackgroundColor',[1 1 1]); 
 
  %% Get previous and current X values
  if ~isfield(gM,'fX1'), Xpc = iHstry(:,end); else Xpc = gM.fX1; end
  if size(iHstry,2)==1, iHstry = [iHstry iHstry]; end;
  if size(iHstry,1)==8, %Deal with affine warps    
    t = iHstry(5:6,end-1:end);
    Xpc = Xpc(5:6)+1;
  else
    t = iHstry(1:2,end-1:end);
    Xpc = Xpc(1:2)+1;
  end
  Xp = t(:,1)+1;
  %Xc = t(:,2);  
  
  %% Selectively draw T0, TN, Tx, and update
  [txi tyi] = size(iI);
  [txf tyf] = size(gM.T0);
  switch gUpd(4),
   case 1,
    labelledImg(tyi-tyf-2,2,gM.T0,...
                'No Update: f_t <- f_x^{(0)}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,4+txf,gM.TN,...
                'Exemplar: f_x^{(m)}',[.3 .7 .7]);
    %line([Xp(2) Xc(2)],[Xp(1) Xc(1)],'Color',[.6 .2 .4]);
    plot(Xp(2),Xp(1),'Color',[1 .2 .4],'Marker','+');
   case 2,
    labelledImg(tyi-tyf-2,2,gM.TNo,...
                'Naive Update: f_t <- f_x^{(m-1)}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,4+txf,gM.TN,...
                'Exemplar: f_x^{(m)}',[.3 .7 .7]);
    plot(Xp(2),Xp(1),'Color',[1 .2 .4],'Marker','+');    
   case 3,
    labelledImg(tyi-tyf-2,2     ,gM.TNo,...
                'Strat. Update: f_t <- f_x^{(m'')}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,4+txf,gM.T0,...
                'Strat. Update: f_t <- f_x^{(0)}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,6+txf*2,gM.TN,...
                'Exemplar: f_x^{(m)}',[.3 .7 .7]);
    plot(Xp(2),Xp(1),'Color',[1 .2 .4],'Marker','+');
    plot(Xpc(2),Xpc(1),'Color',[1 1 .4],'Marker','+');
    th = text( 2,18, ['m'' = ' num2str(gM.updFrame)] );
    set( th,'HorizontalAlignment','Left','VerticalAlignment','Top', ...
          'BackgroundColor',[1 1 1]); 
   case 4,
    labelledImg(tyi-tyf-2,2,gM.TNo,...
                'ZIP Naive Update: f_t <- f_x^{(m-1)}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,4+txf,gM.TN,...
                'Exemplar: f_x^{(m)}',[.3 .7 .7]);
    plot(tyi-tyf/2-2,4+txf*1.5,'Color',[.3 .7 .7],'Marker','+');
    plot(tyi-tyf/2-2+gM.XtX0(iy)*3,4+txf*1.5+gM.XtX0(ix)*3,...
         'Color',[0 1 0],'Marker','+');
    plot(Xp(2),Xp(1),'Color',[1 .2 .4],'Marker','+');    
   case 5,
    labelledImg(tyi-tyf-2,2     ,gM.TNo,...
                'Strat. Update: f_t <- f_x^{(m'')}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,4+txf,gM.T0,...
                'Strat. Update: f_t <- f_x^{(0)}',[.6 .2 .4]);
    labelledImg(tyi-tyf-2,6+txf*2,gM.TN,...
                'Exemplar: f_x^{(m)}',[.3 .7 .7]);
    plot(Xp(2),Xp(1),'Color',[1 .2 .4],'Marker','+');
    plot(Xpc(2),Xpc(1),'Color',[1 1 .4],'Marker','+');
    plot(tyi-tyf/2-2,6+txf*2.5,'Color',[.3 .7 .7],'Marker','+');
    plot(tyi-tyf/2-2+gM.XtX0(iy)*3,6+txf*2.5+gM.XtX0(ix)*3,...
         'Color',[0 1 0],'Marker','+');
    th = text( 2,18, ['m'' = ' num2str(gM.updFrame)] );
    set( th,'HorizontalAlignment','Left','VerticalAlignment','Top', ...
          'BackgroundColor',[1 1 1]); 
  end
  
  switch gUpd(4),
   case 1, tUpd = 'None';
   case 2, tUpd = 'Naive';
   case 3, tUpd = 'Strat';
   case 4, tUpd = 'ZNaive';
   case 5, tUpd = 'ZStrat';
  end
  
  %% Draw the registered position
  t = (size(gM.T0)-1)/2;
  
  fRect1 = [-t(1)  t(1)  t(1) -t(1)
           -t(2) -t(2)  t(2)  t(2)
            1     1     1     1
          ];
  
  fX = iHstry(1:end-2,end);
  % Get rotation matrix
  switch numel(fX),
   case 2, W = [eye(2) fX(1:2)+1];
   case 3, 
    tc = cos(fX(3)); ts = sin(fX(3));
    W = [[tc ts; -ts tc] fX(1:2)+1];
   case 4,
    tc = cos(fX(3)); ts = sin(fX(3));
    W = [[tc ts; -ts tc]*fX(4) fX(1:2)+1];
   case 6, W = reshape(fX,[2 3]); W([1 4 5 6]) = W([1 4 5 6])+1;
  end
  fRect = W*fRect1;
    
  patch( fRect(2,:),fRect(1,:),[0 1 0],'FaceAlpha',0.1,'EdgeColor',[0 1 0]);

  plot(fX(2)+1,fX(1)+1,'g+');
  
  %% Draw a graph of historical scores
  ty = iHstry(end-1,:);
  ty = ty / max([1 abs(ty)]);
  tmy = 0.1*size(iI,1);    
  tcy = size(iI,1);
  
  tyScl = max(1,size(iHstry,2)/size(iI,2));
  tx = [1:size(iHstry,2)]/tyScl;
  
  if ty(end)>0, tmy=-tmy; end
  plot( tx, tcy + tmy*ty, 'k-', 'LineWidth', 3 );
  plot( tx, tcy + tmy*ty, 'y-' );

  %%Printing stuff
  if gUpd(1)==2
    if fFrm<10, fPad = '000';
    elseif fFrm<100, fPad = '00';
    elseif fFrm<1000, fPad = '0';
    else fPad = '';
    end
    stdFontSz(gcf,20);
    set(gcf, 'PaperPositionMode', 'auto');
    set(gcf,'PaperUnits','points');
    t = [480 360]*2;
    set(gcf,'PaperPosition',[0 0 t]);
    set(gcf,'PaperSize',[t]);
    print('-dpng','-r72',['vid/' tUpd tMsr fPad num2str(fFrm) '.png']);
  end
  
  %% Pause
  hold off;
  pause(0.01);

%-----------------------------------------------------------------------
function labelledImg( x,y, I, iText, iColor )

  global gBin;
  
  th = text( x-2, y+4, iText,...
             'HorizontalAlignment','Right','VerticalAlignment','Top',...
             'BackgroundColor',iColor,'Color',[1 1 1]);
  
  imagesc( x+[0 size(I,1)], y+[0 size(I,2)], I*gBin );
  
  tx1 = x-.5; tx2 = x+size(I,1)+.5;
  ty1 = y-.5; ty2 = y+size(I,2)+.5;
  
  patch( [tx1 tx2 tx2 tx1], [ty1 ty1 ty2 ty2], iColor, ...
         'FaceAlpha',0.1, ...
         'EdgeColor', iColor );