/* 
 * Copyright (c) 1997 NextLevel Systems of Delaware, Inc.  All rights reserved.
 * 
 * This software module  was developed by  Bob Eifrig (at NextLevel
 * Systems of Delaware, Inc.), Xuemin Chen (at NextLevel Systems of
 * Delaware, Inc.), and Ajay Luthra (at NextLevel Systems of Delaware,
 * Inc.), in the course of development of the MPEG-4 Video Standard
 * (ISO/IEC 14496-2).   This software module is an implementation of a
 * part of one or more tools as specified by the MPEG-4 Video Standard.
 * 
 * NextLevel Systems of Delaware, Inc. grants the right, under its
 * copyright in this software module, to use this software module and to
 * make modifications to it for use in products which conform to the
 * MPEG-4 Video Standard.  No license is granted for any use in
 * connection with products which do not conform to the MPEG-4 Video
 * Standard.
 * 
 * Those intending to use this software module are advised that such use
 * may infringe existing and unissued patents.  Please note that in
 * order to practice the MPEG-4 Video Standard, a license may be
 * required to certain patents held by NextLevel Systems of Delaware,
 * Inc., its parent or affiliates ("NextLevel").   The provision of this
 * software module conveys no license, express or implied, under any
 * patent rights of NextLevel or of any third party.  This software
 * module is subject to change without notice.  NextLevel assumes no
 * responsibility for any errors that may appear in this software
 * module.  NEXTLEVEL DISCLAIMS ALL WARRANTIES, EXPRESS AND IMPLIED,
 * INCLUDING, BUT NOT LIMITED TO ANY WARRANTY THAT COMPLIANCE WITH OR
 * PRACTICE OF THE SPECIFICATIONS OR USE OF THIS SOFTWARE MODULE WILL
 * NOT INFRINGE THE INTELLECTUAL PROPERTY RIGHTS OF NEXTLEVEL OR ANY
 * THIRD PARTY, AND ANY IMPLIED WARRANTIES OF MERCHANTABILITY AND
 * FITNESS FOR A PARTICULAR PURPOSE.
 * 
 * NextLevel retains the full right to use this software module for its
 * own purposes, to assign or transfer this software module to others,
 * to prevent others from using this software module in connection with
 * products which do not conform to the MPEG-4 Video Standard, and to
 * prevent others from infringing NextLevel's patents.
 * 
 * As an express condition of the above license grant, users are
 * required to include this copyright notice in all copies or derivative
 * works of this software module.
 */
/* MPEG-2 Video TM-5 rate control 
   For the purpose of comparing
   MPEG-4 VM with MPEG-2 TM-5
   27.03.97 By X. Chen in G.I.
 */

#include <stdio.h>
#include <math.h>
#include "mom_structs.h"
#include "vm_config.h"
#include "vm_stats.h"
#include "tm5rc.h"

double tm5_var_sblk(SInt *p, int lx);
void tm5_calc_actj(Vop *vop, double *mbact);


