/**
 * @brief Apply (partial) histogram equalization to the pfs stream
 * 
 * ---------------------------------------------------------------------- 
 * Copyright (C) 2007 Ed Brambley
 * 
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 * ---------------------------------------------------------------------- 
 *
 * @author Ed Brambley, <E.J.Brambley@damtp.cam.ac.uk>
 *
 */

#include <iostream>

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <getopt.h>

#include <pfs.h>

#define PROG_NAME "pfsequalize"

class QuietException 
{
};


static void histogramEqualizeY( const int cols, const int rows, float* const Y, const bool verbose);
static void histogramEqualizeRGB( const int cols, const int rows, float* const R, float* const G, float* const B, const float saturationFactor, const bool verbose);

static void printHelp()
{
  fprintf( stderr, PROG_NAME " [--verbose] [--saturation <NUM>] [--help]\n");
}

static void mainBody( int argc, char* argv[] )
{
  pfs::DOMIO pfsio;

  bool verbose = false;
  float saturationFactor = 1.0f;

  static struct option cmdLineOptions[] = {
    { "help", no_argument, NULL, 'h' },
    { "verbose", no_argument, NULL, 'v' },
    { "saturation", required_argument, NULL, 's' },
    { NULL, 0, NULL, 0 }
  };

  int optionIndex = 0;
  while( 1 ) {
    int c = getopt_long (argc, argv, "hvs:", cmdLineOptions, &optionIndex);
    if( c == -1 ) break;
    switch( c ) {
    case 'h':
      printHelp();
      throw QuietException();
    case 'v':
      verbose = true;
      break;
    case 's':
      saturationFactor = (float)strtod( optarg, NULL );
      if( saturationFactor <= 0.0f )
        throw pfs::Exception("incorrect saturation factor; saturation must be positive.");
      break;
    case '?':
      throw QuietException();
    case ':':
      throw QuietException();
    }
  } 

  while( true ) {
    pfs::Frame *frame = pfsio.readFrame( stdin );
    if( frame == NULL ) break; // No more frames
        
    pfs::Channel *X, *Y, *Z;
    frame->getXYZChannels( X, Y, Z );
    const int cols = frame->getWidth();
    const int rows = frame->getHeight();

    if( X != NULL )
      {
	// Color, XYZ
    
	pfs::transformColorSpace( pfs::CS_XYZ, X, Y, Z, pfs::CS_SRGB, X, Y, Z );
	
	histogramEqualizeRGB(cols, rows, X->getRawData(), Y->getRawData(), Z->getRawData(), saturationFactor, verbose);

	pfs::transformColorSpace( pfs::CS_SRGB, X, Y, Z, pfs::CS_XYZ, X, Y, Z );
    
      }
    else if ((Y = frame->getChannel( "Y" )) != NULL)
      {
	// Luminance only
	
	histogramEqualizeY(cols, rows, Y->getRawData(), verbose);
      }
    else
      {
	throw pfs::Exception( "Missing X, Y, Z channels in the PFS stream" );
      }

    // frame->getTags()->setString("LUMINANCE", "DISPLAY");
    
    pfsio.writeFrame( frame, stdout );
    pfsio.freeFrame( frame );        
  }
}


struct hist_data
{
  float size;
  int index;
};


static int hist_data_comp_size(const void* const v1, const void* const v2)
{
  if (((struct hist_data*) v1)->size < ((struct hist_data*) v2)->size)
    return -1;
  
  if (((struct hist_data*) v1)->size > ((struct hist_data*) v2)->size)
    return 1;

  return 0;
}


static int hist_data_comp_index(const void* const v1, const void* const v2)
{
  return ((struct hist_data*) v1)->index - ((struct hist_data*) v2)->index;
}


void histogramEqualizeY( const int cols, const int rows, float* const Y, const bool verbose )
{
  const int n = cols * rows;
  
  // Allocate memory
  struct hist_data* const hist = (struct hist_data*) malloc(sizeof(struct hist_data) * n);
  if (hist == NULL)
      throw pfs::Exception( "Unable to allocate memory." );
    
  // Build histogram info
  if (verbose) fprintf(stderr, PROG_NAME ": Generating histogram.\n");
  for(int i = 0; i < n; i++)
    {
      hist[i].size = Y[i];
      hist[i].index = i;
    } 
  
  // Generate histogram
  qsort(hist, n, sizeof(struct hist_data), hist_data_comp_size);

  // Calculate cdf
  const float norm = 1.0f / (float) n;
  for (int i = 0; i < n; i++)
    hist[i].size = ((float) i) * norm;

  if (verbose) fprintf(stderr, PROG_NAME ": Applying histogram.\n");  

  // Recalculate in terms of indexes
  qsort(hist, n, sizeof(struct hist_data), hist_data_comp_index);

  //Remap gradient magnitudes
  for( int i = 0; i < n; i++ )
    Y[i] = hist[i].size;

  free(hist);
}


void histogramEqualizeRGB( const int cols, const int rows, float* const R, float* const G, float* const B, const float saturation, const bool verbose )
{
  const int n = cols*rows;

  // Allocate memory
  struct hist_data* const hist = (struct hist_data*) malloc(sizeof(struct hist_data) * n);
  if (hist == NULL)
      throw pfs::Exception( "Unable to allocate memory." );
    
  // Build histogram info
  if (verbose) fprintf(stderr, PROG_NAME ": Generating histogram.\n");
  for(int i = 0; i < n; i++)
    {
      float max = R[i];
      if (G[i]>max) max=G[i];
      if (B[i]>max) max=B[i];
      hist[i].size = max;
      hist[i].index = i;
    } 
  
  // Generate histogram
  qsort(hist, n, sizeof(struct hist_data), hist_data_comp_size);

  // Calculate cdf scaling factors
  const float norm = 1.0f / (float) n;
  for (int i = 0; i < n; i++)
    hist[i].size = ((float) i) * norm / hist[i].size;

  if (verbose) fprintf(stderr, PROG_NAME ": Applying histogram.\n");  

  // Recalculate in terms of indexes
  qsort(hist, n, sizeof(struct hist_data), hist_data_comp_index);

  //Remap gradient magnitudes
  for( int i = 0; i < n; i++ )
    {
      const float factor = powf(hist[i].size, saturation);
      R[i] *= factor;
      G[i] *= factor;
      B[i] *= factor;
    }
  free(hist);
}

int main( int argc, char* argv[] )
{
  try {
    mainBody( argc, argv );
  }
  catch( pfs::Exception ex ) {
    fprintf( stderr, PROG_NAME " error: %s\n", ex.getMessage() );
    return EXIT_FAILURE;
  }        
  catch( QuietException ) {
    return EXIT_FAILURE;
  }        
  return EXIT_SUCCESS;
}