#include "SCT_system.h"
#include "SCT_memory.h"
#include "SCT_image_rw.h"
#include "SCT_movie_rw.h"
#include "SCT_bit_rw.h"
#include "SCT_encoding.h"
#include "SCT_core.h"
#include "SCT_core_mc.h"
#include "SCT_core_vq.h"
#include "SCT_core_wavelet.h"
#include "SCT_initialize.h"
#include <math.h>

#include "RateControl.h"

#ifdef _DEBUG
# include <crtdbg.h>
#endif

#define NHVQ
#define BBAR 2400      /* Average bits per frame */

/* #include "avi2sct.h" */

RateControl *gRATECONTROL;

long total_bits_this_run = 0;
#ifdef RATE_CONTROL
float Tmax = 0.0;
float Tmin = -100.0;
#endif /* RATE_CONTROL */

void FlickerFilterLine(unsigned char *current_line, unsigned char *prev_line, int length)
{
  int i;

  for (i=0; i<length; i++) {
    int current_value = current_line[i];
    int prev_value    = prev_line[i];
    int diff          = current_value - prev_value;

    if (diff<0) diff=-diff;

    if      (diff <2) current_line[i] = prev_value;
    else if (diff <5) current_line[i] = (prev_value*3+current_value)>>2;
//    else if (diff<10) prev_line[i]    = (prev_value+current_value*3)>>2;
//    else if (diff<10) prev_line[i]    = 255;
    else              prev_line[i]    = current_value;   
  }
}

unsigned char filter_memory[320*240];

/* ================================================================================================================== */

