/************************************************************************
 *
 *  coder.c, main coding engine of tmn (TMN encoder)
 *  
 ************************************************************************/

/*****************************************************************
 *
 * Modified by Pat Mulroy, BT Labs to run syntax based arithmetic
 * coding.  SAC option, H.263 (Annex E).
 *
 *****************************************************************/

#include"sim.h"

/**********************************************************************
 *
 *      Name:           CodeOneOrTwo
 *	Description:	code one image normally or two images
 *                      as a PB-frame (CodeTwoPB and CodeOnePred merged)
 *	
 *      Input:          pointer to image, prev_image, prev_recon, Q
 *        
 *	Returns:	pointer to reconstructed image
 *	Side effects:	memory is allocated to recon image
 *
 ***********************************************************************/


void CodeOneOrTwo(PictImage *curr, PictImage *B_image, PictImage *prev, 
          PictImage *pr, int QP, int frameskip, Bits *bits,
          Pict *pic, PictImage *B_recon, PictImage *recon)
{
  unsigned char *prev_ipol,*pi_edge=NULL,*pi,*orig_lum;
  PictImage *prev_recon=NULL, *pr_edge=NULL;
  MotionVector *MV[6][MBR+1][MBC+2];
  MotionVector ZERO = {0,0,0,0,0};
  MB_Structure *recon_data_P; 
  MB_Structure *recon_data_B=NULL; 
  MB_Structure *diff; 
  int *qcoeff_P;
  int *qcoeff_B=NULL;

  int Mode,B;  
  int CBP, CBPB=0;
  int bquant[] = {5,6,7,8};
  int QP_B;
  int newgob;

  int i,j,k;

  /* buffer control vars */
  float QP_cumulative = (float)0.0;
  int abs_mb_num = 0, QuantChangePostponed = 0;
  int QP_new, QP_prev, dquant, QP_xmitted=QP;

  ZeroBits(bits);

  /* interpolate image */
  if (mv_outside_frame) {
    if (long_vectors) {
      /* If the Extended Motion Vector range is used, motion vectors
         may point further out of the picture than in the normal range,
         and the edge images will have to be made larger */
      B = 16;
    }
    else {
      /* normal range */
      B = 8;
    }
    pi_edge = (unsigned char *)malloc(sizeof(char)*(pels+4*B)*(lines+4*B));
    if (pi_edge == NULL) {
      fprintf(stderr,"Couldn't allocate memory for pi_edge\n");
      exit(-1);
    }
    MakeEdgeImage(pr->lum,pi_edge + (pels + 4*B)*2*B+2*B,pels,lines,2*B);
    pi = InterpolateImage(pi_edge, pels+4*B, lines+4*B);
    free(pi_edge);
    prev_ipol = pi + (2*pels + 8*B) * 4*B + 4*B; 
    
    /* luma of non_interpolated image */
    pr_edge = InitImage((pels+4*B)*(lines+4*B));
    MakeEdgeImage(prev->lum, pr_edge->lum + (pels + 4*B)*2*B+2*B,
          pels,lines,2*B);
    orig_lum = pr_edge->lum + (pels + 4*B)*2*B+2*B;
    
    /* non-interpolated image */
    MakeEdgeImage(pr->lum,pr_edge->lum + (pels+4*B)*2*B + 2*B,pels,lines,2*B);
    MakeEdgeImage(pr->Cr,pr_edge->Cr + (pels/2 + 2*B)*B + B,pels/2,lines/2,B);
    MakeEdgeImage(pr->Cb,pr_edge->Cb + (pels/2 + 2*B)*B + B,pels/2,lines/2,B);
    
    prev_recon = (PictImage *)malloc(sizeof(PictImage));
    prev_recon->lum = pr_edge->lum + (pels + 4*B)*2*B + 2*B;
    prev_recon->Cr = pr_edge->Cr + (pels/2 + 2*B)*B + B;
    prev_recon->Cb = pr_edge->Cb + (pels/2 + 2*B)*B + B;
  }
  else {
    pi = InterpolateImage(pr->lum,pels,lines);
    prev_ipol = pi;
    prev_recon = pr;
    orig_lum = prev->lum;
  }

  /* mark PMV's outside the frame */
  for (i = 1; i < (pels>>4)+1; i++) {
    for (k = 0; k < 6; k++) {
      MV[k][0][i] = (MotionVector *)malloc(sizeof(MotionVector));
      MarkVec(MV[k][0][i]);
    }
    MV[0][0][i]->Mode = MODE_INTRA;
  }
  /* zero out PMV's outside the frame */
  for (i = 0; i < (lines>>4)+1; i++) {
    for (k = 0; k < 6; k++) {
      MV[k][i][0] = (MotionVector *)malloc(sizeof(MotionVector));
      ZeroVec(MV[k][i][0]);
      MV[k][i][(pels>>4)+1] = (MotionVector *)malloc(sizeof(MotionVector));
      ZeroVec(MV[k][i][(pels>>4)+1]);
    }
    MV[0][i][0]->Mode = MODE_INTRA;
    MV[0][i][(pels>>4)+1]->Mode = MODE_INTRA;
  }

  /* Integer and half pel motion estimation */
  MotionEstimatePicture(curr->lum,prev_recon->lum,prev_ipol,
        pic->seek_dist,MV, pic->use_gobsync);
  /* note: integer pel motion estimation is now based on previous
     reconstructed image, not the previous original image. We have
     found that this works better for some sequences and not worse for
     others.  Note that it can not easily be changed back by
     substituting prev_recon->lum with orig_lum in the line above,
     because SAD for zero vector is not re-calculated in the half
     pel search. The half pel search has always been based on the
     previous reconstructed image */

#ifndef OFFLINE_RATE_CONTROL
  if (pic->bit_rate != 0) {
    /* Initialization routine for Rate Control */
    QP_new = InitializeQuantizer(PCT_INTER, (float)pic->bit_rate, 
               (pic->PB ? pic->target_frame_rate/2 : pic->target_frame_rate),
               pic->QP_mean);
    QP_xmitted = QP_prev = QP_new; 
  }
  else {
    QP_new = QP_xmitted = QP_prev = QP; /* Copy the passed value of QP */
  }
#else
    QP_new = QP_xmitted = QP_prev = QP; /* Copy the passed value of QP */
#endif
  dquant = 0; 

  for ( j = 0; j < lines/MB_SIZE; j++) {

#ifndef OFFLINE_RATE_CONTROL
    if (pic->bit_rate != 0) {
      /* QP updated at the beginning of each row */
      AddBitsPicture(bits);
      
      QP_new =  UpdateQuantizer(abs_mb_num, pic->QP_mean, PCT_INTER, 
           (float)pic->bit_rate, pels/MB_SIZE, lines/MB_SIZE, 
           bits->total);
    }
#endif

    newgob = 0;

    if (j == 0) {
      pic->QUANT = QP_new;
      bits->header += CountBitsPicture(pic);
      QP_xmitted = QP_prev = QP_new;
    }
    else if (pic->use_gobsync && j%pic->use_gobsync == 0) {
      bits->header += CountBitsSlice(j,QP_new); /* insert gob sync */
      QP_xmitted = QP_prev = QP_new;
      newgob = 1;
    }

    for ( i = 0; i < pels/MB_SIZE; i++) {

      /* Update of dquant, check and correct its limit */
      dquant = QP_new - QP_prev;
      if (dquant != 0 && i != 0 && MV[0][j+1][i+1]->Mode == MODE_INTER4V) {
        /* It is not possible to change the quantizer and at the same
           time use 8x8 vectors. Turning off 8x8 vectors is not
           possible at this stage because the previous macroblock
           encoded assumed this one should use 8x8 vectors. Therefore
           the change of quantizer is postponed until the first MB
           without 8x8 vectors */
        dquant = 0;
        QP_xmitted = QP_prev;
        QuantChangePostponed = 1;
      }
      else {
        QP_xmitted = QP_new;
        QuantChangePostponed = 0;
      }
      if (dquant > 2)  { dquant =  2; QP_xmitted = QP_prev + dquant;}
      if (dquant < -2) { dquant = -2; QP_xmitted = QP_prev + dquant;}

      pic->DQUANT = dquant;
      /* modify mode if dquant != 0 (e.g. MODE_INTER -> MODE_INTER_Q) */
      Mode = ModifyMode(MV[0][j+1][i+1]->Mode,pic->DQUANT);
      MV[0][j+1][i+1]->Mode = Mode;

      pic->MB = i + j * (pels/MB_SIZE);

      if (Mode == MODE_INTER || Mode == MODE_INTER_Q || Mode==MODE_INTER4V) {
        /* Predict P-MB */
        diff = Predict_P(curr,prev_recon,prev_ipol,
         i*MB_SIZE,j*MB_SIZE,MV,pic->PB);

      }
      else {
        diff = (MB_Structure *)malloc(sizeof(MB_Structure));
        FillLumBlock(i*MB_SIZE, j*MB_SIZE, curr, diff);
        FillChromBlock(i*MB_SIZE, j*MB_SIZE, curr, diff);
      }

      /* P or INTRA Macroblock */
      qcoeff_P = MB_Encode(diff, QP_xmitted, Mode);
      CBP = FindCBP(qcoeff_P, Mode, 64);
      if (CBP == 0 && (Mode == MODE_INTER || Mode == MODE_INTER_Q)) 
        ZeroMBlock(diff);
      else
        MB_Decode(qcoeff_P, diff, QP_xmitted, Mode);
      recon_data_P = MB_Recon_P(prev_recon, prev_ipol,diff, 
        i*MB_SIZE,j*MB_SIZE,MV,pic->PB);
      Clip(recon_data_P);
      free(diff);
        

      /* Predict B-MB using reconstructed P-MB and prev. recon. image */
      if (pic->PB) {
        diff = Predict_B(B_image, prev_recon, prev_ipol,i*MB_SIZE, j*MB_SIZE,
         MV, recon_data_P, frameskip, pic->TRB);
        if (QP_xmitted == 0)  
          QP_B = 0;  /* (QP = 0 means no quantization) */
        else 
          QP_B = mmax(1,mmin(31,bquant[pic->BQUANT]*QP_xmitted/4));
        qcoeff_B = MB_Encode(diff, QP_B, MODE_INTER);
        CBPB = FindCBP(qcoeff_B, MODE_INTER, 64);
        if (CBPB)
          MB_Decode(qcoeff_B, diff, QP_B, MODE_INTER);
        else
          ZeroMBlock(diff);
        recon_data_B = MB_Recon_B(prev_recon, diff,prev_ipol,i*MB_SIZE,
          j*MB_SIZE,MV,recon_data_P,frameskip,
          pic->TRB);
        Clip(recon_data_B);

        /* decide MODB */ 

        if (CBPB) {
          pic->MODB = PBMODE_CBPB_MVDB;
        }
        else {
          if (MV[5][j+1][i+1]->x == 0 && MV[5][j+1][i+1]->y == 0)
            pic->MODB = PBMODE_NORMAL;
          else
            pic->MODB = PBMODE_MVDB;
        }

        free(diff);

      }
      else
        ZeroVec(MV[5][j+1][i+1]); /* Zero out PB deltas */

      if ((CBP==0) && (CBPB==0) && (EqualVec(MV[0][j+1][i+1],&ZERO)) && 
          (EqualVec(MV[5][j+1][i+1],&ZERO)) &&
          (Mode == MODE_INTER || Mode == MODE_INTER_Q)) {
        /* Skipped MB : both CBP and CBPB are zero, 16x16 vector is zero,
           PB delta vector is zero and Mode = MODE_INTER */
        if (Mode == MODE_INTER_Q) {
          /* DQUANT != 0 but not coded anyway */
          QP_xmitted = QP_prev;
          pic->DQUANT = 0;
          Mode = MODE_INTER;
        }
        if (!syntax_arith_coding)
          CountBitsMB(Mode,1,CBP,CBPB,pic,bits);
        else
          Count_sac_BitsMB(Mode,1,CBP,CBPB,pic,bits);
      }
      else {
        /* Normal MB */
        if (!syntax_arith_coding) { /* VLC */
          CountBitsMB(Mode,0,CBP,CBPB,pic,bits);

          if (Mode == MODE_INTER  || Mode == MODE_INTER_Q) {
            bits->no_inter++;
            CountBitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          else if (Mode == MODE_INTER4V) {
            bits->no_inter4v++;
            CountBitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          else {
            /* MODE_INTRA or MODE_INTRA_Q */
            bits->no_intra++;
            if (pic->PB)
              CountBitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          
          if (CBP || Mode == MODE_INTRA || Mode == MODE_INTRA_Q)
            CountBitsCoeff(qcoeff_P, Mode, CBP, bits, 64);
          if (CBPB)
            CountBitsCoeff(qcoeff_B, MODE_INTER, CBPB, bits, 64);
        } /* end VLC */

        else { /* SAC */
          Count_sac_BitsMB(Mode,0,CBP,CBPB,pic,bits);
 
          if (Mode == MODE_INTER  || Mode == MODE_INTER_Q) {
            bits->no_inter++;
            Count_sac_BitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          else if (Mode == MODE_INTER4V) {
            bits->no_inter4v++;
            Count_sac_BitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          else {
            /* MODE_INTRA or MODE_INTRA_Q */
            bits->no_intra++;
            if (pic->PB)
              Count_sac_BitsVectors(MV, bits, i, j, Mode, newgob, pic);
          }
          
          if (CBP || Mode == MODE_INTRA || Mode == MODE_INTRA_Q)
            Count_sac_BitsCoeff(qcoeff_P, Mode, CBP, bits, 64);
          if (CBPB)
            Count_sac_BitsCoeff(qcoeff_B, MODE_INTER, CBPB, bits, 64);
        } /* end SAC */

        QP_prev = QP_xmitted;
      }

      abs_mb_num++;
      QP_cumulative += QP_xmitted;     
#ifdef PRINTQ 
      /* most useful when quantizer changes within a picture */
      if (QuantChangePostponed)
        fprintf(stdout,"@%2d",QP_xmitted);