void *tm5rc_init_seq(VolConfig *cfg, Int rc_type)
{
  Tm5Rc *rc;
  Int len;

  if (cfg->shape != RECTANGULAR) {
      fprintf(stderr, "TM5 rate control requaires a rectangular VOP\n");
      exit(1);
  }
  rc = (Tm5Rc *)ecalloc(sizeof(Tm5Rc), 1);
  rc->rc_type = rc_type;
  switch (rc_type) {

  case TM5_RATE_CONTROL:
      break;

  case TM5_RATE_CONTROL+1:
      len = strlen(cfg->bitstream);
      cfg->bitstream[len] = 'Q';
      cfg->bitstream[len + 1] = 0;
      if ((rc->Qfile = fopen(cfg->bitstream, "r")) == NULL) {
          fprintf(stderr, "Cannot open %s\n", cfg->bitstream);
          exit(1);
      }
      cfg->bitstream[len] = 0;
      break;

  case TM5_RATE_CONTROL+2:
      len = strlen(cfg->bitstream);
      cfg->bitstream[len] = 'Q';
      cfg->bitstream[len + 1] = 0;
      if ((rc->Qfile = fopen(cfg->bitstream, "w")) == NULL) {
          fprintf(stderr, "Cannot open %s\n", cfg->bitstream);
          exit(1);
      }
      cfg->bitstream[len] = 0;
      break;

  default:
      fprintf(stderr, "Invalid rate control code: %d\n", rc_type);
      exit(1);
  }
  rc->mb_width = cfg->disk_seq_x / MB_SIZE;
  rc->mb_height = cfg->disk_seq_y / MB_SIZE;
  rc->bitrate = cfg->bit_rate;
  rc->pic_rate = cfg->frame_rate;
  rc->mbact = (double *)ecalloc(rc->mb_width * rc->mb_height, sizeof(double));

  /* reaction parameter (constant) */
  /* if (rc->r_tm5==0)  */
      rc->r_tm5 = (int)floor(2.0*rc->bitrate/rc->pic_rate + 0.5);

  /* average activity */
  if (rc->avg_act==0.0)  rc->avg_act = 400.0;

  /* remaining # of bits in GOP */
  rc->R_tm5 = 0;

  /* global complexity measure */
  if (rc->Xi==0) rc->Xi = (int)floor(160.0*rc->bitrate/115.0 + 0.5);
  if (rc->Xp==0) rc->Xp = (int)floor( 60.0*rc->bitrate/115.0 + 0.5);
  if (rc->Xb==0) rc->Xb = (int)floor( 42.0*rc->bitrate/115.0 + 0.5);

  /* virtual buffer fullness */
  if (rc->d0i==0) rc->d0i = (int)floor(10.0*rc->r_tm5/31.0 + 0.5);
  if (rc->d0p==0) rc->d0p = (int)floor(10.0*rc->r_tm5/31.0 + 0.5);
  if (rc->d0b==0) rc->d0b = (int)floor(1.4*10.0*rc->r_tm5/31.0 + 0.5);

  fprintf(stdout,"\nrate control: sequence initialization\n");
  fprintf(stdout,
    " initial global complexity measures (I,P,B): Xi=%d, Xp=%d, Xb=%d\n",
    rc->Xi, rc->Xp, rc->Xb);
  fprintf(stdout," reaction parameter: r_tm5=%d\n", rc->r_tm5);
  fprintf(stdout,
    " initial virtual buffer fullness (I,P,B): d0i=%d, d0p=%d, d0b=%d\n",
    rc->d0i, rc->d0p, rc->d0b);
  fprintf(stdout," initial average activity: avg_act=%.1f\n", rc->avg_act);

  return (void *)rc;
}

void tm5rc_init_GOP(int np,int nb,void *rcdata)
{
  Tm5Rc *rc = (Tm5Rc *)rcdata;

  if ((rc->Ni + rc->Np + rc->Nb) != 0) {
      fprintf(stderr, "TM5: Short GOP, expected %d/%d/%d more I/P/B-pics=n",
          rc->Ni, rc->Np, rc->Nb);
      exit(1);
  }
  rc->R_tm5 += (int) floor((1 + np + nb) * rc->bitrate / rc->pic_rate + 0.5);
  rc->Np = np;
  rc->Nb = nb;
  rc->Ni = 1;

  fprintf(stdout,"\nrate control: new group of pictures (GOP)\n");
  fprintf(stdout," target number of bits for GOP: R_tm5=%d\n",rc->R_tm5);
  fprintf(stdout," number of P pictures in GOP: Np=%d\n",rc->Np);
  fprintf(stdout," number of B pictures in GOP: Nb=%d\n",rc->Nb);
}

/* Note: we need to substitute K for the 1.4 and 1.0 constants -- this can
   be modified to fit image content */

