/*
    Project  : H.264 encoder test    
    Module   : Context-adaptive variable-length coding (CAVLC)

    Author   : Marko 'Fador' Viitanen

    Date     : 11/5/2009
    Modified : 27/6/2009

*/

#include <math.h>
#include <stdarg.h>
#include <stdio.h>
#include "common.h"
#include "tables.h"
#include "cavlc.h"
#include "bitstream.h"



#define DEBUG 1

void printf_cavlc(char *msg, ...)
{
#if DEBUG == 1
va_list fmtargs;
char buffer[1024];

  va_start(fmtargs,msg);
  vsnprintf(buffer,sizeof(buffer)-1,msg,fmtargs);
  va_end(fmtargs);
  printf("%s",buffer);
#endif

}

//From wikipedia
//http://en.wikipedia.org/wiki/Binary_logarithm#Algorithm
int floorLog2(unsigned int n) {
  int pos = 0;
  if (n >= 1<<16) { n >>= 16; pos += 16; }
  if (n >= 1<< 8) { n >>=  8; pos +=  8; }
  if (n >= 1<< 4) { n >>=  4; pos +=  4; }
  if (n >= 1<< 2) { n >>=  2; pos +=  2; }
  if (n >= 1<< 1) {           pos +=  1; }
  return ((n == 0) ? (-1) : pos);  
}

//Initialize the Exp Golomb code table with desired number of values
void init_exp_golomb(uint32 len)
{
    uint32 code_num;
    uint32 M;
    uint32 info;
    exp_table=(bitTable*)malloc(len*sizeof(bitTable));    
    
    for(code_num=0;code_num<len;code_num++)
    {
        M=(uint32)floorLog2(code_num+1);
        info=code_num+1-(uint32)pow(2,M);        
        exp_table[code_num].len=M*2+1;
        exp_table[code_num].value=(1<<M)|info;
        //printf_cavlc("Len: %i %x\n", M*2+1, (1<<M)|info);
    }
}


uint32 get_exp_golomb_code(sint32 value, uint32* len, uint8 maptype)
{
    uint32 index=0;

    switch(maptype)
    {
        //Unsigned
        case EXP_UNSIGNED:
            index=(uint32)value;
            break;
        //Signed
        case EXP_SIGNED:
            index=(value<=0)?2*(uint32)abs(value):2*(uint32)abs(value)-1;
            break;
        //Intra coding
        case EXP_INTRA:
            index=coded_block_pattern_intra[value];
            break;
        //Inter coding
        case EXP_INTER:
            index=coded_block_pattern_inter[value];
            break;
    }

    len[0]=exp_table[index].len;
    return exp_table[index].value;
}


static uint8 countCoeffs(sint32 *block, uint8 totCount,valueTable *coeffTable)
{
    uint8 i;
    uint8 total=0;
    for(i=0;i<totCount;i++)
    {
        if(block[i]!=0)
        {
            coeffTable[total].value=block[i];
            coeffTable[total].pos=i;
            total++;
        }
    }
    return total;
}

static uint8 countZeros(sint32 *block, uint8 totCount,valueTable *zeroTable)
{
    sint8 i;
    uint8 total=0;
    for(i=totCount-1;i>=0;i--)
    {
        if(block[i]==0)
        {
            zeroTable[total].value=block[i];
            zeroTable[total].pos=i;
            total++;
        }
    }
    return total;
}

//Find trailing ones from the block, maximum of 3 ones
static uint8 findTrailingOnes(sint32 *block, uint8 totCount, valueTable *table)
{
    uint8 oneCount=0;
    sint8 i;
    uint32 tempval;
    
    for(i=totCount-1;i>=0;i--)
    {
        tempval=(uint32)abs(block[i]);
        //Non-zero
        if(tempval!=0)
        {
            //One
            if(tempval==1)
            {
                table[oneCount].value=block[i];
                table[oneCount].pos=i;
                oneCount++;
                if(oneCount==3) return oneCount;
            }
            //Other coeff
            else
            {
                return oneCount;
            }
        }
    }
    //This should newer be executed! Just for the show ;)
    return oneCount;
}

bitTable lev_VLCN[7][4096];


uint8 genLevelVLC()
{
    uint16 code;
    uint8 table;


    for(code=0;code<4096;code++)
    {

        //level VLC0
        if(code<14)
        {
            lev_VLCN[0][code].len=code+1;
            lev_VLCN[0][code].value=1;
        }
        else if(code<30)
        {
            lev_VLCN[0][code].len=19;
            lev_VLCN[0][code].value=(1<<4)|code;
        }
        else
        {
            lev_VLCN[0][code].len=28;
            lev_VLCN[0][code].value=(1<<12)|code;
        }

        //level VLC1-6
        //TODO verify tables!
        for(table=1;table<7;table++)
        {
            if(code<30*(1<<(table-1)))
            {
                lev_VLCN[table][code].len=((uint32)floor(code/(1<<table)))+(table+1);
                lev_VLCN[table][code].value=(1<<table)|(code%(1<<table));
            }
            else
            {
                lev_VLCN[2][code].len=28;
                lev_VLCN[2][code].value=(1<<12)|code;
            }
        }
    }

}

