/*---------------------------------------------------------------------------*/
// Baseline Wavelet Transform Coder Construction Kit
//
// Geoff Davis
// gdavis@cs.dartmouth.edu
// http://www.cs.dartmouth.edu/~gdavis
//
// Copyright 1996 Geoff Davis 9/11/96
//
// Permission is granted to use this software for research purposes as
// long as this notice stays attached to this software.
//
/*---------------------------------------------------------------------------*/

// File: packetize.cc (IB)
// Modification of encode.cc

#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <iostream.h>
#include <fstream.h>
#include "transform.h"
#include "coeffset.h"
#include "allocator.h"
#include "quantizer.h"
#include "packet.h"
#include "conceal.h"


#define RASTER 0  // read raster format - for video

/*---------------------------------------------------------------------------*/
void compress       (Image *image, Wavelet *wavelet, int nStages, 
		     int capacity, Real p, Real *weight,
		     int paramPrecision, int budget, int nQuant, 
		     Real minStepSize, char *outfilename, int monolayer);

void opensfile (char *sfilename, ofstream *sfile, int i);
void openhfile (char *hfilename, ofstream *hfile);
Real variance  (CoeffSet *coeff);
void allocCover(int nSets, int *n, Real *D, Real *R, Real Rmax, Real p, int N);
Real power(Real x, int n);
/*---------------------------------------------------------------------------*/

int main (int argc, char **argv)
{
  char *program = argv[0];
  int monolayer = FALSE;    // TRUE for non-embedded uniform quantizer
                            // FALSE for multilayer quantizer
                            
#ifdef PGM
  if (argc < 4)  {
    fprintf (stderr, 
	     "Usage: %s image output ratio [mono_layer(FALSE)]\n",
	     program);
    fprintf (stderr, 
	     "image: image to be compressed (in PGM format)\n");
    fprintf (stderr, 
	     "output: name of compressed image\n");
    fprintf (stderr, 
	     "ratio: compression ratio\n");
    exit(0);
  }
  char *infile_name = argv[1];
  char *outfile_name = argv[2];
  Real ratio = atof(argv[3]);
  if (argc == 5) {
    monolayer = atoi(argv[4]);
    if((monolayer != 0) && (monolayer != 1)){
          fprintf (stderr, 
	     "mono_layer: 1/0\n");
	  exit(0);
    }
  }
    
  // Load the image to be coded
  Image *image = new Image (infile_name);

  int budget = (int)((Real)(image->hsize*image->vsize)/ratio);  // (assumes 8 bit pixels)
  printf ("Reading %d x %d image %s, writing %s\nCompression ratio %g:1\n", 
	  image->hsize, image->vsize, infile_name, outfile_name, ratio);

#else

  if (argc < 6)  {
    fprintf (stderr, 
	     "Usage: %s [image][width][height][output][ratio]\n",
	     program);
    if (!RASTER)     
      fprintf (stderr, 
	       "image: image to be compressed (in RAW format)\n");
    else
      fprintf (stderr, 
	       "image: image to be compressed (in raster format)\n");
    fprintf (stderr, 
	     "width, height: width and height of image to be compressed\n");
    fprintf (stderr, 
	     "output: name of compressed image\n");
    fprintf (stderr, 
	     "ratio: compression ratio\n");
    exit(0);
  }
  char *infile_name = argv[1];
  int hsize = atoi(argv[2]);
  int vsize = atoi(argv[3]);
  char *outfile_name = argv[4];
  Real ratio = atof(argv[5]);
  int budget = (int)((Real)(hsize*vsize)/ratio);  // (assumes 8 bit pixels)
  printf ("Reading %d x %d image %s, writing %s\nCompression ratio %g:1\n", 
	  hsize, vsize, infile_name, outfile_name, ratio);

  // Load the image to be coded
  //if (!RASTER)
  //  Image *image = new Image (infile_name, hsize, vsize);
  //else
  // Currently using frame 0 - change for loop
    Image *image = new Image (infile_name, hsize, vsize, 0);	  
  
#endif


  // Create a new wavelet from the 7/9 Antonini filterset
  Wavelet *wavelet = new Wavelet (&Antonini);
  // Adelson & Simoncelli 9-tap QMF
  //Wavelet *wavelet = new Wavelet (&Adelson);

  // Coding parameters
  Real p = 2.0;             // exponent for L^p error metric
  int nStages = 4;          // # of stages in the wavelet transform
  int capacity = 512;       // capacity of histogram for arithmetic coder  
  int paramPrecision = 4;   // precision for stored quantizer parameters
  int nQuant = 10;          // # of quantizers to examine for allocation
  Real minStepSize = 0.05;  // smallest quantizer step to consider
  //H monolayer = FALSE;    // TRUE for non-embedded uniform quantizer
  //H                          // FALSE for multilayer quantizer

  Real *weight =            // perceptual weights for coefficient sets
    new Real[3*nStages+1];
  // for now give all sets equal weight
  for (int i = 0; i < 3*nStages+1; i++)
    weight[i] = 1.0;

  compress (image, wavelet, nStages, capacity, p, weight, 
	    paramPrecision, budget, nQuant, minStepSize, outfile_name, 
	    monolayer);

  delete [] weight;
  delete image;
  delete wavelet;
  return 0;
}

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
// Compress an image and save to a file
//
// Image *image          Image to be compressed
// Wavelet *wavelet      Wavelet to use for transform
// int nStages           # of stages to use in transform
// int capacity          Capacity of histograms for arithmetic coder
// Real p                Exponent for L^p error metric
// Real *weight          Perceptual weights for subbands
// int paramPrecision    Precision for storing quantizer parameters --
//                          precision = n means params are stored with
//                          accuracy 2^(-n)
// int budget            Total # of bytes for compressed image
// int nQuant            # of different quantizer resolutions to
//                          consider for each subband.  Step size for
//                          quantizer k is roughly (max coeff - min coeff)/2^k
// Real minStepSize      # minimum quantizer step size to consider --
//                          prevents too much effort from being
//                          expended on fine-scale subbands
// char *filename        name for compressed file
// int monolayer         TRUE for non-embedded quantizer, FALSE for embedded

