/* pamalign.c - align slightly-offset images of the same scene
**
** 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_MAXSHIFT 5
#define DEFAULT_PERCENTAGE 5
#define HUGE_DISTANCE 1e+38 /* no distance can be this large, so this
			       value used as initializer */
#define INTERPOLATION_START 0.5
#define INTERPOLATION_LIMIT 1e-3 /* smallest fraction of a pixel that
				    we search for */

struct cmdlineInfo {
  /* All the information the user supplied in the command line,
     in a form easy for the program to use.
  */
  char *firstInputFileSpec;
  char **otherInputFileSpecs;  /* Filespecs of input files */
  int otherImageCount;
  int maxshift; /* maximum we'll shift the image for a match in any
		   direction */
  unsigned int percentage; /* how much of the image to use for comparison */
  unsigned int verbose;
};

struct cmdlineInfo cmdline;
struct pam pam1, pam2; /* the two input images */
tuple** rows1;   /* array of rows from first input image */
tuple** rows2;   /* array of rows from second input image */
tuple* outputRow;  /* Row of output image */

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;

    option_def_index = 0;   /* incremented by OPTENTRY */
    OPTENT3(0,   "maxshift",    OPT_INT,    &cmdlineP->maxshift,  NULL,      0);
    OPTENT3(0,   "percentage",  OPT_INT,    &cmdlineP->percentage, NULL,      0);

    /* Set the defaults */
    cmdlineP->firstInputFileSpec = NULL;
    cmdlineP->otherInputFileSpecs = NULL;
    cmdlineP->maxshift = DEFAULT_MAXSHIFT;
    cmdlineP->percentage = DEFAULT_PERCENTAGE;

    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. */
    if (argc-1 < 2)
        pm_error("Need at least 2 input images.");
    cmdlineP->firstInputFileSpec = argv[1];
    cmdlineP->otherInputFileSpecs = argv+2;
    cmdlineP->otherImageCount = argc - 2;
}

/* calculate the distance between the colours in colourspace of two
   pixels */
double 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 anyway */
    distancesqr += diff * diff;
  }
  return sqrt(distancesqr);
}

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

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

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

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

/* 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 */
  int otherImage; /* loop variable to iterate through all other
		     images */
  int width, height; /* width and height of the test strip on the
			first image */
  int left_limit, right_limit; /* since they are used a lot,
				  precalculate these */
  int top_limit, bottom_limit; /* ditto */
  int pixels_required;
  int checkx, checky; /* match positions that we iterate through */
  int row, col; /* current row and column */
  int outrow, outcol; /* output row and column */
  double curdist; /* total distance for the current match position */
  double bestdist; /* best distance found */
  int bestx, besty; /* and the position found */
  
  pnm_init( &argc, argv );
  
  parseCommandLine(argc, argv, &cmdline);
  
  /* for now we'll be lazy and just read all of both images, both to
     get the data and to populate the pam* structures */
  pam1.file = pm_openr(cmdline.firstInputFileSpec);
  rows1 = pnm_readpam(pam1.file, &pam1, sizeof(pam1));
  pm_close(pam1.file);
  /* calculate and check the size of the match strip (which is
     overlayed on the second image to find the correct offset */
  width = pam1.width - (cmdline.maxshift * 2 + 1); /* max strip width
						      that allows full
						      movement left
						      and right */
  pixels_required = cmdline.percentage * pam1.width * pam1.height / 100;
  height = pixels_required / width;
  if (height > pam1.height)
    pm_error("%d%% of image with maxshift of %d would require a strip of height %d.\nInput images are only %d high.",
	     cmdline.percentage, cmdline.maxshift, height, pam1.height);
  if (height == 0)
    height = 1;
  left_limit = cmdline.maxshift;
  right_limit = pam1.width - left_limit - 1;
  top_limit = (pam1.height - height) / 2;
  bottom_limit = top_limit + height - 1; /* avoids any rounding errors */

  pm_message("Reference image \"%s\" (%dx%dx%d) match area (%dx%d)",
	     cmdline.firstInputFileSpec, pam1.width, pam1.height, pam1.depth,
	     width, height);

  for (otherImage = 0; otherImage < cmdline.otherImageCount; ++otherImage) {
    pam2.file = pm_openr(cmdline.otherInputFileSpecs[otherImage]);
    rows2 = pnm_readpam(pam2.file, &pam2, sizeof(pam2));
    pm_close(pam2.file);
    /* check that the images are the same size */
    if (pam1.width != pam2.width ||
	pam1.height != pam2.height ||
	pam1.depth != pam2.depth) {
      pm_message("Ignoring \"%s\" because of unmatched size/depth:\n\t\"%s\" (%dx%dx%d) \"%s\" (%dx%dx%d).",
		 cmdline.otherInputFileSpecs[otherImage],
		 cmdline.firstInputFileSpec,
		 pam1.width, pam1.height, pam1.depth,
		 cmdline.otherInputFileSpecs[otherImage],
		 pam2.width, pam2.height, pam2.depth);
    }
    else {
      /* iterate through all possible pixel offsets */
      bestdist = HUGE_DISTANCE;
      /* no need to initialize bestx and besty, since they will
	 automatically be overwritten by the first iteration */
      for (checky = -cmdline.maxshift; checky <= cmdline.maxshift; ++checky) {
	for (checkx = -cmdline.maxshift; checkx <= cmdline.maxshift; ++checkx) {
	  /* iterate through the strip of the image, adding up the total distance */
	  curdist = 0.0;
	  for (row = top_limit; row <= bottom_limit; ++row) {
	    for (col = left_limit; col <= right_limit; ++col) {
	      curdist += distance(rows1[row][col],
				  rows2[row+checky][col+checkx],
				  &pam1);
	    }
	  }
	  if (curdist < bestdist) {
	    bestdist = curdist;
	    bestx = checkx;
	    besty = checky;
	  }
	}
      }
      /* only update the file if the offset is non-zero */
      if (bestx == 0 && besty == 0) {
	pm_message("\"%s\" already aligned",
		   cmdline.otherInputFileSpecs[otherImage]);
      }
      else {
	/* output second image offset by best match */
	pm_message("Aligning \"%s\", offset (%d,%d), distance = %f",
		   cmdline.otherInputFileSpecs[otherImage],
		   -bestx, -besty, bestdist);
	outpam = pam1;    /* Initial value - most fields will be same */
	outpam.file = pm_openw(cmdline.otherInputFileSpecs[otherImage]);
	
	pnm_writepaminit(&outpam);
	
	outputRow = pnm_allocpamrow(&outpam);
	
	for (outrow = 0; outrow < pam2.height; ++outrow) {
	  row = outrow + besty; /* row we'll fetch the data from for this
				   output row */
	  for (outcol = 0; outcol < pam2.width; ++outcol) {
	    col = outcol + bestx;
	    if (row < 0 || row >= pam2.height ||
		col < 0 || col >= pam2.width) {
	      zero_tuple(outputRow[outcol], &outpam);
	    }
	    else {
	      copy_tuple(rows2[row][col], outputRow[outcol], &outpam);
	    }
	  }
	  pnm_writepamrow(&outpam, outputRow);
	}
	pnm_freepamrow(outputRow);
	pm_close(outpam.file);
      }
    }
    pnm_freepamarray(rows2, &pam2);
  }
  pnm_freepamarray(rows1, &pam1);

  exit(0);
}