/*
 * Estimate the egomotion parameters and track them via Kalman filter 
 *-------------------------------------------------------------------
 */

#include "est_and_track.h"

//LSE routine
#include "mathmi.h"

//What is a small and what is a big number?
#define BIG 1e6
#define SMALL 1e-6

//Use nested Kalman or not?
#define KALMAN_NESTED 1

#define ROTATION   1

#define STDOUT 0
//Buffer for the motion field, we always need two consecutive ones!
static vec_field *v_buffer[2];

//Indicate start of routine
static int ego_and_track_egomotion=1;

//Count how many vec_field are currently in buffer
static int count_ego=0;

//Translational component
float **Aest,*best,*xest;

//Rotational component
float **Arot,*brot,*xrot;

//Kalman filter
//Translational
static float uest,vest;
static float uest_p=0, uest_m;
static float vest_p=0, vest_m;
static float variance_p=1e10,variance_m;

//Kalman filter
//Rotational
static float urot_p=0, urot_m;
static float vrot_p=0, vrot_m;
static float wrot_p=0, wrot_m;
static float variance_p_rot=1e10,variance_m_rot;

//Estimate egomotion
int est_and_track_egomotion(vec_field *v_field, 
			    float *foe_x,float *foe_y,
			    float *foe_x_robust,float *foe_y_robust,
			    float *foe_x_kal, float *foe_y_kal,
			    float *Omx_est, 
			    float *Omy_est, 
			    float *Omz_est,
			    float *Omx_robust, 
			    float *Omy_robust, 
			    float *Omz_robust,
			    float *Omx_est_kal, 
			    float *Omy_est_kal, 
			    float *Omz_est_kal,
			    int reset_kal)
  