/*---------------------------------------------------------------------------*/
void compress (Image *image, Wavelet *wavelet, int nStages, 
	       int capacity, Real p, Real *weight,
	       int paramPrecision, int budget, int nQuant, 
	       Real minStepSize, char *outfilename, int monolayer)
{
  int i, j, nSets = 1;
  char hfilename[80];
  
  // Compute the wavelet transform of the given image
  WaveletTransform *transform = 
    new WaveletTransform (wavelet, image, nStages); 

  // For each subband allocate a CoeffSet, an error metric, an
  // EntropyCoder, and a Quantizer
  nSets = transform->nSubbands;
  CoeffSet **coeff = new CoeffSet* [nSets];
  ErrorMetric **err = new ErrorMetric* [nSets];
  EntropyCoder **entropy = new EntropyCoder* [nSets];
  Quantizer **quant = new Quantizer* [nSets];

  //printf("Quantizers defined\n");
  
  for (i = 0; i < nSets; i++) {
    // Use an L^p error metric for each subband
    err[i] = new LpError (p);

    if (monolayer) {
      // Use uniform quantizer and single layer escape coder for each
      //   subband
      entropy[i] = new EscapeCoder  (capacity);
      // Assume all subbands have pdf's centered around 0 except the
      //   low pass subband 0
      quant[i]   = new UniformQuant ((MonoLayerCoder *)entropy[i],
				     paramPrecision, i != 0, nStages, err[i]);
    } else {
      // Use a layered quantizer with dead zones and a layered entropy
      //   coder for each subband
      entropy [i] = new LayerCoder (nQuant, i != 0, capacity);
      // Assume all subbands have pdf's centered around 0 except the
      //   low pass subband 0
      quant [i] = new LayerQuant ((MultiLayerCoder *)entropy[i],
				  paramPrecision, i != 0, nQuant, nStages, err[i]);				  
    }

    //printf("Quantizer created\n");
    
    // Partition the wavelet transformed coefficients into subbands --
    //   each subband will have a different quantizer
    coeff[i]   = new CoeffSet (transform->subband(i),
			       transform->subbandSize[i], quant[i]);

    // For each subband determine the rate and distortion for each of
    //    the possible nQuant quantizers
    coeff[i]->getRateDist (nQuant, minStepSize, i);
  }
  
  Allocator *allocator = new Allocator ();
  // Use rate/distortion information for each subband to find bit
  //    allocation that minimizes total (weighted) distortion subject
  //    to a byte budget
  
  budget -= nSets * 4;  // subtract off approximate size of header info
  //budget -= nSets * 40;   // 65 for peppers, 40 for lena to match PZW
  
  allocator->optimalAllocate (coeff, nSets, budget, TRUE, weight);
  printf ("Target rate = %d bytes\n", budget);
  // Display the resulting allocation
  allocator->print (coeff, nSets);


  // ------------------ domain covering -----------------------------
  Real *R = new Real [nSets];
  Real *D = new Real [nSets];
  int  *n = new int  [nSets];
  for (i = 0; i < nSets; i++)
  {
    D[i] = variance(coeff[i]);
    R[i] = (coeff[i]->rate[allocator->precision[i]])/8.0; // Bytes
    n[i] = 1;
  }
  
  Real ploss = 0.25;
  Real Rmax = 16384.0; // equivalent to 0.5 bpp
  //allocCover(nSets, n, D, R, Rmax, ploss, 16);
  
  printf ("subband   repetitions  var\n");
  for (i = 0; i < nSets; i++)
    printf("%d         %d            %f\n", i, n[i], D[i]);
    
  delete D;
  delete R;
  //delete n; // move this delete later
  // ------------------ end domain covering -------------------------
  
  // Create filename for header file
  sprintf(hfilename, "%c%s",'h', outfilename);
  //printf("Writing header file %s\n",hfilename);
  
  // Open output file for header
  //ofstream hfile (hfilename, ios::out | ios::trunc | ios::bin);
  ofstream hfile (hfilename, ios::out | ios::trunc | ios::binary);
  if (!hfile) {
    error ("Unable to open file %s", hfilename);
  }
  
  //printf("hfile opened\n");
  
  // Create I/O interface for header
  Encoder *hencoder = new Encoder (hfile);
  
  //printf("hencoder created\n");
  
  // Write header data: horizontal and vertical image size 
  hencoder->writePositive(image->hsize);
  hencoder->writePositive(image->vsize);
  
  // write number of repetitions
  for (i = 0; i < nSets; i++)
    hencoder->writePositive(n[i]);
  
  //printf("Image dimensions written\n");
  
  // and subband headers
  for (i = 0; i < nSets; i++)
    coeff[i]->writeHeader (hencoder, allocator->precision[i], i);

  //printf("header written\n");
  
  // Flush bits from header I/O interface and close file
  hencoder->flush ();
  delete hencoder;
  hfile.close ();
  //printf("hfile closed\n");
  // -------------- End of added ------------------------------------

  // ----------------- Added --------------------------------
  const int nP = 16;
  Encoder **sencoder = new Encoder* [nP];
  char sfilename[nP][80];            // Change 16 to number of packets
  ofstream *sfile = new ofstream [nP];
  
  for (i = 0; i < nP; i++)
  {
    // Create output filename
    sprintf(sfilename[i], "%c%s%03d", 's', outfilename, i);
    //printf("Writing subband file %s\n", sfilename[i]);
    
    opensfile(sfilename[i], sfile, i);
    
    if (!sfile[i]) {
      error ("Unable to open file %s\n", sfilename[i]);
    }
    //printf("Packet file %s opened\n", sfilename[i]);
    
    // Create I/O interface
    sencoder[i] = new Encoder (sfile[i]);
    //printf("Interface (sencoder) created\n");
  }
  
  //printf("Packet files opened and I/O buffers created\n");
  
  // define packet objects - 05/17/01
  PacketEnc **pkt = new PacketEnc* [nP];
  EntropyCoder **entLL = new EntropyCoder* [nP];
  EntropyCoder **entHF = new EntropyCoder* [nP];
  for (i = 0; i < nP; i++)
  {
    entLL[i] = new LayerCoder (nQuant, 0, capacity);
    entHF[i] = new LayerCoder (nQuant, 1, capacity);
    pkt[i] = new PacketEnc((MultiLayerCoder *)entLL[i], 
                           (MultiLayerCoder *)entHF[i]);
    if (pkt[i] == NULL)
    {
      printf("Error in defining PacketEnc no. %d in packetize.cc\n", i);
      exit(1);
    }
  }
  //printf("packet objects defined\n");
  
  for (i = 0; i < nSets; i++) 
  {
    // Code and write coefficients
    for (j = 0; j < n[i]; j++)
      coeff[i]->encode (sencoder, allocator->precision[i], i, pkt, j);
    //printf("Coefficients of subband %d written into output buffer\n", i);
  }  
  
    
  //printf("About to write data into files\n");
  
  for (i = 0; i<nP; i++)
  {
    // Flush bits from arithmetic coder and close file
    sencoder[i]->flush ();
    //printf("Flush no. %d\n",i);
    sfile[i].close ();
    //printf("Output file %s closed\n", sfilename[i]);
  }
  
  delete sencoder;
  delete n;
  
  for (i = 0; i < nP; i++)
  {
    delete entLL[i];
    delete entHF[i];
    delete pkt[i];
  }
  
  delete entLL;
  delete entHF;
  delete pkt;
  
  //printf("Encoders deleted\n");
  
  // ------------- End of added -----------------------------


  // Clean up
  for (i = 0; i < nSets; i++) {
    delete err[i];
    delete entropy[i];
    delete quant[i];
    delete coeff[i];
  }
  delete [] err;
  delete [] entropy;
  delete [] quant;
  delete [] coeff;
  delete allocator;
  delete transform;
}

