/*
  RateControl.c
*/

#include <math.h>
#include "SCT_system.h"
#include "SCT_memory.h"
#include "SCT_bit_rw.h"
#include "SCT_encoding.h"
#include "SCT_core.h"
#include "RateControl.h"

#define BIT_TOLERANCE 0.01   /* criterion for convergence of threshold algorithm:
                                actual bits used must differ from the desired budget
				by no more than this proportion */
#define CUTOFF 2             /* Number of successive packets that must be above the
                                marginal error threshold before compression of a
				given frame is cut off */
#define BBAR 2400            /* average bits per frame */
#define MAX_ALLOC 10         /* Maximum multiple of mean number of bits to allocate
                                to any given frame */
#ifdef MIN_MSE				
#define ACCEL 1.0            /* Acceleration factor for negative feedback in iterating */
#else                        /* towards desired bit budget */
#define ACCEL 0.5
#endif

float frame_se;

extern SCT_Block * UseBestSCTPatch ( Heap * error_heap, 
                                     Heap * improvement_heap, 
                                     SCT_Encoding * encoding );

extern void UpdateSCTTreeError ( SCT_Tree * sct_tree );

extern void StoreSCTFrame ( SCT_Frame * sct_frame );

float SCTForestSquareError ( SCT_Forest * sct_forest ) 
{
  int32 x, y;
  float squared_error=0.0f;
  for (y=0;y<sct_forest->rows;y++) for (x=0;x<sct_forest->cols;x++)
    squared_error += (float)sct_forest->sct_tree[y][x].tree_error;
  return squared_error;
}

float EvaluateFrame(SCT_Frame *sct_frame)
#ifdef MIN_MSE
{
  float square_error = SCTForestSquareError (sct_frame->y_forest) + 
                       SCTForestSquareError (sct_frame->u_forest) + 
                       SCTForestSquareError (sct_frame->v_forest);
  /* calculate MSE from all planes */  
  return square_error/(sct_frame->y_forest->source_width*sct_frame->y_forest->source_height*1.5f);
}
#endif
#ifdef MAX_PSNR
{
   float square_error = SCTForestSquareError (sct_frame->y_forest);
   float size = sct_frame->y_forest->source_width*sct_frame->y_forest->source_height*1.5f;
   float psnr = -10.0*log10((255.0*255.0)*size/square_error);
   /* printf("psnr=%f\n", psnr);  */
   return psnr;
}
#endif
#ifdef RATE_CONTROL
#if !defined(MIN_MSE) && !defined(MAX_PSNR)
# error: Must define either MIN_MSE or MAX_PSNR
#endif
#endif /* RATE CONTROL */


float OfflineErrorCalculation(SCT_Frame *sct_frame)
{
  float square_error = SCTForestSquareError (sct_frame->y_forest) + 
                       SCTForestSquareError (sct_frame->u_forest) + 
                       SCTForestSquareError (sct_frame->v_forest);
  
  /* calculate NMSE from all planes ?!?!? */  

  //return square_error = (float)10*log((double)square_error);

  return square_error/(sct_frame->y_forest->source_width*sct_frame->y_forest->source_height*1.5f);
}

/* Compress a single image/frame assigned to the forest structures  */

