#define _INCLUDE_XOPEN_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <time.h>
#include <math.h>
#include "matrix.h"
#include "stax.h"

#define SQR( a ) (( a ) * ( a ))
#define CUBE( a ) (( a ) * ( a ) * ( a ))
#define FOURTH( a ) (( a ) * ( a ) * ( a ) * ( a ))

#define Max( a, b ) (( a > b ) ? ( a ) : ( b ))

double *trueU, *trueV, *infU, *infV, *infGradX, *infGradY;
double **cornerU, **cornerV, **responZ, **gradX, **gradY;
double **errorU, **errorV;
int nPts, nIm;

void
parabolaFit( double *x, double *y, int n, double *parab )
{
    int i;
    static double **a, *b;
    static int first = 1;

    if( first == 1 ) {
      first = 0;
      a = mtxCreate( 3, 3 );
      b = vecCreate( 3 );
    }

    for( i = 0; i < 3; i++ )
      a[ i ][ 0 ] = a[ i ][ 1 ] = a[ i ][ 2 ] = b[ i ] = 0.0;

    for( i = 0; i < n; i++ ) {
      a[ 0 ][ 0 ] += FOURTH( x[ i ] );
      a[ 0 ][ 1 ] += CUBE( x[ i ] );
      a[ 0 ][ 2 ] += SQR( x[ i ] );
      a[ 1 ][ 2 ] += x[ i ];

      b[ 0 ] += y[ i ] * SQR( x[ i ] );
      b[ 1 ] += y[ i ] * x[ i ];
      b[ 2 ] += y[ i ];
    }

    a[ 1 ][ 0 ] = a[ 0 ][ 1 ];
    a[ 2 ][ 0 ] = a[ 1 ][ 1 ] = a[ 0 ][ 2 ];
    a[ 2 ][ 1 ] = a[ 1 ][ 2 ];
    a[ 2 ][ 2 ] = (double) n;

    mtxInvert( a, 3 );

    mtxMultColVec( parab, 3, 3, a, b );
    return;
}

void
parabolaTest( void )
{
    double x[3] = { 3.0, 5.0, 7.0 };
    double y[3] = { 19.0, 41.0, 71.0 };
    double parabola[3];

    parabolaFit( x, y, 3, parabola );
    /* should be 1 x^2 + 3 x + 1 */
    printf( "%g x^2 + %g x + %g\n",
	    parabola[ 0 ], parabola[ 1 ], parabola[ 2 ] );

    y[0] = 16.3, y[1] = 53.3, y[2] = 108.7;
    /* should be 2.3 x^2 + 0.1 x - 4.7 */
    parabolaFit( x, y, 3, parabola );
    printf( "%g x^2 + %g x + %g\n",
	    parabola[ 0 ], parabola[ 1 ], parabola[ 2 ] );
    return;
}

