/* pamblend.c - blend multiple exposures to remove/enhance moving elements
**
** Copyright (C) 2004 by Fraser McCrossan, code outline based on "pamcut.c",
** Copyright (C) 1989 by Jef Poskanzer.
**
** Permission to use, copy, modify, and distribute this software and its
** documentation for any purpose and without fee is hereby granted, provided
** that the above copyright notice appear in all copies and that both that
** copyright notice and this permission notice appear in supporting
** documentation.  This software is provided "as is" without express or
** implied warranty.
*/

#include <pam.h>
#include <math.h>
#include <shhopt.h>

#define DEFAULT_TOLERANCE 6

struct cmdlineInfo {
  /* All the information the user supplied in the command line,
     in a form easy for the program to use.
  */
  char **inputFilespecs;  /* Filespecs of input files */
  int inputFileCount;
  unsigned int closest; /* select the pixel closest to most other pixels */
  unsigned int furthest; /* as above, but furthest */
  unsigned int average; /* select the average of all pixels */
  unsigned int brightest; /* select the brightest pixel */
  unsigned int darkest; /* select the brightest pixel */
  int tolerance; /* how close a pixel must be to be considered
		    "close" */
  unsigned int verbose;
};

struct cmdlineInfo cmdline;
struct pam *pam;   /* Array of input PAM images */

static void
parseCommandLine(int argc, char ** const argv,
                 struct cmdlineInfo * const cmdlineP) {
/*----------------------------------------------------------------------------
   Note that the file spec array we return is stored in the storage that
   was passed to us as the argv array.
-----------------------------------------------------------------------------*/
    optEntry *option_def = malloc(100*sizeof(optEntry));
        /* Instructions to OptParseOptions3 on how to parse our options.
         */
    optStruct3 opt;

    unsigned int option_def_index;
    unsigned int alg_count;

    option_def_index = 0;   /* incremented by OPTENTRY */
    OPTENT3(0,   "closest", OPT_FLAG, &cmdlineP->closest, NULL, 0);
    OPTENT3(0,   "furthest", OPT_FLAG, &cmdlineP->furthest, NULL, 0);
    OPTENT3(0,   "average", OPT_FLAG, &cmdlineP->average, NULL, 0);
    OPTENT3(0,   "brightest", OPT_FLAG, &cmdlineP->brightest, NULL, 0);
    OPTENT3(0,   "darkest", OPT_FLAG, &cmdlineP->darkest, NULL, 0);
    OPTENT3(0,   "tolerance", OPT_INT, &cmdlineP->tolerance, NULL, 0);

    /* Set the defaults */
    cmdlineP->inputFilespecs = NULL;
    cmdlineP->inputFileCount = 0;
    cmdlineP->closest = FALSE;
    cmdlineP->furthest = FALSE;
    cmdlineP->average = FALSE;
    cmdlineP->brightest = FALSE;
    cmdlineP->darkest = FALSE;
    cmdlineP->tolerance = DEFAULT_TOLERANCE;

    opt.opt_table = option_def;
    opt.short_allowed = FALSE;  /* We have no short (old-fashioned)
				   options */
    opt.allowNegNum = FALSE;  /* We may not have parms that are
				 negative numbers */

    optParseOptions3(&argc, argv, opt, sizeof(opt), 0);
        /* Uses and sets argc, argv, and some of *cmdlineP and others. */

    cmdlineP->inputFileCount = argc - 1;
    if (cmdlineP->inputFileCount < 3)
        pm_error("Need at least 3 input images.");
    cmdlineP->inputFilespecs = argv+1;
    alg_count = 0;
    if (cmdlineP->closest)
      ++alg_count;
    if (cmdlineP->furthest)
      ++alg_count;
    if (cmdlineP->average)
      ++alg_count;
    if (cmdlineP->brightest)
      ++alg_count;
    if (cmdlineP->darkest)
      ++alg_count;
    if (alg_count > 1)
      pm_error("Supply only one mode option");
    if (alg_count < 1) {
      pm_message("Selecting default of -closest.");
      cmdlineP->closest = TRUE;
    }
}

/* calculate the distance between the colours in colourspace of two
   pixels */
