#include "include/define.h"
#include "include/entropydec.h"
#include "include/huffman.h"

#include <iostream.h>

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

  //init transmission buffer variables
  use_parity=parity;
  instream.open(filename,ios::in);
  read_buffer();
  s_counter=0;
  frame_num=0;

  //setup VLC trees
  head_MBA=BuildHuffmanTree(MBA_CODE,35);
  head_MTYPE=BuildHuffmanTree(MTYPE_CODE,10);
  head_MVD=BuildHuffmanTree(MVD_CODE,32);
  head_CBP=BuildHuffmanTree(CBP_CODE,63);
  head_TCOEFF=BuildHuffmanTree(TCOEFF_CODE,28,15);

  //extract format from bit 29 of file
  format=((buffer[3] & (char)0x08)?_CIF_:_QCIF_);

  psc(); //subsequently done by mba

}

EntropyDecoder::~EntropyDecoder(void){
  //destructor
  instream.close();

  //kill huffman trees
  KillHuffmanTree(head_MBA);
  KillHuffmanTree(head_MTYPE);
  KillHuffmanTree(head_MVD);
  KillHuffmanTree(head_CBP);
  KillHuffmanTree(head_TCOEFF);

}

void EntropyDecoder::init_ptrs(int gq_vals[],int *mq_vals[],int pic_prop[],
			       int **mb_prop[],int ** (*mot_vec)[2]){
  int i;

  //pointers to data structures
  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];

}

int EntropyDecoder::get_frame(pel_t ** (*frame_data)[3]){

  int i;

  for(i=0;i<3;i++)
    data[i]=(*frame_data)[i];

  return picture_layer();

}

//***************************************************
//***************************************************
// PICTURE LAYER
//***************************************************
//***************************************************
int EntropyDecoder::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;

  remaining_psc();
  tr();
  ptype();
  disgard_pspare();
  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_];
      }
      if (i==0 && j==0) gbsc(); //mba does this subsequently
      if (gob_layer(tmp)==-1){
	for(i=0;i<2;i++)
	  delete tmp[i];
	return TRUE; //eof
      }
    }
  for(i=0;i<2;i++)
    delete tmp[i];
  
  return FALSE; //not done yet

}
  
int EntropyDecoder::psc(void){
  //first 16 bits of picture start code
  char tmpcode[17];
  int i;

  read_bits(tmpcode,16);

  //compare tmpcode to actual code
  for(i=0;i<16;i++)
    if(tmpcode[i]!=PIC_STARTCODE[i])
      return FALSE;
  return TRUE;

} 

int EntropyDecoder::remaining_psc(void){
  //rest of picture start code (after MBA has read part of it)
  char tmpcode[5];
  int i;
  
  read_bits(tmpcode,4);
  
  //compare tmpcode to actual code
  for(i=0;i<4;i++)
    if(tmpcode[i]!=PIC_STARTCODE[i+16])
      return FALSE;
  return TRUE;

} 
  
void EntropyDecoder::tr(void){
  //temporal reference

  int i,num;
  char tmpcode[6];

  num=0;
  read_bits(tmpcode,5);
  for(i=0;tmpcode[i]!=0;i++)
    num|=((int)((tmpcode[i]-48)<<4-i));
  frame_num=num;

}
  
void EntropyDecoder::ptype(void){
  //type information

  char tmpcode[7];
  int i;
  
  read_bits(tmpcode,6);
  for (i=0;i<6;i++)
    pic_properties[i]=tmpcode[i]-48;

}

void EntropyDecoder::disgard_pspare(void){
  char tmpcode[9];
  for(read_bits(tmpcode,1);tmpcode[0]=='1';read_bits(tmpcode,1))
    read_bits(tmpcode,8);
}  

//***************************************************
//***************************************************
// GROUP OF BLOCKS LAYER
//***************************************************
//***************************************************