/* Step 1: compute target bits for current picture being coded */
void tm5rc_init_pict(Vop *vop, void *rcdata)
{
  double Tmin;
  Tm5Rc *rc = (Tm5Rc *)rcdata;
  int type;

  switch (rc->rc_type) {
  
  case TM5_RATE_CONTROL+1:
      type = -1;
      if ((fscanf(rc->Qfile, "%d", &type) != 1) || (type != vop->prediction_type)) {
          fprintf(stderr, "Wrong pictue type: got %d, expected %d\n",
              type, vop->prediction_type);
          exit(1);
      }
      return;

  case TM5_RATE_CONTROL+2:
      fprintf(rc->Qfile, "%d ", vop->prediction_type);
      break;
  }

  switch (vop->prediction_type)
  {
  case I_VOP:
    rc->T_tm5 = (int) floor(rc->R_tm5/(1.0+rc->Np*rc->Xp/(rc->Xi*1.0)+rc->Nb*rc->Xb/(rc->Xi*1.4)) + 0.5);
    rc->d_tm5 = rc->d0i;
    break;
  case P_VOP:
    rc->T_tm5 = (int) floor(rc->R_tm5/(rc->Np+rc->Nb*1.0*rc->Xb/(1.4*rc->Xp)) + 0.5);
    rc->d_tm5 = rc->d0p;
    break;
  case B_VOP:
    rc->T_tm5 = (int) floor(rc->R_tm5/(rc->Nb+rc->Np*1.4*rc->Xp/(1.0*rc->Xb)) + 0.5);
    rc->d_tm5 = rc->d0b;
    break;
  }

  Tmin = (int) floor(rc->bitrate/(8.0*rc->pic_rate) + 0.5);

  if (rc->T_tm5<Tmin)
    rc->T_tm5 = (int)Tmin;

  rc->Q_tm5 = 0;

  tm5_calc_actj(vop, rc->mbact);

  rc->actsum = 0.0;

  fprintf(stdout,"\nrate control: start of picture\n");
  fprintf(stdout," target number of bits: T_tm5=%d\n",rc->T_tm5);
}

void tm5_calc_actj(Vop *vop, double *mbact)
{
  int i,j,k;
  SInt *p, *frame;
  double actj,var;

  k = 0;
  frame = (SInt *)GetImageData(GetVopY(vop));
  for (j=0; j<vop->height; j+=16)
    for (i=0; i<vop->width; i+=16)
    {
      p = frame + i + vop->width*j;

      /* take minimum spatial activity measure of luminance blocks */

      actj = tm5_var_sblk(p,vop->width);
      var = tm5_var_sblk(p+8,vop->width);
      if (var<actj) actj = var;
      var = tm5_var_sblk(p+8*vop->width,vop->width);
      if (var<actj) actj = var;
      var = tm5_var_sblk(p+8*vop->width+8,vop->width);
      if (var<actj) actj = var;

   
         /* field */
        var = tm5_var_sblk(p,vop->width<<1);
        if (var<actj) actj = var;
        var = tm5_var_sblk(p+8,vop->width<<1);
        if (var<actj) actj = var;
        var = tm5_var_sblk(p+vop->width,vop->width<<1);
        if (var<actj) actj = var;
        var = tm5_var_sblk(p+vop->width+8,vop->width<<1);
        if (var<actj) actj = var;
    
      actj+= 1.0;

      mbact[k++]= actj;
    }
}

void tm5rc_update_pict(Vop *vop, BitCount *bc, void *rcdata)
{
  double X;
  Tm5Rc *rc = (Tm5Rc *)rcdata;

  if (rc->rc_type == TM5_RATE_CONTROL+1) {
      switch (vop->prediction_type) {
          case I_VOP: rc->Ni--; break;
          case P_VOP: rc->Np--; break;
          case B_VOP: rc->Nb--; break;
      }
      return;
  }
                        /* total # of bits in picture */
  rc->R_tm5 -= bc->vop; /* remaining # of bits in GOV */
  X = (int) floor(bc->vop*((0.5*(double)(rc->Q_tm5))/
      (rc->mb_width*rc->mb_height)) + 0.5);
  rc->d_tm5+= bc->vop - rc->T_tm5;
  rc->avg_act = rc->actsum/(rc->mb_width*rc->mb_height);

  switch (vop->prediction_type)
  {
  case I_VOP:
    rc->Xi = (int)X;
    rc->d0i = rc->d_tm5;
    if (--(rc->Ni) < 0) {
        fprintf(stderr, "TM5: too many I-pics in GOP\n");
        exit(1);
    }
    break;
  case P_VOP:
    rc->Xp = (int)X;
    rc->d0p = rc->d_tm5;
    if (--(rc->Np) < 0) {
        fprintf(stderr, "TM5: too many P-pics in GOP\n");
        exit(1);
    }
    break;
  case B_VOP:
    rc->Xb = (int)X;
    rc->d0b = rc->d_tm5;
    if (--(rc->Nb) < 0) {
        fprintf(stderr, "TM5: too many B-pics in GOP\n");
        exit(1);
    }
    break;
  }

  fprintf(stdout,"\nrate control: end of picture\n");
  fprintf(stdout," actual number of bits=%ld\n",bc->vop);
  fprintf(stdout," average quantization parameter Q_tm5=%.1f\n",
    (double)rc->Q_tm5/(rc->mb_width*rc->mb_height));
  fprintf(stdout," remaining number of bits in GOP: R_tm5=%d\n",rc->R_tm5);
  fprintf(stdout,
    " global complexity measures (I,P,B): Xi=%d, Xp=%d, Xb=%d\n",
    rc->Xi, rc->Xp, rc->Xb);
  fprintf(stdout,
    " virtual buffer fullness (I,P,B): d0i=%d, d0p=%d, d0b=%d\n",
    rc->d0i, rc->d0p, rc->d0b);
  fprintf(stdout," remaining number of P pictures in GOP: Np=%d\n",rc->Np);
  fprintf(stdout," remaining number of B pictures in GOP: Nb=%d\n",rc->Nb);
  fprintf(stdout," average activity: avg_act=%.1f\n", rc->avg_act);
}

