//Tss.java:  Three Stage Search
//Upper-left corner of macroblock is location of macroblock

class Tss {
  private int width;               //frame width
  private int height;              //frame height
  private Point2 level[][];  

  public Tss ( int w, int h )
  {
    width = w;
    height = h;
    level = new Point2[3][8];      //3 levels
    for ( int i = 0; i < 3; i++ )
      for ( int j = 0; j < 8; j++ )
        level[i][j] = new Point2();
  }

  //set the searching locations of stage n; p is the origin
  public void setLevel ( Point2 p, int n )
  {
    //searching centers for stage 0
    int rx[] = { -4, 0, 4, -4, 4, -4, 0, 4 };
    int ry[] = { -4, -4, -4, 0, 0, 4, 4, 4 };
    int m;                         //scaling factor of (rx, ry)
                                   //  for stage 0, 1, and 2
    if ( n == 0 ) m = 1;           //stage 0            
    else if ( n == 1 ) m = 2;      //stage 1
    else m = 4;                    //stage 2
    if ( p.x < 0 || p.y < 0 ){     //not valid position
      for ( int i = 0; i < 8; ++i )
        level[n][i].set(-1, -1);
      return;
    }
    //set the searching locations
    for ( int i = 0; i < 8; ++i ) {
      int locx = p.x - rx[i] / m;
      int locy = p.y - ry[i] / m;
      //skip the out-of-bound locations
      if (locx < 0||locx >= (width-16)||locy < 0||locy >= (height-16))
        level[n][i].set( -1, -1 );
      else
        level[n][i].set ( locx, locy );
    }
  }

  //Sum of absolute differences between macroblocks
  public int sad( YCbCr_MACRO p1, YCbCr_MACRO p2 )
  {
    int s = 0;
    for ( int i = 0; i < 256; ++i )
      s += Math.abs ( p1.Y[i] - p2.Y[i] );
    for ( int i = 0; i < 64; ++i ) {
      s += Math.abs ( p1.Cb[i] - p2.Cb[i] );
      s += Math.abs ( p1.Cr[i] - p2.Cr[i] );
    }

    return s;
  }

  /*
    Obtain an RGB macroblock at the position p of an RGB image     
      input: rgb_image, p
      ouput: rgb_macro
 */
  void getRGBmacro ( RGBImage rgb_image, Point2 p, RGB_MACRO rgb_macro)
  {
    int k, r;

    //points to specified RGB macroblock   
    k = p.y * width + p.x;         //width is the image width
    r = 0;
    for ( int i = 0; i < 16; ++i ) {
      for ( int j = 0; j < 16; ++j ){
        rgb_macro.rgb[r].R = rgb_image.ibuf[k].R;
        rgb_macro.rgb[r].G = rgb_image.ibuf[k].G;
        rgb_macro.rgb[r].B = rgb_image.ibuf[k].B;
        ++r; ++k;
      }
      k += ( width - 16 );         //points to next row within macroblock 
    }
  }

  //Search for the best mactch using TSS.  The motion vectors of the 'minimal point' for
  //  stages 0, 1, and 2 are returned in mvs]0], mvs[1], and mvs[2] respectively.
  void search ( RGBImage ref_frame, Point2 origin, YCbCr_MACRO yccm,  Vec2 mvs[] )
  {
    Point2 p = new Point2(), cp = new Point2();
    int d, dmin;
    int iwidth = ref_frame.width;
    int iheight = ref_frame.height;
    if ( iwidth != width || iheight != height ) {
      System.out.printf("\nsearch: Inconsistent image dimensions!\n");
      System.exit ( -1 );
    }
    
    RgbYcc rgbYcc = new RgbYcc();
    RGB_MACRO rgb_macro = new RGB_MACRO();
    YCbCr_MACRO yccm_ref = new YCbCr_MACRO();  //YCbCr Macro from reference frame
    for ( int i = 0; i < 3; ++i )
      mvs[i].set ( 0, 0 );         //Set alll motion vectors to (0, 0)
    cp.set ( origin );             //Set current point to origin
    for (int k = 0; k < 3; ++k) {  //Find minimal points at three levels
      p.set( cp );
      setLevel ( cp, k );          //Calculate all searching locations of stage k
      //Get an RGB macroblock from Reference frame
      getRGBmacro (ref_frame, cp, rgb_macro);
      rgbYcc.macroblock2ycbcr(rgb_macro,  yccm_ref);   //convert from RGB to YCbCr
    
      dmin = sad (yccm, yccm_ref);        //calculate SAD between the current and reference macroblocks
      for ( int i = 0; i < 8; ++i ) {     //test all valid search points
        if (level[k][i].x < 0 || level[k][i].y < 0)
          continue;                       //invalid search point
        getRGBmacro (ref_frame, level[k][i], rgb_macro);
        rgbYcc.macroblock2ycbcr (rgb_macro, yccm_ref); //convert from RGB to YCbCr
        d = sad (yccm, yccm_ref);
        if ( d < dmin ) {
          p.set ( level[k][i] );
          dmin = d;
        }
      }
      //p is the 'minimal' point
      if ( p.x == cp.x && p.y == cp.y )   //if origin is minimal point, we're done
        break;
      mvs[k].set ( p.minus( cp ) );       //motion vector is difference between minimal point and origin
      cp.set( p );                        //p becomes the new search center
    } //for k
  }
}
