#include <assert.h>
#include "util.h"
#include "tables.h"
#include "reservoir.h"
#include "quantize-pvt.h"

FLOAT masking_lower=1;
int convert_mdct, reduce_sidechannel;
/*
mt 5/99.  These global flags denote 4 possibilities:
                                                                mode    l3_xmin
1   MDCT input L/R, quantize L/R,   psy-model thresholds: L/R   -m s     either
2   MDCT input L/R, quantize M/S,   psy-model thresholds: L/R   -m j     orig
3   MDCT input M/S, quantize M/S,   psy-model thresholds: M/S   -m f     either
4   MDCT input L/R, quantize M/S,   psy-model thresholds: M/S   -m j -h  m/s

1:  convert_mdct = 0, convert_psy=0,  reduce_sidechannel=0          
2:  convert_mdct = 1, convert_psy=1,  reduce_sidechannel=1
3:  convert_mdct = 0, convert_psy=0,  reduce_sidechannel=1   (this mode no longer used)
4:  convert_mdct = 1, convert_psy=0,  reduce_sidechannel=1

if (convert_mdct), then iteration_loop will quantize M/S data from
the L/R input MDCT coefficients.

if (convert_psy), then calc_noise will compute the noise for the L/R
channels from M/S MDCT data and L/R psy-model threshold information.
Distortion in ether L or R channel will be marked as distortion in
both Mid and Side channels.  
NOTE: 3/00: this mode has been removed.  

if (reduce_sidechannel) then outer_loop will allocate less bits
to the side channel and more bits to the mid channel based on relative 
energies.
*/



/*
  The following table is used to implement the scalefactor
  partitioning for MPEG2 as described in section
  2.4.3.2 of the IS. The indexing corresponds to the
  way the tables are presented in the IS:

  [table_number][row_in_table][column of nr_of_sfb]
*/
unsigned nr_of_sfb_block[6][3][4] =
{
  {
    {6, 5, 5, 5},
    {9, 9, 9, 9},
    {6, 9, 9, 9}
  },
  {
    {6, 5, 7, 3},
    {9, 9, 12, 6},
    {6, 9, 12, 6}
  },
  {
    {11, 10, 0, 0},
    {18, 18, 0, 0},
    {15,18,0,0}
  },
  {
    {7, 7, 7, 0},
    {12, 12, 12, 0},
    {6, 15, 12, 0}
  },
  {
    {6, 6, 6, 3},
    {12, 9, 9, 6},
    {6, 12, 9, 6}
  },
  {
    {8, 8, 5, 0},
    {15,12,9,0},
    {6,18,9,0}
  }
};


/* Table B.6: layer3 preemphasis */
int  pretab[21] =
{
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    1, 1, 1, 1, 2, 2, 3, 3, 3, 2
};

/*
  Here are MPEG1 Table B.8 and MPEG2 Table B.1
  -- Layer III scalefactor bands. 
  Index into this using a method such as:
    idx  = fr_ps->header->sampling_frequency
           + (fr_ps->header->version * 3)
*/

struct scalefac_struct sfBandIndex[6] =
{
  { /* Table B.2.b: 22.05 kHz */
    {0,6,12,18,24,30,36,44,54,66,80,96,116,140,168,200,238,284,336,396,464,522,576},
    {0,4,8,12,18,24,32,42,56,74,100,132,174,192}
  },
  { /* Table B.2.c: 24 kHz */                 /* docs: 332. mpg123: 330 */
    {0,6,12,18,24,30,36,44,54,66,80,96,114,136,162,194,232,278, 332, 394,464,540,576},
    {0,4,8,12,18,26,36,48,62,80,104,136,180,192}
  },
  { /* Table B.2.a: 16 kHz */
    {0,6,12,18,24,30,36,44,54,66,80,96,116,140,168,200,238,284,336,396,464,522,576},
    {0,4,8,12,18,26,36,48,62,80,104,134,174,192}
  },
  { /* Table B.8.b: 44.1 kHz */
    {0,4,8,12,16,20,24,30,36,44,52,62,74,90,110,134,162,196,238,288,342,418,576},
    {0,4,8,12,16,22,30,40,52,66,84,106,136,192}
  },
  { /* Table B.8.c: 48 kHz */
    {0,4,8,12,16,20,24,30,36,42,50,60,72,88,106,128,156,190,230,276,330,384,576},
    {0,4,8,12,16,22,28,38,50,64,80,100,126,192}
  },
  { /* Table B.8.a: 32 kHz */
    {0,4,8,12,16,20,24,30,36,44,54,66,82,102,126,156,194,240,296,364,448,550,576},
    {0,4,8,12,16,22,30,42,58,78,104,138,180,192}
  }
};