int EntropyDecoder::gob_layer(pel_t **gob[3]){
  //group of block layer functions. returns -1 when eof, 0 ow
   
  int i,j,m,gb,mb;
  
  //init GOB
  prev_mb=0;
  prev_mb_mc=FALSE;

  //GOB loop
  gb=gn();
  gquant(gb);
  disgard_gspare();
  for(mb=mb_layer(gob,gb);GOB_START_CODE_IND!=mb;mb=mb_layer(gob,gb))
    if(mb==-1) return -1;

  return 0;

}

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

  char tmpcode[17];
  int i;

  read_bits(tmpcode,16);

  //compare tmpcode to actual code
  for(i=0;i<16;i++)
    if(tmpcode[i]!=MBA[34][i])
      return FALSE;
  return TRUE;

}

int EntropyDecoder::gn(void){
  //group number

  int i,num;
  char tmpcode[5];

  num=0;
  read_bits(tmpcode,4);
  for(i=0;tmpcode[i]!=0;i++)
    num|=((int)((tmpcode[i]-48)<<(3-i)));
  /*
  cout<<"gn:    "<<tmpcode<<" ";
  */
  return num;
  
}

void EntropyDecoder::gquant(int gnum){
  //quantizer information

  int i;
  char tmpcode[6];

  gquant_vals[gnum-1]=0;
  read_bits(tmpcode,5);
  for(i=0;tmpcode[i]!=0;i++)
    gquant_vals[gnum-1]|=((int)((tmpcode[i]-48)<<4-i));

}

void EntropyDecoder::disgard_gspare(void){
  char tmpcode[9];
  for(read_bits(tmpcode,1);tmpcode[0]=='1';read_bits(tmpcode,1))
    read_bits(tmpcode,8);
}  

//***************************************************
//***************************************************
// MACROBLOCK LAYER
//***************************************************
//***************************************************
int EntropyDecoder::mb_layer(pel_t **gob[3],int gnum){
  //maroblock layer functions

  pel_t *mb_Y[16],*mb_U[8],*mb_V[8];
  int mb,a,b,i,m;
  int p[6];
  pel_t **tmp;

  mb=mba();
  //  if (frame_num==1)
  //  cout<<mb<<" "<<flush;
  if (mb==GOB_START_CODE_IND || mb==MBA_STUFFING_IND || mb==-1) return mb;

  /*
  if (gnum==5 && mb==29 && frame_num==1){
    char bits[500];
    read_bits(bits,450);
    cout<<bits<<endl<<flush;
  }
  */

  for(m=0;m<16;m++)
    mb_Y[m]=&gob[_Y_][((mb-1)/11)*_MB_SIZE_+m][((mb-1)%11)*_MB_SIZE_];
  for(m=0;m<8;m++)
    mb_U[m]=&gob[_U_][((mb-1)/11)*_BLOCK_SIZE_+m][((mb-1)%11)*_BLOCK_SIZE_];
  for(m=0;m<8;m++)
    mb_V[m]=&gob[_V_][((mb-1)/11)*_BLOCK_SIZE_+m][((mb-1)%11)*_BLOCK_SIZE_];
  
  if(!mtype(gnum,mb))
    mvd(gnum,mb);
  else{
    if(mquant_vals[gnum-1][mb-1]!=-1)
      mquant(gnum,mb);
    if(mb_properties[gnum-1][mb-1][MC_PRESENT]){
      mvd(gnum,mb);
    }
    if(mb_properties[gnum-1][mb-1][MODE]==INTER){
      //cout<<"("<<gnum<<","<<mb<<")";
      cbp(p);
    }
    else
      for(i=0;i<6;i++)
	p[i]=1;
    tmp=new (pel_t *)[_DCT_SIZE_];  
    //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
    //get Cb
    if (p[4]==1)
      block_layer(mb_U);
    //get Cr
    if (p[5]==1)
      block_layer(mb_V);
    delete tmp;
  }

  prev_mb_mc=mb_properties[gnum-1][mb-1][MC_PRESENT];
  return prev_mb=mb;
}

int EntropyDecoder::mba(void){
  //macroblock address. returns -1 when eof reached
  int info[2];
  if(next_vlc_value(head_MBA,info)==1) return -1;
  if ((info[0]+1)==MBA_STUFFING_IND || (info[0]+1)==GOB_START_CODE_IND)
    return info[0]+1;
  else
    return info[0]+prev_mb+1;
}