int distance(tuple a, tuple b, struct pam *pamP)
{
  int i;
  int distancesqr = 0;
  int diff;
  
  for (i = 0; i < pamP->depth; ++i) {
    diff = a[i] - b[i]; /* no need to calculate absolute value, since
			   we square it down below anyway */
    distancesqr += diff * diff;
  }
  return (int)sqrt(((double)distancesqr)+0.5);
}

void copy_tuple(tuple a, tuple b, struct pam *pamP)
{
  int i;

  for (i = 0; i < pamP->depth; ++i) {
    b[i] = a[i];
  }
}

void zero_tuple(tuple a, struct pam *pamP)
{
  int i;

  for (i = 0; i < pamP->depth; ++i) {
    a[i] = 0;
  }
}

/* add tuple a to tuple b, storing result in b */
void add_tuples(tuple a, tuple b, struct pam *pamP)
{
  int i;

  for (i = 0; i < pamP->depth; ++i) {
    b[i] += a[i];
  }
}

/* divide all elements of a tuple a by a constant divisor, result
   stored in a */
void divide_tuple(tuple a, int divisor, struct pam *pamP)
{
  int i;

  for (i = 0; i < pamP->depth; ++i) {
    a[i] /= divisor;
  }
}

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

  struct pam outpam;  /* Output PAM image */
  tuple** inputRows;   /* array of rows from input image */
  tuple* outputRow;  /* Row of output image */
  tuple averageAll; /* average of tuples in all images */
  unsigned int *withinarray; /* array of counts of pixels within tolerance */
  double *distancearray; /* array of total distances to other pixels */
  double dist;
  int bestcount;
  double brightness, tmpbright, dummy1, dummy2;
  int target;
  int ii, jj; /* input image loop variables */
  int ot; /* output tuple loop variable */
  int row; /* current row */
  int col; /* current column */
  
  pnm_init( &argc, argv );
  
  parseCommandLine(argc, argv, &cmdline);
  
  /* allocate memory */
  /* PAM headers */
  pam = malloc(cmdline.inputFileCount * sizeof(*pam));
  if (pam == NULL)
    pm_error("Unable to allocate memory for header array.");
  /* rows */
  inputRows = malloc(cmdline.inputFileCount * sizeof(*inputRows));
  if (inputRows == NULL)
    pm_error("Unable to allocate memory for row array.");
  if (cmdline.closest || cmdline.furthest) { /* don't need this in average
					  mode */
    withinarray = malloc(cmdline.inputFileCount * sizeof(*withinarray));
    if (withinarray == NULL)
      pm_error("Unable to allocate memory for within array.");
    distancearray = malloc(cmdline.inputFileCount * sizeof(*distancearray));
    if (distancearray == NULL)
      pm_error("Unable to allocate memory for distance array.");
  }

  /* open all input images and check the headers */
  for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
    pam[ii].file = pm_openr(cmdline.inputFilespecs[ii]);
    pnm_readpaminit(pam[ii].file, &(pam[ii]), sizeof(pam[ii]));
  }

  /* check that the images are all the same size */
  for (ii = 1; ii < cmdline.inputFileCount; ++ii) {
    if (pam[ii].width != pam[0].width ||
	pam[ii].height != pam[0].height ||
	pam[ii].depth != pam[0].depth)
      pm_error("Unmatched size/depth: \"%s\" (%dx%dx%d) \"%s\" (%dx%dx%d).",
	       cmdline.inputFilespecs[ii], pam[ii].width, pam[ii].height,
	       pam[ii].depth,
	       cmdline.inputFilespecs[0], pam[0].width, pam[0].height,
	       pam[0].depth);
  }

  /* allocate a single row for each image */
  for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
    inputRows[ii] = pnm_allocpamrow(&pam[ii]);
  }

  /* allocate the averaging tuples */
  averageAll = pnm_allocpamtuple(&(pam[0]));

  outpam = pam[0];    /* Initial value - most fields will be same */
  outpam.file = stdout;

  pnm_writepaminit(&outpam);

  outputRow = pnm_allocpamrow(&outpam);

  /* loop through all rows; since all images are the same size, we can
     use any of the image headers that we like */
  for (row = 0; row < outpam.height; ++row) {
    /* read each input row */
    for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
      pnm_readpamrow(&(pam[ii]), inputRows[ii]);
    }
    for (col = 0; col < outpam.width; ++col) {
      /* find the output pixel that we want */
      if (cmdline.average) {
	/* go through each image, add the current pixel to the averaging tuple */
	zero_tuple(averageAll, &outpam);
	for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
	  add_tuples(inputRows[ii][col], averageAll, &(pam[ii]));
	}
	/* calculate the average */
	divide_tuple(averageAll, cmdline.inputFileCount, &outpam);
	for (ot = 0; ot < outpam.depth; ++ot) {
	  outputRow[col][ot] = averageAll[ot];
	}
      }
      else if (cmdline.brightest || cmdline.darkest) {
	if (outpam.depth == 1) {
	  brightness = (double)(inputRows[0][col][0]);
	}
	else if (outpam.depth == 3) {
	  pnm_YCbCrtuple(inputRows[0][col], &brightness, &dummy1, &dummy2);
	}
	else {
	  pm_error("depth of %d is unsupported for brightness modes",
		   outpam.depth);
	}
	target = 0;
	for (ii = 1; ii < cmdline.inputFileCount; ++ii) {
	  if (outpam.depth == 1) {
	    tmpbright = (double)(inputRows[ii][col][0]);
	  }
	  else {
	    pnm_YCbCrtuple(inputRows[ii][col], &tmpbright, &dummy1, &dummy2);
	  }
	  if ((cmdline.brightest && tmpbright > brightness) ||
	      (cmdline.darkest && tmpbright < brightness)) {
	    brightness = tmpbright;
	    target = ii;
	  }
	}
	for (ot = 0; ot < outpam.depth; ++ot) {
	  outputRow[col][ot] = inputRows[target][col][ot];
	}
      }
      else if (cmdline.closest || cmdline.furthest) {
	/* for each pixel, calculate how many other pixels are within
	   the "tolerance" value, and select the output pixel as the
	   one with the most, or least */
	/* clear the array of counts */
	for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
	  withinarray[ii] = 0;
	  distancearray[ii] = 0;
	}
	/* compute the counts */
	for (ii = 0; ii < cmdline.inputFileCount; ++ii) {
	  for (jj = ii+1; jj < cmdline.inputFileCount; ++jj) {
	    dist = distance(inputRows[ii][col], inputRows[jj][col], &outpam);
	    if ((unsigned int)dist <= cmdline.tolerance) {
	      ++withinarray[ii];
	      ++withinarray[jj];
	    }
	    distancearray[ii] += dist;
	    distancearray[jj] += dist;
	  }
	}
	/* select the one we want */
	target = 0;
	bestcount = withinarray[0];
	dist = distancearray[0]; /* now use this to hold the best
				    distance so far */
	if (cmdline.closest) {
	  for (ii = 1; ii < cmdline.inputFileCount; ++ii) {
	    if (withinarray[ii] > bestcount ||
		(withinarray[ii] == bestcount && distancearray[ii] < dist)) {
	      target = ii;
	      bestcount = withinarray[ii];
	      dist = distancearray[ii];
	    }
	  }
	}
	else {
	  for (ii = 1; ii < cmdline.inputFileCount; ++ii) {
	    if (withinarray[ii] < bestcount ||
		(withinarray[ii] == bestcount && distancearray[ii] > dist)) {
	      target = ii;
	      bestcount = withinarray[ii];
	      dist = distancearray[ii];
	    }
	  }
	}
	for (ot = 0; ot < outpam.depth; ++ot) {
	  outputRow[col][ot] = inputRows[target][col][ot];
	}
      }
      else {
	pm_error("Internal error, nothing to do!");
      }
    }
    pnm_writepamrow(&outpam, outputRow);
  }

  pnm_freepamrow(outputRow);
  pm_close(outpam.file);

  pnm_freepamtuple(averageAll);

  for (ii = 1; ii < cmdline.inputFileCount; ++ii) {
    pnm_freepamrow(inputRows[ii]);
    pm_close(pam[ii].file);
  }
  free(inputRows);
  free(pam);
    
  exit(0);
}