// ----------------------- Added -------------------------------
void opensfile (char sfilename[80], ofstream *sfile, int i)
{
  //sfile[i].open(sfilename, ios::out | ios::trunc | ios::bin);
	sfile[i].open(sfilename, ios::out | ios::trunc | ios::binary);
}  

void openhfile (char hfilename[80], ofstream *hfile)
{
//  *hfile.open(hfilename, ios::out | ios::trunc | ios::bin);
}  

/*---------------------------------------------------------------------------*/

/*---------------------------------------------------------------------------*/
/*                            variance()                                     */
/*---------------------------------------------------------------------------*/
Real variance(CoeffSet *coeff)
{
  int i;
  Real mean, mean2;
  
  mean = mean2 = 0.0;
  for (i = 0; i < coeff->nData; i++)
  {
    mean  += coeff->data[i];
    mean2 += (coeff->data[i])*(coeff->data[i]);
  }
  
  mean  /= coeff->nData;
  mean2 /= coeff->nData;
  
  return (mean2-mean*mean);
}

/*---------------------------------------------------------------------------*/
/*                            allocCover()                                   */
/*---------------------------------------------------------------------------*/
/* nSets = # of subbands                                                     */
/* n[i]  = # of repetitions of subband i                                     */
/* D[i]  = variance of subband i                                             */
/* R[i]  = estimated rate of subband i                                       */
/* p     = loss probability                                                  */
/* N     = # of packets                                                      */
/*---------------------------------------------------------------------------*/
void allocCover(int nSets, int *n, Real *D, Real *R, Real Rmax, Real p, int N)
{
  int i, j, Nrzsubs, Ncompleted;
  Real Rc, Dc, Rmin, Bmax;
  Real *B = new Real [nSets];
  Real *Btemp = new Real [nSets];
  Real *Rtemp = new Real [nSets];
  
  // start off with all n[i]=1
  
  //find current rate and distortion
  Dc = 0.0; Rc = 0.0;
  for (i = 0; i < nSets; i++)
  {
    Rc += (Real)n[i]*R[i];
    if (n[i] > floor(p*((Real)N)))
    {
      Dc += D[i];  // subband received for loss = p
      B[i] = 0.0;  // would not benefit by repeating this subband again
    }
    else
    {
      Dc += (1.0 - power(p,n[i]))*D[i];
      if (n[i]+1 > floor(p*((Real)N))) 
      {
        if (floor(R[i]) > 0)
          B[i] = power(p,n[i])*D[i]; 
        else
          B[i] = 0.0;
      }
      else
      {
        if (floor(R[i]) > 0)
          B[i] = (power(p,n[i]) - power(p,n[i]+1))*D[i];
        else
          B[i] = 0.0;
      }
    }
  }
  
  for (i = 0; i < nSets; i++)
  {
    //Bmax = select(i,nSets,B);
    printf("%d   %f   %f\n", i, R[i], B[i]);
  }
  
  // now loop the following: increase n[i] for highest feasible B[i] until full
  // choose the smallest non-zero rate of the subbands
  Rmin = 0.0; i = 0;
  while (floor(Rmin) == 0  &&  i < nSets)
  {
    for (j = 0; j < nSets; j++)
      Rtemp[j] = R[j];
    // note: 'select' treats array as ascending
    Rmin = select(i, nSets, Rtemp);
    i++;
  }
  Nrzsubs = i-1;  // number of rate-zero subbands
  Ncompleted = Nrzsubs; // these many subbands need not be considerd further
  
  //printf("Nrzsubs = %d, Rmin = %f\n", Nrzsubs, Rmin);
  
  while (Rmax - Rc >= Rmin  &&  Ncompleted < nSets)
  {
    j = nSets-1;  // j is the number on the ascending list
    do
    {
      for (i = 0; i < nSets; i++)
        Btemp[i] = B[i];
        
      Bmax = select(j, nSets, Btemp);
      
      //printf("Bmax = %f  ",Bmax);
      j--;
      // find index of Bmax
      i = 0;
      //while (abv(B[i]-Bmax) > 0.001 && i < nSets-1)
      while (floor(B[i]) != floor(Bmax)  &&  i < nSets-1)
      {
        i++;
      }
      //printf("i = %d  ",i);
    }
    while (j >= Ncompleted && (R[i] > Rmax - Rc  ||  n[i] > floor(p*((Real)N)) 
                               || floor(R[i]) == 0));
    
    //printf("j=%d  ",j);
    
    // index of the best feasible subband found
    n[i]++;
    Rc += R[i];
    Dc += B[i];
    //recalculate B[i]
    if (n[i] > floor(p*((Real)N))) 
    {
      B[i] = 0.0;
      Ncompleted++;
    }
    else if (n[i]+1 > floor(p*((Real)N))) 
    {
      B[i] = power(p,n[i])*D[i]; 
    }
    else
    {
      B[i] = (power(p,n[i]) - power(p,n[i]+1))*D[i];
    } 
    //printf("i = %d, Ncompleted = %d;   ",i,Ncompleted);   
  }
  
  printf("Total rate = %f bytes, D = %f\n", Rc, Dc);
  
  delete B;
  delete Btemp;
  delete Rtemp;
}

/*---------------------------------------------------------------------------*/
/*                            power()                                        */
/*---------------------------------------------------------------------------*/
Real power(Real x, int n)
{
  int i;
  Real res = 1.0;
  
  for(i = 1; i <= n; i++)
    res *= x;  
    
  return res;
}