int EntropyDecoder::mtype(int gnum,int mbnum){
  //reads type information,and returns whether
  //TCOEFF data is available
  
  
  int info[2];
  next_vlc_value(head_MTYPE,info);
  switch(info[0]){

  case INTRA:
    mb_properties[gnum-1][mbnum-1][MODE]=INTRA;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=FALSE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    return TRUE;
    break;
  case INTRA_MQUANT:
    mb_properties[gnum-1][mbnum-1][MODE]=INTRA;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=FALSE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    mquant_vals[gnum-1][mbnum-1]=0;//signal that we should read mquant
    return TRUE;
    break;
  case INTER_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=FALSE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    mquant_vals[gnum-1][mbnum-1]=0;//signal that we should read mquant
    return TRUE;
    break;
  case INTER_MQUANT_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=FALSE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    return TRUE;
    break;
  case INTER_MC_MVD:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    return FALSE;
    break;
  case INTER_MC_MVD_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    return TRUE;
    break;
  case INTER_MC_MQUANT_MVD_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=FALSE;
    mquant_vals[gnum-1][mbnum-1]=0;//signal that we should read mquant
    return TRUE;
    break;
  case INTER_MC_FIL_MVD:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=TRUE;
    return FALSE;
    break;
  case INTER_MC_FIL_MVD_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=TRUE;
    return TRUE;
    break;
  case INTER_MC_FIL_MQUANT_MVD_TCOEFF:
    mb_properties[gnum-1][mbnum-1][MODE]=INTER;
    mb_properties[gnum-1][mbnum-1][MC_PRESENT]=TRUE;
    mb_properties[gnum-1][mbnum-1][FIL_PRESENT]=TRUE;
    mquant_vals[gnum-1][mbnum-1]=0;//signal that we should read mquant
    return TRUE;
    break;
  }  
}

void EntropyDecoder::mquant(int gnum,int mbnum){
  //quantizer

  int i;
  char tmpcode[6];

  read_bits(tmpcode,5);
  for(i=0;tmpcode[i]!=0;i++)
    mquant_vals[gnum-1][mbnum-1]|=((int)((tmpcode[i]-48)<<4-i));

}

void EntropyDecoder::mvd(int gnum,int mbnum){
  //motion vector data

  int info[2],tmp_dx,tmp_dy;
  
  if (0==(mbnum-1)%_GOB_WIDTH_ || prev_mb_mc==FALSE || mbnum-prev_mb!=1){
    next_vlc_value(head_MVD,info);
    motion_vectors[_DX_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11]=info[0]-16;
    next_vlc_value(head_MVD,info);
    motion_vectors[_DY_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11]=info[0]-16;
  }
  else{
    next_vlc_value(head_MVD,info);
    tmp_dx=motion_vectors[_DX_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11-1]+info[0]-16;
    motion_vectors[_DX_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11]=
      ((tmp_dx>15)? tmp_dx-32: ((tmp_dx<-16)?tmp_dx+32:tmp_dx) );
    next_vlc_value(head_MVD,info);
    tmp_dy=motion_vectors[_DY_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11-1]+info[0]-16;
    motion_vectors[_DY_]
      [3*((gnum-1)/2)+(mbnum-1)/11][11*((gnum-1)%2)+(mbnum-1)%11]=
      ((tmp_dy>15)? tmp_dy-32: ((tmp_dy<-16)?tmp_dy+32:tmp_dy) );
  }
    
  //save current data for next call to mvd
  prev_mb=mbnum;
  prev_mb_mc=mb_properties[gnum-1][mbnum-1][MC_PRESENT];

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

  next_vlc_value(head_CBP,info);
  info[0]++;
  /*  if (frame_num==1)
      cout<<info[0]<<" "<<flush;*/
  for(i=0;i<6;i++)
    p[i]=((info[0]&((int)1<<(5-i)))?1:0);

}


