/*
 * Copyright 1993 by the University of Pennsylvania
 *
 * Permission to use, copy, and distribute for non-commercial purposes,
 * is hereby granted without fee, providing that the above copyright
 * notice appear in all copies and that both the copyright notice and this
 * permission notice appear in supporting documentation.
 *
 * The software may be modified for your own purposes, but modified versions
 * may not be distributed.
 *
 * This software is provided "as is" without any expressed or implied warranty.
 */

#include <stdio.h>
#include <fcntl.h>
#include <memory.h>
#include "global.h"
#include "mpeg_code.h"
#include "errorhandle.h"

#ifndef WRITE_LENGTH
#define WRITE_LENGTH 4096  /* buffer size for encoded data */
#endif

unsigned int data[WRITE_LENGTH];
unsigned int datal[WRITE_LENGTH];

int off = 0;
int offl = 0;
int bits_left = 32;
int bits_leftl = 32;
int mb_bits; /* counter for MB generation */

FILE *datafph,*datafpl;

/******************************************************************
  define general value for VLC
 ******************************************************************/
int curr_pri;          /* 0: HP, 1: LP */
int mb_type;    /* Macroblock type 1: intra, 2:pred or Bi-Dir */
/*******************************************************************
  Do VLC puts MB header and MB data into output stream
********************************************************************/
int
  do_vlc(mblock,priority,dctnum,first, picture_type, n_hor_mb)
MB mblock;
int priority; /* 0: HP, 1: LP */
int dctnum;   /* if 64, may assume NON-Priority */
int first;    /* MB is first(1) or not(0) */
int picture_type; /* 0:Intra, 1:Fwd Pred, 2:Bwd Pred, 3:Bi Pred */ 
int n_hor_mb;
{
  int l;
  CODE vlc;
  static int mb_increment = 0;
  static int pat = 0;
  int force_code = 0;
  
  curr_pri = priority; /* make priority of current block global */
  mb_bits = 0; /* reset generation bit counter */

/*------------------------------------------------------------------*/
/*      High Priority packets : contains all header info            */
/*------------------------------------------------------------------*/
  if(priority == 0) { 
/*------------------------------------------------------------------*/
    if (first == 1)   /* reset mb_increment if first MB in slice */
      mb_increment = 0;

    /* first and last mb in slice must be coded*/
    if ((mb_increment == n_hor_mb - 1) || (first == 1)) 
      force_code = 1;

/*---------------- Check to see if MB is coded (get cbp) --------------*/
    pat = get_pattern(mblock.data[1]);

    if (force_code == 1) 
      if ((pat == 0) && (mb_type != 1) && (mblock.mc == 1))
	mblock.mc = 1;
	
/*----------MB address increment ---------------------------------------*/
    if ((pat)||(mb_type == 1)||(mblock.mc == 1)) {
      while(mb_increment > 32) {
	put_bits(position[34]);
	mb_increment = mb_increment - 33;
      }
      put_bits(position[mb_increment]);
      mb_increment = 0;
    }
    else
      mb_increment++;
    
/*----------- put header in each frame type  ----------------------------*/
    if(mb_increment ==0)
    switch(picture_type) {
      case 0: /*--- Intraframe Picture Mode ---*/
        vlc.data = 1; vlc.length = 1;
        put_bits(vlc);
        pat = 63;
        break;
      case 1: /*--- Fwd-Prediction Picture Mode ---*/
        if(mblock.intra == 1) { /* MB intra */
          vlc.data = 3; vlc.length = 5;
          put_bits(vlc);
          pat = 63;
        }
        else { /*--- MB Pred ---*/ 
          vlc.data = 1;
          if(mblock.mc == 1) { /* MB MC */
            if(pat == 0) 
              vlc.length = 3; /* Non patterned */
            else 
              vlc.length = 1; /* Patterned */
            put_bits(vlc);
            /* Motion vector information */
            put_bits(mf[mblock.mv[1]+16]);
            put_bits(mf[mblock.mv[0]+16]);
          }
          else if (pat) { /* No MC  Patterned */ 
            vlc.length = 2;
            put_bits(vlc);
          }  /* there does not exist No-MC, Non-Patterned MB */
	  if (pat) put_bits(cbp[pat-1]); /* Insert pattern if patterned */
        }
        break;
      case 3: /*--- Bidirectional Picture Mode ---*/
        if(mblock.intra == 1) { /* MB intra */
          vlc.data = 3;  vlc.length = 5;
          put_bits(vlc);
          pat = 63;
        }
        else { /*--- Bi-directional MB ---*/ 
          if (pat)
             vlc.data = 3;
          else
             vlc.data = 2;
          switch(mblock.type) {
            case 1: vlc.length = 4; break; /* Fwd Pred */
            case 2: vlc.length = 3; break; /* Bwd Pred */
            case 3: vlc.length = 2; break; /* Bi-Directional */
          }
          put_bits(vlc);
          /*--- put MV information ---*/
          switch(mblock.type) {
            case 1: /* Fwd MV's */
              put_bits(mf[mblock.mv[1]+16]);
              put_bits(mf[mblock.mv[0]+16]);
              break;
            case 2: /* Bwd MV's */
              put_bits(mf[mblock.mv[3]+16]);
              put_bits(mf[mblock.mv[2]+16]);
              break;
            case 3: /* Bi-Directional (Fwd and Bwd) MV's */
              put_bits(mf[mblock.mv[1]+16]);
              put_bits(mf[mblock.mv[0]+16]);
              put_bits(mf[mblock.mv[3]+16]);
              put_bits(mf[mblock.mv[2]+16]);
              break;
	    }
	  /* Insert pattern if patterned */
	  if (pat) put_bits(cbp[pat-1]);
	}
	break;
      } /* end of header info for each MB */
/*------------ actual coded data ------------*/
    for(l=0;l<6;l++)  
      if (pat & (1 << (5-l))) {
        code_vl(mblock.data[priority+1]+l*DCT_SIZE_2,l,dctnum);
      }
  }
  else if(dctnum<64)  {  /* LP stream */
    put_bits(position[mb_increment]);
    for(l=0;l<6;l++) 
      if (pat & (1 << (5-l))) {
	code_vl(mblock.data[priority+1]+l*DCT_SIZE_2 + dctnum, l,
		(DCT_SIZE_2 - dctnum));
      }
  }
  curr_pri = 0;
  return mb_bits;
}