void CompressTestSeq ( char * source_file, char * output_sct, 
                       int32 first_frame, int32 max_frame, int32 bits_per_frame, int32 frame_step,
                       int32 real_time,
                       float target_mse, char * first_frame_file ) {

  SCT_Frame * frame;
  SCT_Initialization sct_init;
  YUV_File * yuv_file, * yuv_init;
  int32 f,bits,budget,l,m;
  Boolean compress=TRUE, padded=FALSE;
  float psnr, bandwidth, delay=(float)0;
  IO_Stream io;

  int total_frames = 0;
  float total_psnr = 0.0f;


  /* Open source and first frame */

  if ((yuv_file=OpenYUVMovie(source_file))==NULL) FatalError ("CompressTestSeq","%s",GetWarning());

  if (first_frame_file!=NULL) {
    if ((yuv_init=OpenYUVMovie(first_frame_file))==NULL) 
      FatalError ("CompressTestSeq","Cannot read first frame file %s",first_frame_file);
  } else {
    yuv_init = NULL;
  }

  /* Initialize CODEC */

  sct_init.do_compression = TRUE;
  sct_init.height         = yuv_file->height;
  sct_init.width          = yuv_file->width;
  sct_init.format         = yuv_file->format;
  sct_init.block_levels   = 4;
  sct_init.stream_id      = output_sct;
  sct_init.frame_delay    = (1000*(abs(frame_step)))/yuv_file->frame_rate;
  sct_init.real_time      = real_time & 1;

  io.OpenStreamForReading = OpenBitFileRead;
  io.OpenStreamForWriting = OpenBitFileWrite;
  io.ReadDataFromStream   = ReadBitFile;
  io.WriteDataToStream    = WriteBitFile;
  io.SeekPositionInStream = BitFileSeek;
  io.CloseDataStream      = CloseBitFile;

  if (max_frame==0) max_frame=yuv_file->frames;

  if ((frame=InitializeCODEC(&io,&sct_init))==NULL) {
    fprintf (stderr,"Unable to initialize CODED - %s\n",GetWarning());
    FreeAllMemory ();
    exit (0);
  }

  bandwidth = (float)bits_per_frame*(float)yuv_file->frame_rate/(float)(abs(frame_step));

  /* Preset to first frame */

  if (yuv_init!=NULL) ReadYUVFrame (yuv_init,first_frame,frame->yuv_output);

  budget = bits_per_frame;

  if (frame_step<0) {
    padded     = TRUE;
    frame_step = 0-frame_step;
  }

  frame->sct_encoding->first_frame = FALSE;
  frame->sct_encoding->lock_all_methods = TRUE;

  EnableAllEncodingMethods (frame->sct_encoding);

  for (f=(yuv_init==NULL?first_frame:first_frame+1);(compress)&&(f<max_frame);f++) {

    int32 tot_time;

    if ((f%frame_step)==0) {

      if (compress=(ReadYUVFrame(yuv_file,f,frame->yuv_source)!=NULL)) {
        if (real_time&2) FlickerFilterLine(frame->yuv_source->y_plane[0], filter_memory, frame->yuv_source->height*frame->yuv_source->width);
        if ((delay-=(float)bits_per_frame/bandwidth)<(float)0) delay=(float)0;
        if (target_mse<(float)1.0) {
          fprintf (stderr,"Compress frame %3d with %5d bits: ",f,budget);
          ReportTimeUsed();

#ifdef RATE_CONTROL
          total_bits_this_run += bits = 
	     CompressSCTFrameOfflineRateControl (frame,gRATECONTROL);
#else
          bits = CompressSCTFrame (frame,budget);
#endif /* RATE CONTROL */

          tot_time = ReportTimeUsed();

          fprintf (stderr,"%5d bits written, %0.3f secs - ",bits,(float)tot_time/1000000.0);
        } else {
          fprintf (stderr,"Compressing frame %d with %d bits : ",f,budget);
          fprintf (stderr,"%d bits written - ",bits=CompressSCTQualityFrame (frame,budget,target_mse));
        }

        if (bits<0) {
          fprintf (stderr,"WriteSCTFrame fail - %s\n",GetWarning());
          FreeAllMemory();
          exit (0);
        }


        if (real_time&2) ReadYUVFrame(yuv_file,f,frame->yuv_source);

        budget = budget-bits+bits_per_frame;
        delay += (float)bits/bandwidth;

        psnr   = GetImagePSNR(frame->yuv_output->y_plane,frame->yuv_source->y_plane);
      
        total_psnr += psnr;
        total_frames++;

        fprintf (stderr,"PSNR=%0.3f: delay = %0.4f se = %f\n",
	   psnr,delay, OfflineErrorCalculation(frame));
        printf ("%d %f %f\n",f,psnr,(float)tot_time/1000000.0);

      } else {
        fprintf (stderr,"%s - %s\n",GetWarning());
      }

    } else if (padded) {

      if (compress=(ReadYUVFrame(yuv_file,f,frame->yuv_source)!=NULL)) {
        psnr   = GetImagePSNR(frame->yuv_output->y_plane,frame->yuv_source->y_plane);
        fprintf (stderr,"PSNR of Y frame %d is %f\n",f,psnr);
        printf ("%d %f 0 0\n",f,psnr);
      }
    }
  }

  CloseYUVMovie (yuv_file);

  if (yuv_init!=NULL) CloseYUVMovie (yuv_init);

  TerminateCODEC (frame);

  if (total_frames) {
    printf("%f\n", total_psnr/total_frames);
    fprintf(stderr, "mean psnr = %f\n", total_psnr/total_frames);
  }
}

/* ================================================================================================================== */