{
  int a,i;
  int x0,y0,x1,y1;
  int xnew,ynew,g;
  float u1,v1;
  float residual;
  double K;
  int ok;
  
  //Translational part
  //How many areas to sum for LSE translational components?
  //LSE = least square estimation
  int count_t=4, sqrt_count_t=2;
  float u[2],v[2];
  float sum_xv[2], sum_yu[2];
  
  //Rotation part
  //How many vectors to use for the
  //LSE estimation of the rotational components?
  int count_r=30;
  
  //Image Center  
  int xCenter=v_field->width>>1;
  int yCenter=v_field->height>>1;
  
  int ops=0;
 
  //Reset the Kalman values
  if(reset_kal==1)
    {
      uest_p=0;
      vest_p=0;
      variance_p=1e10;
      
      urot_p=0;
      vrot_p=0;
      wrot_p=0;
      variance_p_rot=1e10;
    }
  
  //First time the routins is called => Do init stuff
  if(ego_and_track_egomotion==1)
    {
      count_ego=0;
      
      //Copy vector field
      v_buffer[count_ego]=copy_vec_field(v_field);
      
      //Malloc for LSE
      Aest=matrixf(1,count_t,1,2);
      best=vectorf(1,count_t);
      xest=vectorf(1,2);
      
      Arot=matrixf(1,count_r,1,3);
      brot=vectorf(1,count_r);
      xrot=vectorf(1,3);
      
      //Indicate start of routine
      uest_p=0; vest_p=0; variance_p=1e10;
      urot_p=0; vrot_p=0; wrot_p=0; variance_p_rot=1e10;
      
      count_ego++;
      
      //Init stuff done
      ego_and_track_egomotion=0;
      ok=0;
    }
  
  ////////////////////////////
  //Continous processing then
  ////////////////////////////
  else
    {
      //Copy new vector field into ring buffer
      v_buffer[count_ego]=copy_vec_field(v_field);
      
      /////////////////////////////////////
      //Solving for the translational part
      /////////////////////////////////////
      //calculate step size
      a=(MIN(v_field->width,v_field->height)-
	 2*MAX(v_field->xoffset,v_field->yoffset)) / sqrt_count_t;
      
      //Start coordinates
      x0=v_field->xoffset;
      y0=v_field->yoffset;
      
      ops=0;
      //Sum rectangular areas
      for(i=1; i <= count_t ; i++)
	{
	  x1=x0 + a; y1=y0 + a;
	  
	  //Sum components => increases robustness
	  sum_it_now(v_buffer[(count_ego + 1)%2], //t1
		     v_buffer[count_ego], //t2
		     x0, y0, x1, y1, //area to sum 
		     u, v, //output  
		     sum_xv, sum_yu);
	  
	  //build Matrix equation system
	  Aest[i][1] = v[1]-v[0];
	  Aest[i][2] = u[1]-u[0];
	  best[i] = sum_yu[1]-sum_yu[0]-sum_xv[1]+sum_xv[0];
	  
	  if(FABS(Aest[i][1])<SMALL && FABS(Aest[i][2])<SMALL && 
	     FABS(best[i])<SMALL)
	    ops++;
	  
	  //next area
	  if(i%sqrt_count_t == 0 )
	    {
	      y0=y1; 
	      x0 = v_field->xoffset;
	    }
	  else
	    x0=x1;
	}
      
      if(ops==0 )
	//Solve the equation system!
	ok=least_square(Aest, 2, count_t,
			best, count_t,
			xest, 2,
			&residual);
      else if (ops == count_t)
	{//Everything zero!
	  xest[0]=0.0;
	  xest[1]=0.0;
	  
	  residual=0.0;
	  
	  uest_p=0;
	  vest_p=0;
	  variance_p=1e10;
	  
	  ok=1;
	}
      else
	ok=0;
      
      if(ok==1 && (FABS(xest[1]) < 200.0) && (FABS(xest[2]) < 200.0) )
	{
	  *foe_x= xest[1];
	  *foe_y= xest[2];
	}
      else
	//Could not estimate FOE
	ok=0;
      
      //Kalman Filter for translational component
      //-----------------------------------------
      //Measurements
      if(ok==1)
	{
	  //the actual variance of the measurement
	  variance_m=residual*residual;
	  
	  //measured values
	  uest_m=xest[1]; vest_m=xest[2];
	  
	  //Kalman - Filter gain  
	  K = variance_p / (variance_p + variance_m);
	  
	  //predict
	  uest_p = uest_p + K * (uest_m - uest_p);
	  vest_p = vest_p + K * (vest_m - vest_p);
	  
	  //Update = 2 ways
#if KALMAN_NESTED
	  variance_p = K * variance_m; 
#else
	  variance_p = variance_m - K * variance_m; 
#endif	 
	  *foe_x_kal= uest_p;
	  *foe_y_kal= vest_p;
	}
      else
	{/*
	  uest_p=0;
	  vest_p=0;
	  variance_p=1e10;
	 */
	  }

#if STDOUT
      fprintf(stdout,"----------------------------------------\n");
      fprintf(stdout,"Est:    Tx/Tz:%2.3f Ty/Tz:%2.3f\n",uest_m,vest_m);
      fprintf(stdout,"Kalman: Tx/Tz:%2.3f Ty/Tz:%2.3f <= K: %1.4f \n",
	      uest_p,vest_p, K);
#endif

#if ROTATION
      //////////////////////////////////
      //Rotation components  estimation
      //////////////////////////////////
      if(ok==1)
	{
	  //Calculate step size
	  a = v_buffer[count_ego]->datasize / (5*count_r) ;
	  g=0; i=1;
	  int fog=0;
	  
	  //uest_p=vest_p=0;
	  //Get vectors 
	  ops=0;
	  do
	    { 
	      //we can not find enough vectors, 
	      //look again but this time starting with the one to the right
	      if((unsigned int) (g+=a) > v_buffer[count_ego]->datasize)
		{
		  //Still can not find enough vectors -> Get out!
		  if((fog++) > a)
		    {
		      fprintf(stdout,"Can not find enough vectors!\n");
		      ok=0;
		      break;
		    }
		  
		  //Indicate how many times we started with the 
		  //neighbour to the right. Maximum a-times
		  g=fog;
		}
	      
	      if(v_buffer[count_ego]->valid[g]==1 && ok==1)
		{
		  //optic flow vectors
		  u1=v_buffer[count_ego]->u[g];
		  v1=-v_buffer[count_ego]->v[g];
		  
		  //x,y components expressed in an centered image
		  // coordinates
		  xnew=g%v_buffer[count_ego]->width-xCenter;
		  ynew=-(g/v_buffer[count_ego]->width-yCenter);
		  
		  //Build equation system and thereby eleminating the
		  //translational part!
		  //Equ. system: Arot*xrot=brot 
		  //(=A'rot'ational,b'rot'ational,etc.)
		  brot[i]=(xnew-uest_p)*v1+(vest_p-ynew)*u1;
		  
		  Arot[i][1]=(vest_p-ynew)*(xnew*ynew)+
		    (xnew-uest_p)*(1+ynew*ynew); //Omx
		  
		  Arot[i][2]=(ynew-vest_p)*(1+xnew*xnew)
		    +(uest_p-xnew)*(xnew*ynew); //Omy
		  
		  Arot[i][3]=(vest_p-ynew)*ynew+(uest_p-xnew)*xnew; //Omz  
		  
		  if(FABS(Arot[i][1])<SMALL && FABS(Arot[i][2])<SMALL &&  
		     FABS(brot[i])<SMALL && FABS(Arot[i][3])<SMALL)
		    ops++;
		  
		  i++;
		}
	      
	    } while( i <= count_r);
	  
	  if(ops==0 && ok==1)
	    //Solve equation system
	    ok=least_square(Arot, 3, count_r,
			    brot, count_r,
			    xrot, 3,
			    &residual);
	  else if (ops == count_r && ok==1)
	    {//Everything zero!
	      xrot[0]=0.0;
	      xrot[1]=0.0;
	      xrot[2]=0.0;
	      
	      residual=0.0;
	      
	      urot_p=0;
	      vrot_p=0;
	      wrot_p=0;
	      variance_p_rot=1e10;
	      
	      ok=1;
	    }
	  else
	    ok=0;
	  
	  //Kalman Filter for rotational component
	  //--------------------------------------
	  //Measurements
	  if(residual<BIG && residual>SMALL && ok==1)
	    {
	      *Omx_est=xrot[1];
	      *Omy_est=xrot[2];
	      *Omz_est=xrot[3];

	      variance_m_rot=residual*residual;
	      
	      urot_m=xrot[1]; 
	      vrot_m=xrot[2]; 
	      wrot_m=xrot[3];
	      
	      //Kalman - Filter gain 
	      K = variance_p_rot / (variance_p_rot + variance_m_rot);
	      
	      //Predict <=> Correct!
	      urot_p = urot_p + K * (urot_m - urot_p);
	      vrot_p = vrot_p + K * (vrot_m - vrot_p);
	      wrot_p = wrot_p + K * (wrot_m - wrot_p);
	  
	      //Update
#if KALMAN_NESTED
	      variance_p_rot = K * variance_m_rot; 
#else
	      variance_p_rot = variance_m_rot - K * variance_m_rot;
#endif	
	    }
	  
	  //Write information out
	  *Omx_est_kal=urot_p;
	  *Omy_est_kal=vrot_p;
	  *Omz_est_kal=wrot_p;
#if STDOUT             
	  fprintf(stdout,"Est.:   Omx:%1.4f Omy:%1.4f Omz:%1.4f\n",
		  urot_m,vrot_m,wrot_m);
	  fprintf(stdout,"Kalman: Omx:%1.4f Omy:%1.4f Omz:%1.4f\n",
		  urot_p,vrot_p,wrot_p);
	  fprintf(stdout,"----------------------------------------\n");
#endif
	}
#endif //ROTATION

      //Here comes the next vector field
      count_ego=(count_ego+1)%2;
      
      //but make space for new vector field
      freeVec_field(v_buffer[count_ego]);      
      
    } 
  
  //Everything o.k.!
  return ok;
}
//////////////////////////////////////////////////////////////////////
// Sums the horizontal and vertical components of the OF
// in the rectangular area defined by 'int x0, int y0, int x1, int y1'
//////////////////////////////////////////////////////////////////////
void sum_it_now(vec_field *field_t1,//t1
		vec_field *field_t2,//t2
		int x0,int y0,int x1,int y1,
		float  *u, float *v, 
		float *sum_xv, float *sum_yu)
{
  int i;
  int x,y;
  float sum_u_t1=0.0;float sum_v_t1=0.0;
  float sum_x_t1=0.0;float sum_y_t1=0.0;
  
  float sum_u_t2=0.0;float sum_v_t2=0.0;
  float sum_x_t2=0.0;float sum_y_t2=0.0;
 
  int x_Center=field_t1->width>>1;
  int y_Center=field_t1->height>>1;
  
  for(y=y0; y<y1; y++)
    { 
      for(x=x0; x<x1; x++)
	{
	  i=y * field_t1->width +x;
	  
	  if(field_t1->valid[i]==1 && field_t2->valid[i]==1)
	    {
	      sum_u_t1+=field_t1->u[i];
	      sum_u_t2+=field_t2->u[i];
	     
	      sum_v_t1+=field_t1->v[i];
	      sum_v_t2+=field_t2->v[i];
	      
	      sum_x_t1+=(-y + y_Center )*field_t1->u[i];
	      sum_x_t2+=(-y + y_Center )*field_t2->u[i];
	      
	      sum_y_t1+=(x - x_Center )*(-field_t1->v[i]);
	      sum_y_t2+=(x - x_Center )*(-field_t2->v[i]);
	    }
	}
    }
  
  u[0]=sum_u_t1;
  u[1]=sum_u_t2;
  
  v[0]=sum_v_t1;
  v[1]=sum_v_t2;
  
  sum_xv[0] = sum_y_t1;
  sum_xv[1] = sum_y_t2;
  
  sum_yu[0] = sum_x_t1;
  sum_yu[1] = sum_x_t2;
   
}

