// *************************************************************************************
// *************************************************************************************
// Cabac.c	 CABAC entropy coding routines 
//
// Main contributors (see contributors.h for copyright, address and affiliation details)
//
// Detlev Marpe                    <marpe@hhi.de>
// *************************************************************************************
// *************************************************************************************
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <memory.h>
#include "cabac.h"
#include "global.h"



/************************************************************************
*
*  Name :       create_contexts_MotionInfo()
*
*  Description: Allocates of contexts models for the motion info
*								used for arithmetic encoding 
*
************************************************************************/
MotionInfoContexts* create_contexts_MotionInfo(void)
{
		
	int j;
	MotionInfoContexts *enco_ctx;	
	

  enco_ctx = (MotionInfoContexts*) calloc(1, sizeof(MotionInfoContexts) );
	if( enco_ctx == NULL )
		no_mem_exit(1);

	for (j=0; j<2; j++)
	{
		enco_ctx->mb_type_contexts[j] = (BiContextTypePtr) malloc(NUM_MB_TYPE_CTX  * sizeof( BiContextType ) );
				
		if( enco_ctx->mb_type_contexts[j] == NULL ) 
			no_mem_exit(1);
		
		enco_ctx->mv_res_contexts[j] = (BiContextTypePtr) malloc(NUM_MV_RES_CTX  * sizeof( BiContextType ) );
				
		if( enco_ctx->mv_res_contexts[j] == NULL ) 
			no_mem_exit(1);		
	}

	enco_ctx->ref_no_contexts = (BiContextTypePtr) malloc(NUM_REF_NO_CTX * sizeof( BiContextType ) );
				
	if( enco_ctx->ref_no_contexts == NULL ) 
		no_mem_exit(1);
						
	return enco_ctx;
}  


/************************************************************************
*
*  Name :       create_contexts_TextureInfo()
*
*  Description: Allocates of contexts models for the texture info
*								used for arithmetic encoding
*
************************************************************************/
TextureInfoContexts* create_contexts_TextureInfo(void)
{
		
	int j,k;
	TextureInfoContexts *enco_ctx;	

	enco_ctx = (TextureInfoContexts*) calloc(1, sizeof(TextureInfoContexts) );
	if( enco_ctx == NULL )
		no_mem_exit(1);

	for (j=0; j < 6; j++)
	{
		
		enco_ctx->ipr_contexts[j] = (BiContextTypePtr) malloc(NUM_IPR_CTX  * sizeof( BiContextType ) );
		
		if( enco_ctx->ipr_contexts[j] == NULL ) 
			no_mem_exit(1);
		
	}

	for (k=0; k<2; k++)
		for (j=0; j<3; j++)
		{
			
			enco_ctx->cbp_contexts[k][j] = (BiContextTypePtr) malloc(NUM_CBP_CTX  * sizeof( BiContextType ) );
			
			if( enco_ctx->cbp_contexts[k][j] == NULL ) 
				no_mem_exit(1);
		}

		
	for (j=0; j < NUM_TRANS_TYPE; j++)
	{
		
		enco_ctx->level_context[j] = (BiContextTypePtr) malloc(NUM_LEVEL_CTX  * sizeof( BiContextType ) );
		
		if( enco_ctx->level_context[j] == NULL ) 
			no_mem_exit(1);
		
		
		enco_ctx->run_context[j] = (BiContextTypePtr) malloc(NUM_RUN_CTX  * sizeof( BiContextType ) );
		
		if( enco_ctx->run_context[j] == NULL ) 
			no_mem_exit(1);
		
	}				
	return enco_ctx;
}  


