#include "include/define.h"
#include "include/entropyenc.h"

EntropyEncoder::EntropyEncoder(const char filename[],int parity){
  //constructor

  int i;

  //init transmission buffer variables
  use_parity=parity;
  outstream.open(filename,ios::out);
  for(i=0;i<=BUFFER_FRAME_WIDTH/8;i++)
    buffer[i]=0;
  bit_counter=BUFFER_FRAME_WIDTH-1;
  s_counter=0;
  frame_num=0;

  num_pspare=0;
  num_gspare=0;

}

EntropyEncoder::~EntropyEncoder(void){
  //destructor
  flush_frames();
  outstream.close();

}


void EntropyEncoder::send_frame(pel_t ** (*frame_data)[3],int gq_vals[],
				int *mq_vals[],int pic_prop[],
				int **mb_prop[],int ** (*mot_vec)[2]){
  int i;

  for(i=0;i<3;i++)
    data[i]=(*frame_data)[i];
  gquant_vals=gq_vals;
  mquant_vals=mq_vals;
  pic_properties=pic_prop;
  mb_properties=mb_prop;
  for(i=0;i<2;i++)
    motion_vectors[i]=(*mot_vec)[i];

  picture_layer();

}

//***************************************************
//***************************************************
// PICTURE LAYER
//***************************************************
//***************************************************

void EntropyEncoder::picture_layer(void){
  //picture layer functions
  
  pel_t **tmp[3];
  tmp[_Y_]=new (pel_t *)[16*3];
  tmp[_U_]=new (pel_t *)[8*3];
  tmp[_V_]=new (pel_t *)[8*3];
  int i,j,m;

  psc();
  tr(frame_num==32?frame_num=0:frame_num++);
  ptype(pic_properties[SPLIT_SCREEN],pic_properties[DOC_CAM],
	pic_properties[FREEZE_PIC],pic_properties[FORMAT],
	pic_properties[STILL_IMAGE],pic_properties[SPARE]);
  for(i=0;i<num_pspare;i++){
    pei(TRUE);
    pspare(i);
  }
  pei(FALSE);
  for(i=0;i<(pic_properties[FORMAT]==_CIF_?6:3);i++)
    for(j=0;j<(pic_properties[FORMAT]==_CIF_?2:1);j++){
      for(m=0;m<16*3;m++)
	tmp[_Y_][m]=&data[_Y_][i*3*_MB_SIZE_+m][j*11*_MB_SIZE_];
      for(m=0;m<8*3;m++){
	tmp[_U_][m]=&data[_U_][i*3*_BLOCK_SIZE_+m][j*11*_BLOCK_SIZE_];
	tmp[_V_][m]=&data[_V_][i*3*_BLOCK_SIZE_+m][j*11*_BLOCK_SIZE_];
      }
      gob_layer(tmp,i*2+j+1);
    }
  for(i=0;i<2;i++)
    delete tmp[i];

}
  
void EntropyEncoder::psc(void){
  //picture start code
  buffer_bits(PIC_STARTCODE);
} 
  
void EntropyEncoder::tr(int num){
  //temporal reference

  char code[6];
  int i;

  //turn bits into string
  for(i=0;i<5;i++)
    code[i]=((num&(0x10>>i))?'1':'0');
  code[5]='\0';

  buffer_bits(code);
 
}
  
void EntropyEncoder::ptype(int split_screen,int doc_camera,int freeze_pic,
           int format,int still_image,int spare){
  //type information

  char code[7];

  if (split_screen==ON) code[0]='1';
  else code[0]='0';
  if (doc_camera==ON)   code[1]='1';
  else code[1]='0';
  if (freeze_pic==ON)   code[2]='1';
  else code[2]='0';
  if (format==_CIF_)    code[3]='1';
  else code[3]='0';
  if (still_image==ON)  code[4]='0';
  else code[4]='1';
  if (spare==1)         code[5]='1';
  else code[5]='0';

  code[6]='\0';

  buffer_bits(code);

}