int
main( int argc, char *argv[] )
{
    int i, j, oPts, s, argStart = 1;
    char *ch, label[ 20 ];
    double dot, norm, normG, magn, atanGrad, u_avg, v_avg, pEval = 1.0;
    double *winSize, *hMin, *hMax, *hWid, *fitCoord, parabola[3];
    FILE *fp;

    /* optional arguments */
    for( i = 1; i < argc && argv[ i ][ 0 ] == '-'; i++ ) {
      switch( argv[ i ][ 1 ] ) {
      case 'p': /* set parabola evaluation parameter */
	if( sscanf( argv[ i ], "-p%lf", &pEval ) != 1 ) {
	  printf( "Error parsing -p argument\n" );
	  exit( -1 );
	}
	break;
      } /* end switch */
    } /* end for( i ) */
    argStart = i;

    /* read the analytically computed locations */
    fp = fopen( argv[ argStart ], "r" );
    fscanf( fp, "%d", &nPts );
    if( nPts > 0 ) {
      trueU = (double *) malloc( nPts * sizeof( double ));
      trueV = (double *) malloc( nPts * sizeof( double ));
      for( i = 0; i < nPts; i++ )
	fscanf( fp, "%lf %lf", &trueU[ i ], &trueV[ i ] );
    }
    fclose( fp );

    nIm = argc - 3 - ( argStart - 1 );

    stxInit( 2 * nIm + 1 ); /* 1 for extrap corner, one per im for err:grad */
    hMin = (double *) malloc(( nIm + 1 ) * sizeof( double ));
    hMax = (double *) malloc(( nIm + 1 ) * sizeof( double ));
    hWid = (double *) malloc(( nIm + 1 ) * sizeof( double ));
    for( j = 0; j < nIm + 1; j++ )
      hMin[ j ] = -1.0, hMax[ j ] = 1.0, hWid[ j ] = 0.05;
    stxHistInit( nIm + 1, hMin, hMax, hWid ); /* ditto on +1 */

    cornerU = (double **) malloc( nIm * sizeof( double * ));
    cornerV = (double **) malloc( nIm * sizeof( double * ));
    responZ = (double **) malloc( nIm * sizeof( double * ));
    gradX   = (double **) malloc( nIm * sizeof( double * ));
    gradY   = (double **) malloc( nIm * sizeof( double * ));
    if( nPts > 0 ) {
      errorU  = (double **) malloc( nIm * sizeof( double * ));
      errorV  = (double **) malloc( nIm * sizeof( double * ));
    }
    winSize = (double *) malloc( nIm * sizeof( double ));
    for( ch = argv[ argStart + 1 ], j = 0; j < nIm; j++ ) {
      /* parse the second argument for window sizes */
      sscanf( ch, "%d", &s );
      ch++;  ch++;  /* assume 0 <= s <= 9, one for comma */
      winSize[ j ] = (double) s;

      /* read the data and compute error vectors */
      fp = fopen( argv[ argStart + 2 + j ], "r" );
      fscanf( fp, "%d", &oPts );
      if( oPts != nPts && nPts > 0 ) {
	printf( "multiple input files with different number of points!\n" );
	exit( -1 );
      }
      cornerU[ j ] = (double *) malloc( oPts * sizeof( double ));
      cornerV[ j ] = (double *) malloc( oPts * sizeof( double ));
      responZ[ j ] = (double *) malloc( oPts * sizeof( double ));
      gradX[ j ]   = (double *) malloc( oPts * sizeof( double ));
      gradY[ j ]   = (double *) malloc( oPts * sizeof( double ));
      if( nPts > 0 ) {
	errorU[ j ]  = (double *) malloc( oPts * sizeof( double ));
	errorV[ j ]  = (double *) malloc( oPts * sizeof( double ));
      }
      for( i = 0; i < oPts; i++ ) {
	fscanf( fp, "%lf %lf %lf %lf %lf",
		&cornerU[ j ][ i ], &cornerV[ j ][ i ], &responZ[ j ][ i ],
		&gradX[ j ][ i ], &gradY[ j ][ i ] );

	normG = sqrt( SQR( gradX[ j ][ i ] ) + SQR( gradY[ j ][ i ] ));

	/* this assumes that points are corresponded correctly!!! */
	if( nPts > 0 ) {
	  errorU[ j ][ i ] = cornerU[ j ][ i ] - trueU[ i ];
	  errorV[ j ][ i ] = cornerV[ j ][ i ] - trueV[ i ];
	  /* distance */
	  norm = sqrt( SQR( errorU[ j ][ i ] ) + SQR( errorV[ j ][ i ] ));
	  stxUpdate( j, norm );
	  /* correlation of direction with gradient */
	  dot = errorU[ j ][ i ] * gradX[ j ][ i ] +
	    errorV[ j ][ i ] * gradY[ j ][ i ];
	  stxHistUpdate( j, dot / ( norm * normG ));

	  stxUpdate( nIm + 1 + j, norm / normG );
	}
      }
      fclose( fp );
    } /* end for( j ) */

    /*
    ** a final pass: compute corner for infinitely small window by
    ** extrapolating from other window sizes and then compute this error
    */
    fitCoord = (double *) malloc( nIm * sizeof( double ));
    infU = (double *) malloc( oPts * sizeof( double ));
    infV = (double *) malloc( oPts * sizeof( double ));
    infGradX = (double *) malloc( oPts * sizeof( double ));
    infGradY = (double *) malloc( oPts * sizeof( double ));
    if( nPts > 0 )
      printf( "infinitely small window extrapolated corners\n" );
    else
      printf( "%d\n", oPts );
    for( i = 0; i < oPts; i++ ) {
      /* first: a sanity checker */
      for( j = 0, u_avg = v_avg = 0.0; j < nIm; j++ ) {
	u_avg += cornerU[ j ][ i ];
	v_avg += cornerV[ j ][ i ];
      }
      u_avg *= 1.0 / (double) nIm;
      v_avg *= 1.0 / (double) nIm;
      for( j = 0; j < nIm; j++ ) {
	if( fabs( cornerU[ j ][ i ] - u_avg ) > 2.0 /* pixels */ )
	  printf( "check sanity on U coord %d, img %d [ a = %lf, u = %lf ]\n",
		  i, j, u_avg, cornerU[ j ][ i ] );
	if( fabs( cornerV[ j ][ i ] - v_avg ) > 2.0 /* pixels */ )
	  printf( "check sanity on V coord %d, img %d [ a = %lf, v = %lf ]\n",
		  i, j, v_avg, cornerV[ j ][ i ] );
      }

      /* compute parabolic fit for U coord */
      for( j = 0; j < nIm; j++ )
	fitCoord[ j ] = cornerU[ j ][ i ];
      parabolaFit( winSize, fitCoord, nIm, parabola );
      infU[ i ] = SQR( pEval ) * parabola[ 0 ] +
	pEval * parabola[ 1 ] + parabola[ 2 ];

      /* ... and for V coord */
      for( j = 0; j < nIm; j++ )
	fitCoord[ j ] = cornerV[ j ][ i ];
      parabolaFit( winSize, fitCoord, nIm, parabola );
      infV[ i ] = SQR( pEval ) * parabola[ 0 ] +
	pEval * parabola[ 1 ] + parabola[ 2 ];

      /* are those coordinates at all reasonable? */
      if( infU[ i ] < cornerU[ 0 ][ i ] - 10.0 ||
	  infU[ i ] > cornerU[ 0 ][ i ] + 10.0 ||
	  infV[ i ] < cornerV[ 0 ][ i ] - 10.0 ||
	  infV[ i ] > cornerV[ 0 ][ i ] + 10.0 )
	printf( "corner %d moved from image 0 by ( %.2lf,%.2lf )\n", i,
		fabs( infU[ i ] - cornerU[ 0 ][ i ] ),
		fabs( infV[ i ] - cornerV[ 0 ][ i ] ));

      /* ... and the strength and direction of the gradient */
      for( j = 0; j < nIm; j++ )
	fitCoord[ j ] = sqrt( SQR( gradX[ j ][ i ] ) + SQR( gradY[ j ][ i ] ));
      parabolaFit( winSize, fitCoord, nIm, parabola );
#if 0 /* scale gradient by parabola value */
      magn = Max( SQR( pEval ) * parabola[ 0 ] +
		  pEval * parabola[ 1 ] + parabola[ 2 ], 0.0 );
#else /* just the direction, strength set to 1.0 */
      magn = 1.0;
#endif
      for( j = 0; j < nIm; j++ )
	fitCoord[ j ] = atan2( gradY[ j ][ i ], gradX[ j ][ i ] );
      parabolaFit( winSize, fitCoord, nIm, parabola );
      atanGrad = SQR( pEval ) * parabola[ 0 ] +
	pEval * parabola[ 1 ] + parabola[ 2 ];

      /* set up extrapolated gradient */
      infGradX[ i ] = cos( atanGrad ), infGradY[ i ] = sin( atanGrad );
#if 0
      norm = magn / sqrt( SQR( infGradX[ i ] ) + SQR( infGradY[ i ] ));
      infGradX[ i ] *= norm, infGradY[ i ] *= norm;
#else
      infGradX[ i ] *= magn, infGradY[ i ] *= magn;
#endif

      printf( "%8.4lf %8.4lf  0.0 %8.4lf %8.4lf\n",
	      infU[ i ], infV[ i ], infGradX[ i ], infGradY[ i ] );
      /* compute error */
      if( nPts > 0 ) {
	normG = sqrt( SQR( infGradX[ i ] ) + SQR( infGradY[ i ] ));
	norm = sqrt( SQR( infU[ i ] - trueU[i] ) + SQR( infV[ i ] - trueV[i] ));
	stxUpdate( nIm, norm );
	dot = ( infU[ i ] - trueU[ i ] ) * infGradX[ i ] +
	  ( infV[ i ] - trueV[ i ] ) * infGradY[ i ];
	if( normG > 1e-04 )
	  stxHistUpdate( nIm, dot / ( norm * normG ));
      }
    }

    if( nPts > 0 ) {
      stxPrint( nIm, NULL, "extrapolated  " );
      for( j = 0; j < nIm; j++ ) {
	sprintf( label, "%2d x %2d window",
		 (int) winSize[ j ], (int) winSize[ j ] );
	stxPrint( j, NULL, label );
      }
      for( j = 0; j < nIm; j++ )
	stxPrint( nIm + 1 + j, NULL, "err:grad      " );

      for( j = 0; j < nIm; j++ ) {
	printf( "histogram for correlation of gradient to error for %d x %d\n",
		(int) winSize[ j ], (int) winSize[ j ] );
	stxHistPrint( j );
      }
      printf( "histogram for correlation of gradient to error for limit\n" );
      stxHistPrint( nIm );

      for( j = 0; j < nIm; j++ ) {
	sprintf( label, "error%02dx%02d.dat",
		 (int) winSize[ j ], (int) winSize[ j ] );
	fp = fopen( label, "w" );
	for( i = 0; i < nPts; i++ )
	  fprintf( fp, "%8.4lf %8.4lf\n", errorU[ j ][ i ], errorV[ j ][ i ] );
	fclose( fp );
      }

      sprintf( label, "error%02dx%02d.dat", 0, 0 );
      fp = fopen( label, "w" );
      for( i = 0; i < nPts; i++ )
	fprintf( fp, "%8.4lf %8.4lf\n",
		 infU[ i ] - trueU[ i ], infV[ i ] - trueV[ i ] );
      fclose( fp );
    }

    return 0;
}