struct scalefac_struct scalefac_band;


FLOAT8 pow20[Q_MAX];
FLOAT8 ipow20[Q_MAX];
FLOAT8 pow43[PRECALC_SIZE];
static FLOAT8 adj43[PRECALC_SIZE];
static FLOAT8 adj43asm[PRECALC_SIZE];
static FLOAT8 ATH_l[SBPSY_l];
static FLOAT8 ATH_s[SBPSY_l];

FLOAT8 ATH_mdct_long[576];
FLOAT8 ATH_mdct_short[192];


/************************************************************************/
/*  initialization for iteration_loop */
/************************************************************************/
void
iteration_init( lame_global_flags *gfp,III_side_info_t *l3_side, int l3_enc[2][2][576])
{
  gr_info *cod_info;
  int ch, gr, i;

  l3_side->resvDrain = 0;

  if ( gfp->frameNum==0 ) {
    for (i = 0; i < SBMAX_l + 1; i++) {
      scalefac_band.l[i] =
	sfBandIndex[gfp->samplerate_index + (gfp->version * 3)].l[i];
    }
    for (i = 0; i < SBMAX_s + 1; i++) {
      scalefac_band.s[i] =
	sfBandIndex[gfp->samplerate_index + (gfp->version * 3)].s[i];
    }

    l3_side->main_data_begin = 0;
    compute_ath(gfp,ATH_l,ATH_s);

    for(i=0;i<PRECALC_SIZE;i++)
        pow43[i] = pow((FLOAT8)i, 4.0/3.0);

    for (i = 0; i < PRECALC_SIZE-1; i++)
	adj43[i] = (i + 1) - pow(0.5 * (pow43[i] + pow43[i + 1]), 0.75);
    adj43[i] = 0.5;


    adj43asm[0] = 0.0;
    for (i = 1; i < PRECALC_SIZE; i++)
      adj43asm[i] = i - 0.5 - pow(0.5 * (pow43[i - 1] + pow43[i]),0.75);

    for (i = 0; i < Q_MAX; i++) {
	ipow20[i] = pow(2.0, (double)(i - 210) * -0.1875);
	pow20[i] = pow(2.0, (double)(i - 210) * 0.25);
    }
  }


  convert_mdct=0;
  reduce_sidechannel=0;
  if (gfp->mode_ext==MPG_MD_MS_LR) {
    convert_mdct = 1;
    reduce_sidechannel=1;
  }
  
  /* some intializations. */
  for ( gr = 0; gr < gfp->mode_gr; gr++ ){
    for ( ch = 0; ch < gfp->stereo; ch++ ){
      cod_info = (gr_info *) &(l3_side->gr[gr].ch[ch]);

      if (cod_info->block_type == SHORT_TYPE)
        {
	  cod_info->sfb_lmax = 0; /* No sb*/
	  cod_info->sfb_smax = 0;
        }
      else
	{
	  /* MPEG 1 doesnt use last scalefactor band */
	  cod_info->sfb_lmax = SBPSY_l;
	  cod_info->sfb_smax = SBPSY_s;    /* No sb */
	}

    }
  }


  /* dont bother with scfsi. */
  for ( ch = 0; ch < gfp->stereo; ch++ )
    for ( i = 0; i < 4; i++ )
      l3_side->scfsi[ch][i] = 0;
}