/************************************************************************
*
*  Name :       init_contexts_MotionInfo()
*
*  Description: Initializes an array of contexts models with some pre-defined 
*								counts (ini_flag = 1) or with a flat histogram (ini_flag = 0)
*
************************************************************************/
void init_contexts_MotionInfo(MotionInfoContexts *enco_ctx, int ini_flag)
{
		
	int i,j;
	int scale_factor;

	if ( (img->width*img->height) <=  (IMG_WIDTH * IMG_HEIGHT) ) /*  format <= QCIF */                        
   scale_factor=1;
  else                                      
   scale_factor=2; 


	for (j=0; j<2; j++)
	{		
		if (ini_flag)
		{
			for (i=0; i < NUM_MB_TYPE_CTX; i++)
				biari_init_context(enco_ctx->mb_type_contexts[j] + i,MB_TYPE_Ini[j][i][0]*scale_factor,MB_TYPE_Ini[j][i][1]*scale_factor,MB_TYPE_Ini[j][i][2]*scale_factor);
		}
		else
		{
			for (i=0; i < NUM_MB_TYPE_CTX; i++)
				biari_init_context(enco_ctx->mb_type_contexts[j] + i,1,1,1000);
		}
				
		if (ini_flag)
		{
			for (i=0; i < NUM_MV_RES_CTX; i++)
				biari_init_context(enco_ctx->mv_res_contexts[j] + i,MV_RES_Ini[j][i][0]*scale_factor,MV_RES_Ini[j][i][1]*scale_factor,MV_RES_Ini[j][i][2]*scale_factor);
		}
		else
		{
			for (i=0; i < NUM_MV_RES_CTX; i++)
				biari_init_context(enco_ctx->mv_res_contexts[j] + i,1,1,1000);
		}
	}
								
	if (ini_flag)
	{		
		for (i=0; i < NUM_REF_NO_CTX; i++)
			biari_init_context(enco_ctx->ref_no_contexts + i,REF_NO_Ini[i][0]*scale_factor,REF_NO_Ini[i][1]*scale_factor,REF_NO_Ini[i][2]*scale_factor);
	}
	else
	{
		for (i=0; i < NUM_REF_NO_CTX; i++)
			biari_init_context(enco_ctx->ref_no_contexts + i,1,1,1000);		
	}
}  

/************************************************************************
*
*  Name :       init_contexts_TextureInfo()
*
*  Description: Initializes an array of contexts models with some pre-defined 
*								counts (ini_flag = 1) or with a flat histogram (ini_flag = 0)
*
************************************************************************/
void init_contexts_TextureInfo(TextureInfoContexts *enco_ctx, int ini_flag)
{
		
	int i,j,k;
	int scale_factor;
	 
	if ( (img->width*img->height) <=  (IMG_WIDTH * IMG_HEIGHT) ) /*  format <= QCIF */                        
   scale_factor=1;
  else                                      
   scale_factor=2; 

	for (j=0; j < 6; j++)
	{
		if (ini_flag)
		{
			for (i=0; i < NUM_IPR_CTX; i++)
				biari_init_context(enco_ctx->ipr_contexts[j] + i,IPR_Ini[j][i][0]*scale_factor,IPR_Ini[j][i][1]*scale_factor,IPR_Ini[j][i][2]*scale_factor);
		}
		else
		{
			for (i=0; i < NUM_IPR_CTX; i++)
				biari_init_context(enco_ctx->ipr_contexts[j] + i,2,1,50);
		}
	}

	for (k=0; k<2; k++)
		for (j=0; j<3; j++)
		{	
			if (ini_flag)
			{
				for (i=0; i < NUM_CBP_CTX; i++) 
					biari_init_context(enco_ctx->cbp_contexts[k][j] + i,CBP_Ini[k][j][i][0]*scale_factor,CBP_Ini[k][j][i][1]*scale_factor,CBP_Ini[k][j][i][2]*scale_factor);
			}
			else
			{
				for (i=0; i < NUM_CBP_CTX; i++) 
					biari_init_context(enco_ctx->cbp_contexts[k][j] + i,1,1,100);
			}
		}

		
	for (j=0; j < NUM_TRANS_TYPE; j++)
	{
				
		if (ini_flag)
		{
			for (i=0; i < NUM_LEVEL_CTX; i++)
				biari_init_context(enco_ctx->level_context[j] + i,Level_Ini[j][i][0]*scale_factor,Level_Ini[j][i][1]*scale_factor,Level_Ini[j][i][2]*scale_factor);
		}
		else
		{
			for (i=0; i < NUM_LEVEL_CTX; i++)
				biari_init_context(enco_ctx->level_context[j] + i,2,1,100);
		}
				
		if (ini_flag)
		{
			for (i=0; i < NUM_RUN_CTX; i++)
				biari_init_context(enco_ctx->run_context[j] + i,Run_Ini[j][i][0]*scale_factor,Run_Ini[j][i][1]*scale_factor,Run_Ini[j][i][2]*scale_factor);
		}
		else
		{
			for (i=0; i < NUM_RUN_CTX; i++)
				biari_init_context(enco_ctx->run_context[j] + i,2,1,100);
		}
	}				
}  