void EntropyEncoder::pei(int sparebits){
  //Extra insertion information

  if (sparebits==FALSE)
    buffer_bits("0");  
  else
    buffer_bits("1");  

}

void EntropyEncoder::pspare(int which){
  //spare information

  char spare_bits[9];

  switch (which){
    case 0://spare bits 0
           break;
  }

  buffer_bits(spare_bits);  

}


//***************************************************
//***************************************************
// GROUP OF BLOCKS LAYER
//***************************************************
//***************************************************
 
void EntropyEncoder::gob_layer(pel_t **gob[3],int num){
  //group of block layer functions
   
  int i,j,m;
  pel_t *tmp_Y[16],*tmp_U[8],*tmp_V[8];
  
  //init GOB
  prev_dx=0;
  prev_dy=0;
  prev_mb=0;
  prev_mb_mc=FALSE;

  //GOB loop
  gbsc();
  gn(num);
  gquant(gquant_vals[num-1]);
  for(i=0;i<num_pspare;i++){
    gei(TRUE);
    gspare(i);
  }
  gei(FALSE);
  for(i=0;i<3;i++){
    for(j=0;j<11;j++){
      for(m=0;m<16;m++)
        tmp_Y[m]=&gob[_Y_][i*_MB_SIZE_+m][j*_MB_SIZE_];
      for(m=0;m<8;m++)
        tmp_U[m]=&gob[_U_][i*_BLOCK_SIZE_+m][j*_BLOCK_SIZE_];
      for(m=0;m<8;m++)
        tmp_V[m]=&gob[_V_][i*_BLOCK_SIZE_+m][j*_BLOCK_SIZE_];
      mb_layer(tmp_Y,tmp_U,tmp_V,
	       mb_properties[num-1][i*11+j][MODE],
               mb_properties[num-1][i*11+j][MC_PRESENT],
	       mb_properties[num-1][i*11+j][FIL_PRESENT],
               mquant_vals[num-1][i*11+j],
	       motion_vectors[_DX_][3*((num-1)/2)+i][11*((num-1)%2)+j],
               motion_vectors[_DY_][3*((num-1)/2)+i][11*((num-1)%2)+j],
	       i*11+j+1);
    }
  }

}

void EntropyEncoder::gbsc(void){
  //group of block start code

  buffer_bits(MBA[34]);

}

void EntropyEncoder::gn(int num){
  //group number

  char code[5];
  int i;

  //turn bits into string
  for(i=0;i<4;i++)
    code[i]=((num&(0x08>>i))?'1':'0');
  code[4]='\0';

  //  cout<<"gn:     "<<code<<" ";

  buffer_bits(code);
  
}

void EntropyEncoder::gquant(int stepsize){
  //quantizer information

  char code[6];
  int i;

  //turn bits into string
  for(i=0;i<5;i++)
    code[i]=((stepsize&(0x10>>i))?'1':'0');
  code[5]='\0';

  buffer_bits(code);
  
}

void EntropyEncoder::gei(int sparebits){
  //extra insertion information

  if (sparebits==FALSE)
    buffer_bits("0");  
  else
    buffer_bits("1");  
  
}

void EntropyEncoder::gspare(int which){
  //spare information

  char spare_bits[9];

  switch (which){
    case 0://spare bits 0
           break;
  }

  buffer_bits(spare_bits);  
  
}


//***************************************************
//***************************************************
// MACROBLOCK LAYER
//***************************************************
//***************************************************
   