//////////////////////////////////////////////
// Derotate the vector field = subtract the 
// rotational part from the translational part
///////////////////////////////////////////////
vec_field *derotate(vec_field *field, 
		    int x0, int y0, int x1, int y1,
		    float Tx, float Ty, float Tz,float Z,
		    float Omx, float Omy, float Omz)
{
  int x,y,xnew,ynew;
  int i;
  float u,v;
  int xCenter=(field->width>>1);
  int yCenter=(field->height>>1);
  vec_field *tmp;
  
  tmp=newVec_field();
  
  resizeVec_field(tmp,field->width, field->height,
		  field->xoffset, field->yoffset);

  clearVec_field(tmp);
  
  for(y=y0; y<y1; y++)
    { 
      for(x=x0; x<x1; x++)
	{
	  i= y * field->width + x;
	  
	  xnew=x-xCenter;
	  ynew=-y+yCenter;
	  
	  if(field->valid[i]==1)
	    {
	      //Derotate
	      tmp->u[i]=field->u[i]-xnew*ynew*Omx+(1+xnew*xnew)*Omy-ynew*Omz;
	      tmp->v[i]=field->v[i]+(1+ynew*ynew)*Omx-xnew*ynew*Omy-xnew*Omz;
	      
	      ////////////////////////////////// 
	      //Cross check - Debuging purpose
	      //u=(xnew-Tx/Tz)*Tz/Z
	      //	+xnew*ynew*Omx-(1+xnew*xnew)*Omy+ynew*Omz;
	      //v=(ynew-Ty/Tz)*Tz/Z
	      //+(1+ynew*ynew)*Omx-xnew*ynew*Omy-xnew*Omz;
	      //////////////////////////////////
	      
	      tmp->valid[i]=1;
	    }
	}
    }

  return tmp;

}