/************************************************************************
*
*  Name :       delete_contexts_MotionInfo()
*
*  Description: Frees the memory of the contexts models 
*								used for arithmetic encoding of the motion info.
*
************************************************************************/
void delete_contexts_MotionInfo(MotionInfoContexts *enco_ctx)
{
	
	int j;

	if( enco_ctx == NULL ) 
		return;

	for (j=0; j<2; j++)
	{
		if (enco_ctx->mb_type_contexts[j] != NULL)
			free(enco_ctx->mb_type_contexts[j] );
	
		if (enco_ctx->mv_res_contexts[j]  != NULL)
			free(enco_ctx->mv_res_contexts[j] );
	}

	if (enco_ctx->ref_no_contexts != NULL)
		free(enco_ctx->ref_no_contexts);
	
	free( enco_ctx );
	
	return;
	
}  

/************************************************************************
*
*  Name :       delete_contexts_TextureInfo()
*
*  Description: Frees the memory of the contexts models 
*								used for arithmetic encoding of the texture info.
*
************************************************************************/
void delete_contexts_TextureInfo(TextureInfoContexts *enco_ctx)
{
		
	int j,k;

	if( enco_ctx == NULL ) 
		return;

	for (j=0; j < 6; j++)
	{
		if (enco_ctx->ipr_contexts[j] != NULL)
			free(enco_ctx->ipr_contexts[j]);
	}

	for (k=0; k<2; k++)
		for (j=0; j<3; j++)
		{
			if (enco_ctx->cbp_contexts[k][j] != NULL)
				free(enco_ctx->cbp_contexts[k][j]);
		}

	for (j=0; j < NUM_TRANS_TYPE; j++)
	{
		if (enco_ctx->level_context[j] != NULL)
			free(enco_ctx->level_context[j]);
	
		if (enco_ctx->run_context[j] != NULL)
			free(enco_ctx->run_context[j]);
	}
	free( enco_ctx );
	
	return;
	
}  

/************************************************************************
*
*  Routine      writeSyntaxElement_CABAC
*
*  Description: generates arithmetic code and passes the code to the buffer
*
************************************************************************/
int	writeSyntaxElement_CABAC(SyntaxElement *se, DataPartition *this_dataPart)
{
	int curr_len;
	EncodingEnvironmentPtr eep_dp = &(this_dataPart->ee_cabac);

	curr_len = arienco_bits_written(eep_dp);

	/* perform the actual coding by calling the appropriate method */
	se->writing(se, eep_dp);

	return (se->len = (arienco_bits_written(eep_dp) - curr_len));
}

/****************************************************************************
*
* Function:        writeMB_typeInfo2Buffer_CABAC()
*
* Purpose:         This function is used to arithmetically encode the macroblock
*                  type info of a given MB. 
*
*
***************************************************************************/