/* 
compute the ATH for each scalefactor band 
cd range:  0..96db

Input:  3.3kHz signal  32767 amplitude  (3.3kHz is where ATH is smallest = -5db)
longblocks:  sfb=12   en0/bw=-11db    max_en0 = 1.3db
shortblocks: sfb=5           -9db              0db

Input:  1 1 1 1 1 1 1 -1 -1 -1 -1 -1 -1 -1 (repeated)
longblocks:  amp=1      sfb=12   en0/bw=-103 db      max_en0 = -92db
            amp=32767   sfb=12           -12 db                 -1.4db 

Input:  1 1 1 1 1 1 1 -1 -1 -1 -1 -1 -1 -1 (repeated)
shortblocks: amp=1      sfb=5   en0/bw= -99                    -86 
            amp=32767   sfb=5           -9  db                  4db 


MAX energy of largest wave at 3.3kHz = 1db
AVE energy of largest wave at 3.3kHz = -11db
Let's take AVE:  -11db = maximum signal in sfb=12.  
Dynamic range of CD: 96db.  Therefor energy of smallest audible wave 
in sfb=12  = -11  - 96 = -107db = ATH at 3.3kHz.  

ATH formula for this wave: -5db.  To adjust to LAME scaling, we need
ATH = ATH_formula  - 103  (db)
ATH = ATH * 2.5e-10      (ener)

*/
FLOAT8 ATHformula(lame_global_flags *gfp,FLOAT8 f)
{
  FLOAT8 ath;
  f  = Max(0.02, f);
  /* from Painter & Spanias, 1997 */
  /* minimum: (i=77) 3.3kHz = -5db */
  ath=(3.640 * pow(f,-0.8)
       -  6.500 * exp(-0.6*pow(f-3.3,2.0))
       +  0.001 * pow(f,4.0));
  /* convert to energy */
  if (gfp->noATH)
    ath -= 200; /* disables ATH */
  else {
    ath -= 114;    /* MDCT scaling.  From tests by macik and MUS420 code */
    /* ath -= 109; */
  }
#ifdef RH_QUALITY_CONTROL 
  /* purpose of RH_QUALITY_CONTROL:
   * at higher quality lower ATH masking abilities   => needs more bits
   * at lower quality increase ATH masking abilities => needs less bits
   * works together with adjusted masking lowering of GPSYCHO thresholds
   * (Robert.Hegemann@gmx.de 2000-01-30)
   */
  ath -= (4-gfp->VBR_q)*4.0; 
#endif
  ath = pow( 10.0, ath/10.0 );
  return ath;
}
 

void compute_ath(lame_global_flags *gfp,FLOAT8 ATH_l[SBPSY_l],FLOAT8 ATH_s[SBPSY_l])
{
  int sfb,i,start,end;
  FLOAT8 ATH_f;
  FLOAT8 samp_freq = gfp->out_samplerate/1000.0;
#ifdef RH_ATH
  /* going from average to peak level ATH masking
   */
  FLOAT8 adjust_mdct_scaling = 10.0; 
#endif
  

  /* last sfb is not used */
  for ( sfb = 0; sfb < SBPSY_l; sfb++ ) {
    start = scalefac_band.l[ sfb ];
    end   = scalefac_band.l[ sfb+1 ];
    ATH_l[sfb]=1e99;
    for (i=start ; i < end; i++) {
      ATH_f = ATHformula(gfp,samp_freq*i/(2*576)); /* freq in kHz */
      ATH_l[sfb]=Min(ATH_l[sfb],ATH_f);
#ifdef RH_ATH
      ATH_mdct_long[i] = ATH_f*adjust_mdct_scaling;
#endif
    }
    /*
    printf("sfb=%i %f  ATH=%f %f  %f   \n",sfb,samp_freq*start/(2*576),
10*log10(ATH_l[sfb]),
10*log10( ATHformula(samp_freq*start/(2*576)))  ,
10*log10(ATHformula(samp_freq*end/(2*576))));
    */
  }

  for ( sfb = 0; sfb < SBPSY_s; sfb++ ){
    start = scalefac_band.s[ sfb ];
    end   = scalefac_band.s[ sfb+1 ];
    ATH_s[sfb]=1e99;
    for (i=start ; i < end; i++) {
      ATH_f = ATHformula(gfp,samp_freq*i/(2*192));     /* freq in kHz */
      ATH_s[sfb]=Min(ATH_s[sfb],ATH_f);
#ifdef RH_ATH
      ATH_mdct_short[i] = ATH_f*adjust_mdct_scaling;
#endif
    }
  }
}