//***************************************************
//***************************************************
// BLOCK LAYER
//***************************************************
//***************************************************
void EntropyDecoder::block_layer(pel_t *block[]){
  //block layer functions
  
  get_tcoeff(block);

}

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

  int i,j,k,run,level;
  int info[2]={0,0};//run,level
  char tmpcode[9];

  for(i=0;info[0]!=EOB;i++){
    if (i==0 && bit_peak()){
      //special case for first coefficient
      read_bits(tmpcode,2);
      if (tmpcode[1]=='0')
	tc[0][0]=1;
      else
	tc[0][0]=-1;
    }
    else{
      next_vlc_value(head_TCOEFF,info);
      if (info[0]<EOB){
	read_bits(tmpcode,1);//sign bit
	i+=info[0];//run
	tc[order[i][0]][order[i][1]]=(info[1]+1)*((tmpcode[0]=='0')?1:-1);//level
      }
      else if (info[0]==ESCAPE){
        //20-bit code: 6-bit ESC, 6-bit run, 8-bit level

	//RUN
	run=0;
	read_bits(tmpcode,6);
        for(j=0;tmpcode[j]!=0;j++)
	  run|=((int)((tmpcode[j]-48)<<5-j));
	i+=run;

	//LEVEL
	level=0;
	read_bits(tmpcode,8);
        for(j=0;tmpcode[j]!=0;j++)
	  level|=((int)((tmpcode[j]-48)<<7-j));
	if (i==0 && frame_num==0)
	  tc[0][0]=(unsigned char)level;//level of INTRA-DC
	else
	  tc[order[i][0]][order[i][1]]=(char)level;//level
	
      }//else if
    }//else
  }//for i
}
   
//***************************************************
//***************************************************
// TRANSMISSION BUFFER
//***************************************************
//***************************************************
int EntropyDecoder::bit_peak(void){
  //return next bit w/out incrementing pointer
  return ((buffer[BUFFER_FRAME_WIDTH/8-bit_counter/8-1]&
	  (char)(1<<(bit_counter%8))?1:0));
  
}
  
int EntropyDecoder::read_bits(char bits[],int num){
  //gets a series of bits from a buffer
  
  int i,done=FALSE;

  if (use_parity){
  }
    
  //read bits from each byte of buffer
  for(i=0;i<num && done==FALSE;i++){
    if (instream.eof() && BUFFER_FRAME_WIDTH/8-bit_counter/8-1>=eofbytes)
      done=TRUE;
    else{
      if(buffer[BUFFER_FRAME_WIDTH/8-bit_counter/8-1]&
	 (char)(1<<(bit_counter%8)))
	bits[i]='1';
      else
	bits[i]='0';
      if((use_parity && bit_counter==BUFFER_PARITY_WIDTH) ||
	 (!use_parity && bit_counter==0)){//continue reading bits
	if (use_parity)
	  check_parity();
	if (1==read_buffer())
	  eofbytes=instream.gcount();//eof reached
      }//if counter
      else bit_counter--;
    }

  }//for i
  
  bits[num]='\0';
  return done;
  
}

int EntropyDecoder::next_vlc_value(nodep head,int info[2]){
//returns values corresponding to next VLC in bitstream in "info"
  //returns 0 if everything is fine, 1 if eof reached
  
  nodep currentnode=head;
  char bit[2];
  while(currentnode->left!=NULL || currentnode->right!=NULL){
    if (1==read_bits(bit,1)) return 1;
    if(bit[0]=='1')
      if (currentnode->left!=NULL)
	currentnode=currentnode->left;
      else
	return 1;
    else
      if (currentnode->right!=NULL)
	currentnode=currentnode->right;
      else
	return 1;
  }
  info[0]=currentnode->info[0];
  info[1]=currentnode->info[1];

  return 0;
}


void EntropyDecoder::check_parity(){
  

}

int EntropyDecoder::read_buffer(){
  //reads enough bits from a file to fill 'buffer'
  //returns 0 if !eof, 1 ow

  //reset counter
  bit_counter=BUFFER_FRAME_WIDTH-1;

  //read in new buffer
  instream.read(buffer,BUFFER_FRAME_WIDTH/8);

  return instream.eof();
  
}