void EntropyEncoder::mb_layer(pel_t *mb_Y[],pel_t *mb_U[],pel_t *mb_V[],
              int mode,int mc_present,int fil_present,
              int q_stepsize,int mv_dx,int mv_dy,int mb){
  //maroblock layer functions
  //
  // mb_Y: 16x16 luminance block
  // mb_U: 8x8 Cb block
  // mb_V: 8x8 Cr block
  // mode: {0(Intra),1(Inter)}

  int p[6]; //p[i]: whether to transmit block i
  int i,m,n,a,b,tcoeff_present;
  pel_t **tmp;

  //***********************************
  //figure out which blocks to transmit
  //***********************************
  //loop through luminance blocks
  tcoeff_present=FALSE;
  if (mode==INTER){
    for(a=0;a<2;a++)
      for(b=0;b<2;b++){
        //loop through coefficients
        p[a*2+b]=0;
        for(i=0;i<64 && p[a*2+b]==0;i++)
          if (0!=(int)mb_Y[a*_BLOCK_SIZE_+order[i][0]]
                          [b*_BLOCK_SIZE_+order[i][1]]){
            tcoeff_present=TRUE;
            p[a*2+b]=1;
          }
      }//for b
    //look at Cb
    p[4]=0;
    for(i=0;i<64 && p[4]==0;i++)
      if (0!=(int)mb_U[order[i][0]][order[i][1]]){
        tcoeff_present=TRUE;
        p[4]=1;
      }
    //look at Cr
    p[5]=0;
    for(i=0;i<64 && p[5]==0;i++)
      if (0!=(int)mb_V[order[i][0]][order[i][1]]){
        tcoeff_present=TRUE;
        p[5]=1;
      }
  }
  else{ //Intra (transmit all blocks)
    tcoeff_present=TRUE;
    for(i=0;i<6;i++)
      p[i]=1;
  }//else
  //***********************************

  //check whether we should send anything
  if(tcoeff_present || mc_present){
    tmp=new (pel_t *)[_DCT_SIZE_];  

    mba(mb);
    mtype(mode,mc_present,fil_present,
          (q_stepsize==-1?FALSE:TRUE),tcoeff_present);
    if (!tcoeff_present)
      mvd(mv_dx,mv_dy,mb,mc_present);
    else{
      if (q_stepsize!=-1)
        mquant(q_stepsize);
      mvd(mv_dx,mv_dy,mb,mc_present);
      if (mode==INTER){
        cbp(p);
      }
      //loop through luminance blocks
      for(a=0;a<2;a++)
	for(b=0;b<2;b++)
	  if (p[a*2+b]==1){
	    for(i=0;i<_DCT_SIZE_;i++)  
	      tmp[i] = &mb_Y[a*_BLOCK_SIZE_+i][b*_BLOCK_SIZE_];
	    block_layer(tmp);
	  }//if p
      //send Cb
      if (p[4]==1)
        block_layer(mb_U);
      //send Cr
      if (p[5]==1)
        block_layer(mb_V);
    }//else
    delete tmp;
  }
  
}

void EntropyEncoder::mba(int mb){
  //macroblock address

  buffer_bits(MBA[mb-prev_mb-1]);

}

void EntropyEncoder::mtype(int pred,int mc_present,int fil_present,
           int mquant_present,int tcoeff_present){
  //type information
  //pred = {0,1}, 0: Intra, 1: Inter
  //all other values are {FALSE,TRUE}

  int ind;

  //get index of VLC
  if (pred==0){ //Intra
    if (mquant_present)
      ind=INTRA_MQUANT;
    else
      ind=INTRA;
  }else{ //Inter
    if (mc_present){
      if (fil_present){
        if(tcoeff_present){
          if(mquant_present)
            ind=INTER_MC_FIL_MQUANT_MVD_TCOEFF;
          else
            ind=INTER_MC_FIL_MVD_TCOEFF;
        }else// tcoeff
           ind=INTER_MC_FIL_MVD;
      }else{ //fil
        if(tcoeff_present){
          if(mquant_present)
            ind=INTER_MC_MQUANT_MVD_TCOEFF;
          else
            ind=INTER_MC_MVD_TCOEFF;
        }else// tcoeff
           ind=INTER_MC_MVD;
      }//else fil
    }else{ //mc
      if (mquant_present)
        ind=INTER_MQUANT_TCOEFF;
      else      
        ind=INTER_TCOEFF;
    }//else mc
  }// else Inter

  buffer_bits(MTYPE[ind]);
  
}