int
  get_pattern(macroblock)
short *macroblock;
{

  int i,j,coded,cbp;

  for(cbp = 0, j = 0; j < 6; j++){
      for(coded = 0, i = 0; i < 64; i++)
	if( macroblock[i + j*DCT_SIZE_2] ) {
	  coded = 1;
	  break;
	}
      cbp += (coded == 1 ? (1 << (5 - j)) : 0);
    }

  return cbp;
}

void
  code_vl(block, type, dctnum)
short *block;
int type;
int dctnum;
{
  int i=0,sign;
  int ct=0,val;
  CODE vlc;
  /* Variable length code for dc component in intra macroblocks*/
  if(mb_type == 1) {
    code_dc_coeff(block[0],type);
    i = 1;
  }
  for(;i < dctnum; i++) {
    if (block[i] == 0) 
      ct++;
    else {
      sign = (block[i] < 0);
      val = (sign ? -block[i] : block[i]);
      if ((i == 0) && (val == 1)) {
	vlc.data = 1;
	vlc.length = 1;
	put_vl_code(vlc,sign);
      }
      else {
        if ((ct < 32) && (val <= 40))
          if (dct_codes[ct][val - 1].length == 0)
            put_fl_code(ct,block[i]);
          else {
            vlc = dct_codes[ct][val - 1];
            put_vl_code(vlc,sign);
          }
        else 
          put_fl_code(ct,block[i]);
      }
      ct = 0;
    }
  }
  put_bits(eob);
}

void
  put_fl_code(run, level)
int run;
short level;
{
  CODE flc;
  put_bits(esc);
  /* write fixed length code (FLC) for length of run */
  flc.data = run;
  flc.length = 6;
  put_bits(flc);
  /* write FLC for next level */
  if (abs(level) >= 128) {
    if ( level < 0 )
      flc.data = 128;
    else
      flc.data = 0;
    flc.length = 8;
    put_bits(flc);
  }
  flc.data = (level & 0xff);
  flc.length = 8;
  put_bits(flc);
}

void 
  put_vl_code( bcode, sign_val)
CODE bcode;
int sign_val;
{
  CODE sign;

  /* add vlc to the data stream */
  put_bits(bcode);

  /* add sign bit to the data stream */
  sign.data = sign_val;
  sign.length = 1;
  put_bits(sign);
}