/* convert from L/R <-> Mid/Side */
void ms_convert(FLOAT8 xr[2][576],FLOAT8 xr_org[2][576])
{
  int i;
  for ( i = 0; i < 576; i++ ) {
    FLOAT8 l = xr_org[0][i];
    FLOAT8 r = xr_org[1][i];
    xr[0][i] = (l+r)*(SQRT2*0.5);
    xr[1][i] = (l-r)*(SQRT2*0.5);
  }
}



/************************************************************************
 * allocate bits among 2 channels based on PE
 * mt 6/99
 ************************************************************************/
void on_pe(lame_global_flags *gfp,FLOAT8 pe[2][2],III_side_info_t *l3_side,
int targ_bits[2],int mean_bits, int gr)
{
  gr_info *cod_info;
  int extra_bits,tbits,bits;
  int add_bits[2]; 
  int ch;

  /* allocate targ_bits for granule */
  ResvMaxBits( mean_bits, &tbits, &extra_bits, gr);
    

  for (ch=0 ; ch < gfp->stereo ; ch ++) {
    /******************************************************************
     * allocate bits for each channel 
     ******************************************************************/
    cod_info = &l3_side->gr[gr].ch[ch].tt;
    
    targ_bits[ch]=tbits/gfp->stereo;
    
    /* allocate extra bits from reservoir based on PE */
    bits=0;
    
    /* extra bits based on PE > 700 */
    add_bits[ch]=(pe[gr][ch]-750)/1.55;  /* 1.4; */
    
    /* short blocks need extra, no matter what the pe */
    if (cod_info->block_type==SHORT_TYPE) 
      if (add_bits[ch]<500) add_bits[ch]=500;
    
    if (add_bits[ch] < 0) add_bits[ch]=0;
    bits += add_bits[ch];
    
    if (bits > extra_bits) add_bits[ch] = (extra_bits*add_bits[ch])/bits;
    if ((targ_bits[ch]+add_bits[ch]) > 4095) 
      add_bits[ch]=4095-targ_bits[ch];

    targ_bits[ch] = targ_bits[ch] + add_bits[ch];
    extra_bits -= add_bits[ch];
  }
}

void reduce_side(int targ_bits[2],FLOAT8 ms_ener_ratio,int mean_bits)
{
int ch;
int numchn=2;
    /*  ms_ener_ratio = 0:  allocate 66/33  mid/side  fac=.33  
     *  ms_ener_ratio =.5:  allocate 50/50 mid/side   fac= 0 */
    /* 75/25 split is fac=.5 */
    /* float fac = .50*(.5-ms_ener_ratio[gr])/.5;*/
    float fac = .33*(.5-ms_ener_ratio)/.5;
    if (fac<0) fac=0;
    
    if (targ_bits[1] >= 125) {
      /* dont reduce side channel below 125 bits */
      if (targ_bits[1]-targ_bits[1]*fac > 125) {
	targ_bits[0] += targ_bits[1]*fac;
	targ_bits[1] -= targ_bits[1]*fac;
      } else {
	targ_bits[0] += targ_bits[1] - 125;
	targ_bits[1] = 125;
      }
    }
    
    /* dont allow to many bits per channel */  
    for (ch=0; ch<numchn; ch++) {
      int max_bits = Min(4095,mean_bits/2 + 1200);
      if (targ_bits[ch] > max_bits) {
	targ_bits[ch] = max_bits;
      }
    }

}

/*************************************************************************** 
 *         inner_loop                                                      * 
 *************************************************************************** 
 * The code selects the best global gain for a particular set of scalefacs */
 
int
inner_loop( lame_global_flags *gfp,FLOAT8 xrpow[576],
	    int l3_enc[576], int max_bits,
	    gr_info *cod_info)
{
    int bits;
    assert( max_bits >= 0 );
    cod_info->global_gain--;
    do
    {
      cod_info->global_gain++;
      bits = count_bits(gfp,l3_enc, xrpow, cod_info);
    }
    while ( bits > max_bits );
    return bits;
}