bitTable levelVLC(sint32 levelCode, uint8 table)
{
    bitTable returnValue;   
          
    returnValue=lev_VLCN[table][((levelCode<=0)?2*(uint32)abs(levelCode+1)+1:2*(uint32)abs(levelCode-1))];
    printf_cavlc("levelVLC(%i,%i) = {%i,%i} val: %i\n", levelCode,table,returnValue.value,returnValue.len,
        ((levelCode<=0)?2*(uint32)abs(levelCode+1)+1:2*(uint32)abs(levelCode-1)));

    return returnValue;
}

uint8 cavlc_code_coeff_4x4(sint32 block[4*4], bitstream *stream, uint8 nC)
{
    sint32 reordered[4*4];
    sint16 i;
    valueTable oneTable[3];
    valueTable zeroTable[16];
    valueTable coeffTable[16];
    uint8 TotalCoeffs;
    uint8 TrailingOnes;
    //Which table to select for nC
    uint8 nC_to_table[16] = {0,0,1,1,2,2,2,2,3,3,3,3,3,3,3,3};
    //Thresholds to move to next suffixLength
    uint8 suffixThresholds[7] = {0,3,6,12,24,48,255};
    uint8 tableToUse=nC_to_table[nC];
    uint8 suffixLength=0;
    uint8 total_zeros;
    uint8 ZerosLeft;
    uint8 run_before;
    bitTable tempValue;
    

    genLevelVLC();

    printf_cavlc("Reorder: ");
    //Zig-zag reorder
    for(i=0;i<4*4;i++)
    {
        reordered[i]=block[zigzag_4x4[i]];
        printf_cavlc("%i,",reordered[i]);
    }
    printf_cavlc("\n");

    //Stage 1 Page 201
    //coeff token
        
    TrailingOnes=findTrailingOnes(&reordered, 4*4, &oneTable);
    TotalCoeffs=countCoeffs(&reordered, 4*4, &coeffTable);
    printf_cavlc("TotalCoeffs: %i\n", TotalCoeffs);

    //Value from table for TotalCoeffs and TrailingOnes
    bitstream_put(stream, coeff_token_tables[TotalCoeffs][TrailingOnes][tableToUse].value,
                          coeff_token_tables[TotalCoeffs][TrailingOnes][tableToUse].len);

    //Init suffixLength as 1 when more than 10 coeffs and less than 3 trailing ones
    //Page 202 book, page 223 specification
    if(TotalCoeffs>10 && TrailingOnes<3) suffixLength=1;

    //TrailingOnes signs to stream
    for(i=0;i<TrailingOnes;i++)
    {
        bitstream_put(stream,(oneTable[i].value==-1)?1:0, 1);
    }

    //Levels to bitstream (not the trailing ones)
    for(i=TotalCoeffs-1-TrailingOnes;i>=0;i--)
    {
        //If less than 3 trailing ones the first level to encode
        //Cannot be 1, so we can save bits to set +-2=+-1, +-3=+-2 etc
        if(TrailingOnes<3 && i==TotalCoeffs-1-TrailingOnes){  coeffTable[i].value=((coeffTable[i].value<0)?coeffTable[i].value+1:coeffTable[i].value-1);}

        //Level VLC from tables
        tempValue=levelVLC(coeffTable[i].value,suffixLength);
        bitstream_put(stream,tempValue.value,tempValue.len);        

        //Check if we need to increase suffix Length
        if((uint32)abs(coeffTable[i].value)>suffixThresholds[suffixLength])
            suffixLength++;
    }

    //Total number of zeros
    total_zeros=countZeros(&reordered,(TrailingOnes?oneTable[0].pos:coeffTable[TotalCoeffs-1].pos),zeroTable);
    printf_cavlc("total_zeros: %i\n", total_zeros);
    bitstream_put(stream,total_zeros_4x4[total_zeros][TotalCoeffs-1].value,total_zeros_4x4[total_zeros][TotalCoeffs-1].len);
    
    //Run before

    if(total_zeros)
    {
        run_before=0;
        ZerosLeft=total_zeros;

        for(i=zeroTable[0].pos;i>=0&&i>coeffTable[0].pos&&ZerosLeft>0;i--)
        {
            if(reordered[i]==0)
            {
                run_before++;
                printf_cavlc("run_before++\n");
            }
            else
            {                
                printf_cavlc("run_before: %i zerosleft: %i\n", run_before,ZerosLeft);
                bitstream_put(stream,run_before_table[run_before][((ZerosLeft>6)?6:ZerosLeft-1)].value, 
                                     run_before_table[run_before][((ZerosLeft>6)?6:ZerosLeft-1)].len);
                ZerosLeft-=run_before;
                run_before=0;
            }
        }
        if(run_before)
        {
           bitstream_put(stream,run_before_table[run_before][((ZerosLeft>6)?6:ZerosLeft-1)].value, 
                                     run_before_table[run_before][((ZerosLeft>6)?6:ZerosLeft-1)].len);
        }
    }


    for(i=0;i<stream->cur_bit;i++)
    {
        printf_cavlc("%i", (((stream->data[0]&(1<<(31-i)) ))?1:0));
    }
    
    printf_cavlc("\n");
    //printf_cavlc("\n%x %x %i %i\n", stream->data[0],stream->data[1], stream->cur_bit, stream->cur_byte);

    return TotalCoeffs;
}