void EntropyEncoder::mquant(int stepsize){
  //quantizer

  char code[6];
  int i;

  //turn bits into string
  for(i=0;i<5;i++)
    code[i]=((stepsize&(0x10>>i))?'1':'0');
  code[5]='\0';

  buffer_bits(code);

}

void EntropyEncoder::mvd(int dx,int dy,int mb,int mb_mc){
  //motion vector data
  // dx: horizontal motion vector component
  // dy: vertical motion vector component
  // mb: current macroblock
  // mb_mc: {FALSE,TRUE} tells whether current mb is motion compensated

  int tmp_dx=dx;
  int tmp_dy=dy;

  //calculate MV differences
  if (0==(mb-1)%_GOB_WIDTH_ || prev_mb_mc==FALSE || mb-prev_mb!=1)
    prev_dx=prev_dy=0;
  dx-=prev_dx;
  dy-=prev_dy;

  //save current data for next call to mvd
  prev_dx=tmp_dx;
  prev_dy=tmp_dy;
  prev_mb=mb;
  prev_mb_mc=mb_mc;

  //set dx and dy to correct corresponding index
  if (mb_mc){
    dx=(dx>15? dx-32: (dx<-16?dx+32:dx) );
    dy=(dy>15? dy-32: (dy<-16?dy+32:dy) );
  
    buffer_bits(MVD[dx+16]);//horizontal component;
    buffer_bits(MVD[dy+16]);//vertical component;
  }

}
  
void EntropyEncoder::cbp(int p[6]){
  //coded block pattern (p[i] = {0,1} to signify whether
  //block i is being transmitted)

  int ind=0;
  ind=(p[0]<<5)|(p[1]<<4)|(p[2]<<3)|(p[3]<<2)|(p[4]<<1)|p[5];
  /*
  if (frame_num==2)
    cout<<ind<<" "<<flush;
  */
  buffer_bits(CBP[ind-1]);
}

//***************************************************
//***************************************************
// BLOCK LAYER
//***************************************************
//***************************************************
   
void EntropyEncoder::block_layer(pel_t *block[]){
  //block layer functions
  
  tcoeff(block);
  send_eob();

}

void EntropyEncoder::tcoeff(pel_t *tc[]){
  //transform coefficients for an 8x8 block

  int i,j,k,level,run=0;
  char tmpcode[15];
  tmpcode[14]='\0';
  /*
  if (frame_num==2)
    cout<<" [[[";
  */
  for(i=0;i<64;i++){
    level=(int)tc[order[i][0]][order[i][1]];
    if(level!=0){
      /*
      if (frame_num==2){
	if(run!=0)
	  cout<<"_"<<run<<"_";
	cout<<level<<" ";
      }
      */
      //get coded string
      if (i==0 && abs(level)==1){
        //special case for first coefficient
        tmpcode[0]='1';tmpcode[1]=sign(level)+48;tmpcode[2]='\0';
        buffer_bits(tmpcode);  
      }
      else if(run<=26 && abs(level)<=15 && TCOEFF[run][abs(level)-1][0]!='\0'){
        //VLC
        buffer_bits(TCOEFF[run][abs(level)-1]);
        tmpcode[0]=sign(level)+48;tmpcode[1]='\0';
        buffer_bits(tmpcode);  
      }
      else{
        //20-bit code: 6-bit ESC, 6-bit run, 8-bit level

        //ESC
        buffer_bits(TCOEFF[27][0]);  
        //RUN
        for(j=1<<5,k=0;k<6;j>>=1,k++)
          tmpcode[k]=(((char)run & j)?'1':'0');
        //LEVEL
        for(j=1<<7,k=0;k<8;j>>=1,k++)
          tmpcode[6+k]=(((char)level & j)?'1':'0');
        buffer_bits(tmpcode);  
      }//else
      run=0;
    }//if level
    else
      run++;  
  }//for i
  /*
  if (frame_num==2)
    cout<<"]]] ";
  */

}
  