/* compute initial quantization stepsize (at the beginning of picture) */
int tm5rc_start_mb(void *rcdata)
{
    Tm5Rc *rc = (Tm5Rc *)rcdata;
    int mquant;

    if (rc->rc_type == TM5_RATE_CONTROL+1) {
        if (fscanf(rc->Qfile, "%d", &mquant) != 1) {
            fprintf(stderr, "Q file read error\n");
            exit(1);
        }
        return mquant;
    }

    mquant = (int) floor(rc->d_tm5*31.0/rc->r_tm5 + 0.5);
    
    /* clip mquant to legal range */
    if (mquant<1)
      mquant = 1;
    if (mquant>31)
      mquant = 31;

    rc->prev_mquant = rc->mquant = mquant;

    if (rc->rc_type == TM5_RATE_CONTROL+2) {
        fprintf(rc->Qfile, "%2d\n", mquant);
        rc->linectr = 0;
    }

    return mquant;
}

/* Step 2: measure virtual buffer - estimated buffer discrepancy */
int tm5rc_calc_mquant(int j, UInt size, void *rcdata)
{
  int mquant;
  double dj, Qj, actj, N_actj;
  Tm5Rc *rc = (Tm5Rc *)rcdata;

  if (rc->rc_type == TM5_RATE_CONTROL+1) {
      if (fscanf(rc->Qfile, "%d", &mquant) != 1) {
         fprintf(stderr, "Q file read error\n");
         exit(1);
      }
      return mquant;
  }

  /* measure virtual buffer discrepancy from uniform distribution model */
  dj = rc->d_tm5 + size - j*(rc->T_tm5/(rc->mb_width*rc->mb_height));

  /* scale against dynamic range of mquant and the bits/picture count */
  Qj = dj*31.0/rc->r_tm5;

  actj = rc->mbact[j];
  rc->actsum += actj;

  /* compute normalized activity */
  N_actj = (2.0*actj+rc->avg_act)/(actj+2.0*rc->avg_act);

    /* modulate mquant with combined buffer and local activity measures */
    mquant = (int) floor(Qj*N_actj + 0.5);
   
    /* clip mquant to legal range */
    if (mquant<1)
      mquant = 1;
    if (mquant>31)
      mquant = 31;

#if 0
    /* ignore small changes in mquant */
    if (mquant>=8 && (mquant-rc->mquant)>=-4 && (mquant-rc->mquant)<=4)
      mquant = rc->mquant;
#endif

    rc->mquant = mquant;
 
  rc->Q_tm5 += mquant; /* for calculation of average mquant */

  if (rc->rc_type == TM5_RATE_CONTROL+2) {
      fprintf(rc->Qfile, "%2d ", mquant);
      if (++(rc->linectr) >= rc->mb_width) {
          fprintf(rc->Qfile, "\n");
          rc->linectr = 0;
      }
  }

  return mquant;
}

/* compute variance of 8x8 block */
double tm5_var_sblk(SInt *p, int lx)
{
  int i, j;
  unsigned int v, s, s2;

  s = s2 = 0;

  for (j=0; j<8; j++)
  {
    for (i=0; i<8; i++)
    {
      v = *p++;
      s+= v;
      s2+= v*v;
    }
    p+= lx - 8;
  }

  return s2/64.0 - (s/64.0)*(s/64.0);
}