void
  code_dc_coeff( value, type)
short value;
int type;
{
  int size=0; /* size of dc_coeff */
  short tmp = value;
  CODE val;
  if(value < 0)
    tmp = -tmp;
  if(value)
    while(tmp >> size)
      size++ ;
  if(type < 4)
    put_bits(dct_dc_size_y[size]);
  else
    put_bits(dct_dc_size_uv[size]);
  if(size) {
    if(value < 0)
      val.data = value - 1 + (1 << size);
    else
      val.data = value ;
    val.length = size;
    put_bits(val);
  }
}

void
  put_bits(ocode)
CODE ocode;
{
  mb_bits += ocode.length;
  if (curr_pri == 0) {
    if (ocode.length >= bits_left) {
      data[off++] +=  (ocode.data >> (ocode.length - bits_left));
      if (off >= WRITE_LENGTH) 
	reset_data(0);
      bits_left = 32 - (ocode.length - bits_left);
      if (bits_left != 32)
	data[off] = (ocode.data << bits_left);
    }
    else {
      data[off] += (ocode.data << (bits_left - ocode.length));
      bits_left = bits_left - ocode.length;
    }
  }
  else {
    if (ocode.length >= bits_leftl) {
      datal[offl++] +=  (ocode.data >> (ocode.length - bits_leftl));
      if (offl >= WRITE_LENGTH) 
	reset_data(1);
      bits_leftl = 32 - (ocode.length - bits_leftl);
      if (bits_left != 32)
	datal[offl] = (ocode.data << bits_leftl);
    }
    else {
      datal[offl] += (ocode.data << (bits_leftl - ocode.length));
      bits_leftl = bits_leftl - ocode.length;
    }
  }
}

void
  align()
{ /* only used for HP */
  bits_left += -(bits_left % 8);
  
  if (bits_left == 0) {
    off++; bits_left = 32;
  }

  if (off >= WRITE_LENGTH)
    reset_data(0);
}
/*************************************************************
  MPEG coded data related functions
 *************************************************************/
/*************************************/
/* OPEN output file for encoded data */
/*************************************/
int
  open_mpeg_file(fname,dctnum)
char *fname;
int dctnum;
{
  char fnamel[256];

  /* init global value */
  datafph = NULL;
  datafpl = NULL;

  off = 0;
  offl = 0;
  bits_left = 32;
  bits_leftl = 32;

  memset(data, 0,WRITE_LENGTH*sizeof(int));
  memset(datal,0,WRITE_LENGTH*sizeof(int));

  if(strlen(fname) == 0) return 0; /*--- non record mode ---*/

  if((datafph = fopen(fname,"w")) == NULL) { /* MPEG data for H code */
    errnum = 25;
    strcpy(errmsg,fname);
    return -1;
  }
#ifdef DEBUG
  else 
    fprintf(stderr,"MPEG data file opened (%s)\n",fname);
#endif
  if((dctnum < 64)&&(!strncmp(fname,"/dev/",5))) {
    sprintf(fnamel,"%sl",fname);
    if((datafpl = fopen(fnamel,"w")) == NULL) { /* MPEG data for L code */
      errnum = 25;
      strcpy(errmsg,fname);
      return -1;
    }
    else {
      fprintf(stderr,"MPEG data file opened (%s)\n",fname);
    }
  }
#ifdef DEBUG
  else 
    fprintf(stderr,"No Low priority MPEG data file.\n");
#endif
  return 0;
}

int
  close_mpeg_file() /* flush and clsoe MPEG encode data */ 
{
  if(datafph != NULL) { /* H code */
    if (bits_left  != 32) off++;
    reset_data(0);
    fclose(datafph);
  }
  if(datafpl != NULL) { /* L code */
    if (bits_leftl != 32) offl++;
    reset_data(1);
    fclose(datafpl);
  }
}

void
  reset_data(pri)
int pri;
{
  /* This routine is used to dump and reset data when array limit
     has been reached */
  if(pri==0) {
    if(datafph != NULL) {
      fwrite(data,sizeof(int),off,datafph);
    }
    off  =  0;
    memset(data,0,WRITE_LENGTH*sizeof(int));
  }
  else if(pri ==1) {
    if(datafpl != NULL) {
      fwrite(datal,sizeof(int),offl,datafpl);
    }
    offl =  0;
    memset(datal,0,WRITE_LENGTH*sizeof(int));
  }
}