void EntropyEncoder::send_eob(void){
  //end of block
  buffer_bits(TCOEFF[EOB][1]);  
}

//***************************************************
//***************************************************
// TRANSMISSION BUFFER
//***************************************************
//***************************************************
   
void EntropyEncoder::buffer_bits(const char bits[]){
  //sends a series of bits to a buffer
  
  int i;

  if (use_parity){
    if (bit_counter==BUFFER_FRAME_WIDTH-1){
      //buffer frame fill bits
      if(BUFFER_S[s_counter++]=='1')
        buffer[BUFFER_FRAME_WIDTH/8-(bit_counter--)/8-1]=(unsigned char)0x80;
      if (s_counter==8) s_counter=0;
    }
  
    if (bit_counter==BUFFER_FRAME_WIDTH-2){
      if(bits[0]!='\0')
        buffer[BUFFER_FRAME_WIDTH/8-(bit_counter--)/8-1]|=(unsigned char)0x40;
    }
  }
    
  //write bits to each byte of buffer
  for(i=0;bits[i]!='\0';i++){

    if(bits[i]=='1')
      buffer[BUFFER_FRAME_WIDTH/8-bit_counter/8-1]|=
	(unsigned char)(1<<(bit_counter%8));
    if((use_parity && bit_counter==BUFFER_PARITY_WIDTH) ||
       (!use_parity && bit_counter==0)){//write encoding to file
      if (use_parity)
        calc_parity();
      flush_buffer();
    }//if counter
    else bit_counter--;

  }//for i

}

void EntropyEncoder::flush_frames(){
  //flush remaining frames filled w/ 1's and close file.  Should only
  //be called when all frames have been encoded and sent.

  int i;

  if(bit_counter<BUFFER_FRAME_WIDTH-1){

    //only flush as many bytes as we need to
    buffer_bits(MBA[MBA_STUFFING_IND-1]);
    outstream.write(buffer,BUFFER_FRAME_WIDTH/8-bit_counter/8-1);

    //fill rest of buffer with MBA stuffing
    //for(i=bit_counter/MBA_STUFFING_LEN;i>=0;i--)
    //  buffer_bits(MBA[MBA_STUFFING_IND-1]);
  }

  if (use_parity) //send rest of 8 frames
    while(s_counter<8){
      //fill bytes w/ 1's
      for(i=0;i<(BUFFER_FRAME_WIDTH-BUFFER_PARITY_WIDTH)/8-1;i++)
        buffer[i]|=(unsigned char)0xFF;
      if (use_parity)
        calc_parity();
      flush_buffer();
      s_counter++;
    }
  
}

void EntropyEncoder::calc_parity(){
  
  //parity bits for a frame of all 1's
  buffer[(BUFFER_FRAME_WIDTH-BUFFER_PARITY_WIDTH)/8]=0xFD;
  buffer[(BUFFER_FRAME_WIDTH-BUFFER_PARITY_WIDTH)/8+1]=0xB5;
  buffer[(BUFFER_FRAME_WIDTH-BUFFER_PARITY_WIDTH)/8+2]=0x1B;

}

void EntropyEncoder::flush_buffer(){

  //flushes the contents of the buffer
  outstream.write(buffer,64);

  //set buffer to all zeros
  for(int i=0;i<BUFFER_FRAME_WIDTH/8;i++)
    buffer[i]=0;

  //reset counter
  bit_counter=BUFFER_FRAME_WIDTH-1;
  
}