void writeMB_typeInfo2Buffer_CABAC(SyntaxElement *se, EncodingEnvironmentPtr eep_dp)
{
	int l;
	int a, b;
	int act_ctx;
	int act_sym;
	int log_sym;
	int mode_sym=0;
	int mask;
	int mode16x16;

	MotionInfoContexts *ctx = (img->currentSlice)->mot_ctx;
	Macroblock *currMB = &img->mb_data[img->current_mb_nr];
	int curr_mb_type = se->value1;
	
	if(img->type == INTRA_IMG)  // INTRA-frame 
	{			
		if (currMB->mb_available[0][1] == NULL)
			b = 0;
		else 
			b = (( (currMB->mb_available[0][1])->mb_type != 0) ? 1 : 0 );
		if (currMB->mb_available[1][0] == NULL)
			a = 0;
		else 
			a = (( (currMB->mb_available[1][0])->mb_type != 0) ? 1 : 0 );
		
		act_ctx = a + 2*b;
		act_sym = curr_mb_type;
		se->context = act_ctx; // store context
				
		if (act_sym==0)	/* 4x4 Intra */	
			biari_encode_symbol(eep_dp, 0, ctx->mb_type_contexts[0] + act_ctx );
		else /* 16x16 Intra */
		{
			biari_encode_symbol(eep_dp, 1, ctx->mb_type_contexts[0] + act_ctx );
			mode_sym=act_sym-1; /* Values in the range of 0...23 */
			act_ctx = 4;
			act_sym = mode_sym/12;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[0] + act_ctx ); /* coding of AC/no AC */
			mode_sym = mode_sym % 12;
			act_sym = mode_sym / 4; // coding of cbp: 0,1,2 
			act_ctx = 5;
			if (act_sym==0) 
			{
				biari_encode_symbol(eep_dp, 0, ctx->mb_type_contexts[0] + act_ctx );
			} 
			else 
			{
				biari_encode_symbol(eep_dp, 1, ctx->mb_type_contexts[0] + act_ctx );
				act_ctx=6;
				if (act_sym==1) 
				{
					biari_encode_symbol(eep_dp, 0, ctx->mb_type_contexts[0] + act_ctx );
				}	 
				else 
				{
					
					biari_encode_symbol(eep_dp, 1, ctx->mb_type_contexts[0] + act_ctx );
				}	
				
			}
			mode_sym = mode_sym % 4; /* coding of I pred-mode: 0,1,2,3 */
			act_sym = mode_sym/2;
			act_ctx = 7;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[0] + act_ctx );
			act_ctx = 8;
			act_sym = mode_sym%2;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[0] + act_ctx );
		}
	}
	else 
	{
		
		if (currMB->mb_available[0][1] == NULL)
			b = 0;
		else 
			b = (( (currMB->mb_available[0][1])->mb_type != 0) ? 1 : 0 );
		if (currMB->mb_available[1][0] == NULL)
			a = 0;
		else 
			a = (( (currMB->mb_available[1][0])->mb_type != 0) ? 1 : 0 );
		
		
		act_sym = curr_mb_type;
				
		if(act_sym>=(mode16x16=(8*img->type+1))) // 16x16 Intra-mode: mode16x16=9 (P-frame) mode16x16=17 (B-frame)
		{
			mode_sym=act_sym-mode16x16;
			act_sym=mode16x16; // 16x16 mode info 
		}
		
		act_sym++;
				
		for (log_sym = 0; (1<<log_sym) <= act_sym; log_sym++);
		log_sym--;
		
		act_ctx = a + 2*b;
		se->context = act_ctx; // store context
		
		if (log_sym==0) 
		{
			biari_encode_symbol(eep_dp, 0, &ctx->mb_type_contexts[1][act_ctx] );
			
		} 
		else 
		{
			// code unary part 
			biari_encode_symbol(eep_dp, 1, &ctx->mb_type_contexts[1][act_ctx] );
			act_ctx=4;
			if (log_sym==1) 
			{
				biari_encode_symbol(eep_dp, 0, &ctx->mb_type_contexts[1][act_ctx ] );
			}	 
			else 
			{
				for(l=0;l<log_sym-1;l++)
				{
					biari_encode_symbol(eep_dp, 1, &ctx->mb_type_contexts[1][act_ctx] );
					if (l==0) act_ctx=5;
				}
				if ( log_sym < (3+((img->type == B_IMG)?1:0)) ) // maximum mode no. is 9 (P-frame) or 17 (B-frame)
					biari_encode_symbol(eep_dp, 0, &ctx->mb_type_contexts[1][act_ctx ] );	
			}
			
			// code binary part 
			act_ctx=6;
			if (log_sym==(3+((img->type == B_IMG)?1:0)) ) log_sym=2; // only 2 LSBs are actually set for mode 7-9 (P-frame) or 15-17 (B-frame)
			mask = (1<<log_sym); // MSB 
			for(l=0;l<log_sym;l++)
			{
				mask >>=1;
				biari_encode_symbol(eep_dp, (unsigned char) (act_sym & mask), &ctx->mb_type_contexts[1][act_ctx] );
			}
		}
				
				
		if(act_sym==(mode16x16+1)) // additional info for 16x16 Intra-mode 
		{
			act_ctx = 7;
			act_sym = mode_sym/12;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[1] + act_ctx ); // coding of AC/no AC 
			mode_sym = mode_sym % 12;
			
			act_sym = mode_sym / 4; // coding of cbp: 0,1,2 
			act_ctx = 8;
			if (act_sym==0) 
			{
				biari_encode_symbol(eep_dp, 0, ctx->mb_type_contexts[1] + act_ctx );
			} 
			else 
			{
				biari_encode_symbol(eep_dp, 1, ctx->mb_type_contexts[1] + act_ctx );
				if (act_sym==1) 
				{
					biari_encode_symbol(eep_dp, 0, ctx->mb_type_contexts[1] + act_ctx );
				}	 
				else 
				{
					
					biari_encode_symbol(eep_dp, 1, ctx->mb_type_contexts[1] + act_ctx );
				}	
				
			}
			
			mode_sym = mode_sym % 4; // coding of I pred-mode: 0,1,2,3 
			act_ctx = 9;
			
			act_sym = mode_sym/2;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[1] + act_ctx );
			act_sym = mode_sym%2;
			biari_encode_symbol(eep_dp, (unsigned char) act_sym, ctx->mb_type_contexts[1] + act_ctx );
						
		}	
	}
}

/****************************************************************************
*
* Function:        writeIntraPredMode2Buffer_CABAC()
*
* Purpose:         This function is used to arithmetically encode a pair of