void DecompressTestSeq ( char * fname ) {

  SCT_Frame * frame;
  int32 x,f,t=0;
  RGB_Image * rgb;
  SCT_Initialization sct_init;
  IO_Stream io;

  /* Decompress sequence */

  sct_init.do_compression = FALSE;
  sct_init.stream_id      = fname;

  io.OpenStreamForReading = OpenBitFileRead;
  io.OpenStreamForWriting = OpenBitFileWrite;
  io.ReadDataFromStream = ReadBitFile;
  io.WriteDataToStream = WriteBitFile;
  io.SeekPositionInStream = BitFileSeek;
  io.CloseDataStream = CloseBitFile;

  printf ("Initial memory usage = %d bytes\n",GetMemoryUsage());

  if ((frame=InitializeCODEC(&io,&sct_init))==NULL) {
    fprintf (stderr,"Unable to start CODEC - %s\n",GetWarning());
    exit (0);
  }

  printf ("Memory usage after init = %d bytes\n",GetMemoryUsage());
  x=0;
  frame->sct_encoding->lock_all_methods = TRUE;

  do {
    f=DecompressSCTFrame (frame);
    if (f>0) {
/*
      
      if ((x%10)==0 || 1) {
      rgb=YUV2RGBImage (frame->yuv_output);
      sprintf (fname,"out.bmp",x);
      printf ("Writing frame %s - %d bits read\n",fname,f);
      WriteRGBImageAsBMP (fname,rgb);

      FreeRGBImage (rgb);

    }
*/
      FILE *out = fopen("out.yuv", "a+b");
      fwrite(frame->yuv_output->y_plane[0], frame->yuv_output->width*frame->yuv_output->height, 1, out);
      fwrite(frame->yuv_output->u_plane[0], (frame->yuv_output->width*frame->yuv_output->height)>>2, 1, out);
      fwrite(frame->yuv_output->v_plane[0], (frame->yuv_output->width*frame->yuv_output->height)>>2, 1, out);
      fclose(out);

      x++;
    }
    if (f==-1) { fprintf (stderr,"SYNC ERROR!\n"); f=1; }
  } while (f>0);

  if (f!=0) fprintf (stderr,"Fail on Decompress (return value %d) - %s\n",f,GetWarning());

  TerminateCODEC (frame);
}


/* ================================================================================================================== */

void Usage ( char * name ) {

  fprintf (stderr,"Compress:   %s ",name);
  fprintf (stderr,"<source> <output> <frames> <bits/frame> <skip> [<target_mse>] [<first_frame_file>]\n");
  fprintf (stderr,"Decompress: %s <sct_file>\n",name);
}
  
/* ================================================================================================================== */

int main ( int argc, char ** argv ) {

#ifdef RATE_CONTROL
#ifdef MIN_MSE
  float T = -0.01;
#else
  float T = -0.0002;  /* for MAX_PSNR */
#endif     
  int iteration = 0;
  gRATECONTROL = RateControlAllocate(100,BBAR,T); 

  do {
    printf("iteration %d\n", iteration);
    ++iteration;
    CompressTestSeq (argv[1],argv[2],0,atoi(argv[3]),atoi(argv[4]),atoi(argv[5]),FALSE,(float)0,NULL); 
    T = RateControlCalculateNewThreshold(gRATECONTROL);
    gRATECONTROL->marginal_error_threshold = T;
    RateControlReset(gRATECONTROL);
    printf("T=%f\n", T);
  } while (T < 0 && iteration < 15);

  return 0; 
#endif /* RATE_CONTROL */

  switch (argc) {
    case 2  : DecompressTestSeq (argv[1]); break;
    case 6  : CompressTestSeq (argv[1],argv[2],0,atoi(argv[3]),atoi(argv[4]),atoi(argv[5]),FALSE,(float)0,NULL); break;
    case 7  : CompressTestSeq (argv[1],argv[2],atoi(argv[3]),atoi(argv[4]),atoi(argv[5]),atoi(argv[6]),FALSE,(float)0,NULL); break;
    case 8  : CompressTestSeq (argv[1],argv[2],atoi(argv[3]),atoi(argv[4]),atoi(argv[5]),atoi(argv[6]),atoi(argv[7]),(float)0,NULL); break;
    default : Usage (argv[0]); break;

  }

  fprintf (stderr,"CODEC Terminated - %d bytes deallocated\n",FreeAllMemory());
  return 0;
}