int32 CompressSCTFrameOfflineRateControl ( SCT_Frame * sct_frame, 
                                           /*int32 bits_avail,*/ 
                                           RateControl *rate_control ) {
  /* Routine variable declarations */

  SCT_Encoding * encoding = sct_frame->sct_encoding;
  SCT_Block * best_block;

  float last_frame_se=0.0f;
  float marginal_error=0.0f, last_marginal_error=255.0f*255.0f*32.0f*32.0f;
  int last_bits_used=0;
  int cutoff_count=0;
  frame_se=0.0f;

  /* Store source as quad and determine error of new source against last reconstructed */
  SCTFrameNewFrame (sct_frame);

  last_frame_se = EvaluateFrame(sct_frame);

  /* Compress until bandwidth used up or no valid patches left */

  do {
    /* Determine patch that gives biggest improvement in error */
    best_block = UseBestSCTPatch (sct_frame->block_error_heap,sct_frame->improvement_heap,encoding);

    /* If patch found, update bits required to encode frame and error of frame */
    if (best_block!=NULL) {
      UpdateSCTTreeBitsUsed (best_block->sct_tree, encoding);
      /* we must add bits to the encoding to test */
      if (last_bits_used < encoding->bits_used) {
         UpdateSCTTreeError (best_block->sct_tree);
         frame_se = EvaluateFrame(sct_frame);
         /* we must be able to improve the frame */
         if (frame_se < last_frame_se) {
            marginal_error = (frame_se-last_frame_se)/(encoding->bits_used-last_bits_used);
            last_marginal_error = marginal_error;
            last_frame_se = frame_se;
            last_bits_used = encoding->bits_used;
         }
       }
     }

     printf("\n%d %f %f", encoding->bits_used, 
                        marginal_error, 
                        rate_control->marginal_error_threshold); 

     if (marginal_error > rate_control->marginal_error_threshold) {
        ++cutoff_count;
     } else {
        cutoff_count = 0;
     }	

} while ((best_block!=NULL) 
          && ((cutoff_count < (CUTOFF+1)) || (encoding->complete<0))
	  && encoding->bits_used < (MAX_ALLOC * BBAR)); 

  /* Update reconstructed output based on patch encodings */

  StoreSCTFrame (sct_frame);
  rate_control->total_bits+=encoding->bits_used;
  rate_control->history[rate_control->total_frames].error = last_frame_se;
  rate_control->history[rate_control->total_frames].bits_allocated = encoding->bits_used;
  rate_control->total_frames++;

  // rate_control->history[rate_control->total_frames].y_psnr = /* can calculate this if required */;

  /* Write encoded frame and return the number of bits used to encode frame */
  return (WriteSCTFrame(sct_frame));
}


/* reset values ready for another rate control */
void RateControlReset(RateControl *r)
{
  r->total_frames=0;
  r->total_bits=0;
}


/* allocate and initialize for rate control */
RateControl * RateControlAllocate(int num_frames, 
                                  int target_bits_per_frame, 
                                  float initial_marginal_error_threshold)
{
  RateControl *temp = (RateControl*)AllocateMemory(sizeof(RateControl),1);
  RateControlReset(temp);
  temp->history_size             = num_frames;
  temp->target_bits_per_frame    = target_bits_per_frame;
  temp->marginal_error_threshold = initial_marginal_error_threshold;
  temp->history                  = AllocateMemory(sizeof(RateControlHistory),temp->history_size);
  return temp;
}

/* free rate control memory */
void RateControlFree(RateControl *r)
{
  FreeMemory(r->history);
  FreeMemory(r);
}

#ifdef RATE_CONTROL
/* Modified by Allin Cottrell */
float RateControlCalculateNewThreshold(RateControl *r)
{
  extern float Tmin, Tmax;
  int diff_bit_allocation = (r->target_bits_per_frame*r->total_frames)-r->total_bits;
  float target = r->target_bits_per_frame*r->total_frames;
  float diff_bits = target - r->total_bits;
  float diff_bits_percent = (target - r->total_bits)/target;
  float new_threshold;
  float T = r->marginal_error_threshold;
  if (diff_bits > 0 && T > Tmin) Tmin = T; 
  if (diff_bits < 0 && T < Tmax) Tmax = T; 
  new_threshold = (1.0 - ACCEL * diff_bits_percent) * T;
  if (new_threshold < Tmin) new_threshold = 0.5 * (T + Tmin);
  if (new_threshold > Tmax) new_threshold = 0.5 * (T + Tmax);

  printf("B-hat: %d proportional error: %f threshold: %f\n\n", 
       r->total_bits, diff_bits_percent, r->marginal_error_threshold);

{
  FILE * out = fopen("T_smoother_vs_B.txt", "ab");
  fprintf(out, "%f %d\n", r->marginal_error_threshold, r->total_bits);
  fclose(out);
}

  if (fabs(diff_bits_percent) > BIT_TOLERANCE) {
     return new_threshold;
  } else {
     return 1.0;
  }

}
#endif /* RATE_CONTROL */