/*****************************************************************************
**
** zerror3D.c = compute a 3D error metric for Zhang's calibration parameters
**
** This program reads in a set of points and calibration data for some number
** of camera image planes.  The points are assumed to be in correspondence
** according to the order in which they are read.  (Each file must have the
** same number of points.)  Each point, along with the calibration data for
** that camera, defines a ray in space.  Call the set of rays for a set of
** corresponding points a "bundle."
**
** If the calibration is perfect for each camera, the bundle will intersect
** in exactly one point.  Due to error, it will not.  In fact, some of the
** rays in a bundle may not intersect at all.  For each bundle, this program
** finds the nearest point of intersection for the bundle, then computes the
** average distance of the intersection point from each ray in the bundle.
** Over all bundles, this gives an error metric for the calibration.
**
*****************************************************************************/
#include <stdio.h>
#include <math.h>
#include "stax.h"
#include "optimize.h"
#include "matrix.h"

#define MAX_IMAGE_POINTS 200

int verbosity = 0;
typedef double Matrix[4][4];
typedef double Coord[4];

#define Normalize( v ) { \
  double norm; \
  norm = 1.0 / sqrt( SQR( v[0] ) + SQR( v[1] ) + SQR( v[2] )); \
  v[0] *= norm; \
  v[1] *= norm; \
  v[2] *= norm; \
  v[3]  = 0.0; \
}

#ifndef SQR
#define SQR( a ) (( a ) * ( a ))
#endif

/*
** These variables are global so that the parameter list to the minimizer
** doesn't have to be changed.
*/
int nIm;
Coord *l_org, *l_dir;

void
readCamera( FILE *fp, Matrix A, Matrix *Rt, double *k1, double *k2 )
{
    int i;

    /* the intrinsic matrix */
    fscanf( fp, "%lf %lf %lf %lf %lf",
	    &A[0][0], &A[0][1], &A[1][1], &A[0][2], &A[1][2] );
    A[1][0] = A[2][0] = A[2][1] = 0.0;
    A[3][0] = A[3][1] = A[3][2] = A[3][3] = A[2][3] = A[1][3] = A[0][3] = 0.0;
    A[2][2] = 1.0;

    /* the distortion parameters */
    fscanf( fp, "%lf %lf", k1, k2 );

    for( i = 0; i < nIm; i++ ) {
      /* rotation and translation - reading column matrix */
      fscanf( fp, "%lf %lf %lf   %lf %lf %lf   %lf %lf %lf   %lf %lf %lf",
              &Rt[i][0][0], &Rt[i][0][1], &Rt[i][0][2],
              &Rt[i][1][0], &Rt[i][1][1], &Rt[i][1][2],
              &Rt[i][2][0], &Rt[i][2][1], &Rt[i][2][2],
              &Rt[i][0][3], &Rt[i][1][3], &Rt[i][2][3] );
    }

    return;
}

void
readImageData( char *filename, Coord *image, int nPts )
{
    int i;
    FILE *fp;

    fp = fopen( filename, "r" );
    if( fp == NULL ) {
      printf( "unable to open file '%s'\n", filename );
      return;
    }

    for( i = 0; i < nPts; i++ ) {
      fscanf( fp, "%lf %lf", &image[i][0], &image[i][1] );
      image[i][2] = 1.0;
    }

    return;
}

#define Swap( a, b, t ) { t = a; a = b; b = t; }
void
sortByDistance( Coord *d, int n, double u, double v )
{
    int i, pass, changed;
    double *dist, tmp;

    /* compute distances */
    dist = (double *) malloc( n * sizeof( double ));
    for( i = 0; i < n; i++ )
      dist[i] = SQR( d[i][0] - u ) + SQR( d[i][1] - v );

    if( verbosity > 1 ) {
      printf( "PtNo    %22s     distance\n", "image point" );
      for( i = 0; i < n; i++ )
        printf( "[%2d]: < %10lf, %10lf >  %10lf\n",
		i, d[i][0], d[i][1], dist[i] );
    }

    /* bubble sort for easier code, since efficiency is almost irrelevant */
    pass = 1;
    do {
      changed = 0;
      for( i = 0; i < n - pass; i++ )
        if( dist[ i ] > dist[ i + 1 ] ) {
          Swap( dist[ i ], dist[ i + 1 ], tmp );
          Swap( d[ i ][ 0 ], d[ i + 1 ][ 0 ], tmp );
          Swap( d[ i ][ 1 ], d[ i + 1 ][ 1 ], tmp );
          changed = 1;
        }
      pass++;
    } while( changed == 1 );

    if( verbosity > 1 ) {
      printf( "PtNo    %22s     distance\n", "image point" );
      for( i = 0; i < n; i++ )
        printf( "[%2d]: < %10lf, %10lf >   %10lf\n",
		i, d[i][0], d[i][1], dist[i] );
    }

    return;
}

void
DistortPixel( double *u_, double *v_, double u, double v,
              double k1, double k2, Matrix A )
{
    double u0 = A[0][2], v0 = A[1][2];
    double alpha = A[0][0], beta = A[1][1], gamma = A[0][1];
    double x, y, r2;

    y = ( v - v0 ) / beta;
    x = ( u - u0 - gamma * y ) / alpha;
    r2 = SQR( x ) + SQR( y );

    *u_ = u + ( u - u0 ) * ( k1 * r2 + k2 * SQR( r2 ));
    *v_ = v + ( v - v0 ) * ( k1 * r2 + k2 * SQR( r2 ));

    return;
}

Matrix distort_A;
double distort_k1, distort_k2;
double distort_ud, distort_vd;
double
UndistortionError( double *crd )
{
    double test_u, test_v;

    DistortPixel( &test_u, &test_v, crd[0], crd[1],
		  distort_k1, distort_k2, distort_A );
    return sqrt( SQR( test_u - distort_ud ) + SQR( test_v - distort_vd ));
}

void
UndistortPixelCoords( Coord *un, Coord *d, int n,
                      Matrix *A, double k1, double k2 )
{
    int i, moved;
    double est[2];
    double min[2], max[2];

#if 0
    /* sort distorted and model coords by increasing distance from center */
    sortByDistance( d, n, (*A)[0][2], (*A)[1][2] );
#endif

    /*
    ** Set a few terms constant for each point.
    ** These should be parameters, but optimizer isn't set for fixed params.
    */
    memcpy( distort_A, *A, sizeof( Matrix ));
    distort_k1 = k1, distort_k2 = k2;

    /* now for each point, undistort it via a search for the right coords */
    for( i = 0; i < n; i++ ) {
      /* set coordinates to match in the optimizer's error function (above) */
      distort_ud = d[i][0];
      distort_vd = d[i][1];

      /* initial estimate: no distortion */
      est[0] = d[i][0];
      est[1] = d[i][1];

      /* bounds for the optimizer */
      min[0] = est[0] - 30.0, max[0] = est[0] + 30.0;
      min[1] = est[1] - 30.0, max[1] = est[1] + 30.0;

      optOptimize( est, 2, min, max, 0, UndistortionError );

      un[i][0] = est[0];
      un[i][1] = est[1];
    }
    return;
}

/*
** The distance between a point P and a line ( org + t * dir ) in 3D is
** computed in two steps.  First, find the closest point Q on the line to the
** given point.  Second, compute the distance between P and Q.  Finding Q
** involves setting up the general distance function between P and any point
** on the line as a function of t, the parameter of the line.  This function
** is quadratic with a positive leading coefficient.  Thus it has a minimum,
** which can be easily computed from the derivative.
**
** The general distance formula for l = ( a_0, b_0, c_0 ) + t * ( a, b, c )
** and P = ( x, y, z ) is
**
** d( t ) = ( a_0 + a * t - x )^2 +
**          ( b_0 + b * t - y )^2 +
**          ( c_0 + c * t - z )^2
**
**        = ( a^2 + b^2 + c^2 ) t^2 +
**          2 ( a ( a_0 - x ) + b ( b_0 - y ) + c ( c_0 - z ) ) t +
**          ( a_0 - x )^2 + ( b_0 - y )^2 + ( c_0 - z )^2
**
** which I rewrite as d( t ) = D t^2 + E t + F.  This is minimized by
** t = -E / ( 2 D ).  It is minimized because D - which determines the sign
** of the second derivative - is greater than or equal to zero.  This fact
** can in turn be seen from the long, multi-line expression for d( t ) above.
** This value of t enables us to find Q using the parametric equations for
** the line.
*/
double
pointToLineDistance( double *P, Coord org, Coord dir )
{
    double D, E, F, t;
    Coord Q;

    /* find the coefficients in the general distance formula */
    D = SQR( dir[0] ) + SQR( dir[1] ) + SQR( dir[2] );
    E = 2.0 * ( dir[0] * ( org[0] - P[0] ) +
                dir[1] * ( org[1] - P[1] ) +
                dir[2] * ( org[2] - P[2] ));
    F = SQR( org[0] - P[0] ) + SQR( org[1] - P[1] ) + SQR( org[2] - P[2] );

    /* find the parameter value for the closest point */
    t = -E / ( 2.0 * D );

    /* find the closest point */
    Q[0] = org[0] + t * dir[0];
    Q[1] = org[1] + t * dir[1];
    Q[2] = org[2] + t * dir[2];

    /* return the distance */
    return sqrt( SQR( P[0] - Q[0] ) + SQR( P[1] - Q[1] ) + SQR( P[2] - Q[2] ));
}

/* this is the error function for the minimizer */
double
bundleError( double *P )
{
    int i;
    double err;

    for( i = 0, err = 0.0; i < nIm; i++ )
      err += pointToLineDistance( P, l_org[i], l_dir[i] );

    return err / (double) nIm;
}

/*
** Return in Pt the nearest point of intersection for the rays defined by
** [ org, dir ]
*/
void
findNearestIntersection( Coord Pt, Coord *org, Coord *dir, int nRays )
{
    double est[3];
    double min[3], max[3];

    static double **A;
    static int first = 1;
    double b[3] = { 0.0, 0.0, 0.0 };
    int i;

    if( first ) {
      A = mtxCreate( 3, 3 );
      first = 0;
    }

    A[0][0] = A[0][1] = A[0][2] = 0.0;
    A[1][0] = A[1][1] = A[1][2] = 0.0;
    A[2][0] = A[2][1] = A[2][2] = 0.0;

    for( i = 0; i < nRays; i++ ) {
      if( verbosity > 1 ) {
	printf( "org[ %d ] = { %10.6lf, %10.6lf, %10.6lf },\n", i,
		org[ i ][ 0 ], org[ i ][ 1 ], org[ i ][ 2 ] );
	printf( "dir[ %d ] = { %10.6lf, %10.6lf, %10.6lf },\n", i,
		dir[ i ][ 0 ], dir[ i ][ 1 ], dir[ i ][ 2 ] );
      }

      A[0][0] += 1.0 - SQR( dir[ i ][ 0 ] );
      A[0][1] -= dir[ i ][ 0 ] * dir[ i ][ 1 ];
      A[0][2] -= dir[ i ][ 0 ] * dir[ i ][ 2 ];
      A[1][1] += 1.0 - SQR( dir[ i ][ 1 ] );
      A[1][2] -= dir[ i ][ 1 ] * dir[ i ][ 2 ];
      A[2][2] += 1.0 - SQR( dir[ i ][ 2 ] );

      b[0] += ( 1.0 - SQR( dir[ i ][ 0 ] )) * org[ i ][ 0 ]
            - dir[ i ][ 0 ] * dir[ i ][ 1 ] * org[ i ][ 1 ]
            - dir[ i ][ 0 ] * dir[ i ][ 2 ] * org[ i ][ 2 ];
      b[1] += ( 1.0 - SQR( dir[ i ][ 1 ] )) * org[ i ][ 1 ]
            - dir[ i ][ 0 ] * dir[ i ][ 1 ] * org[ i ][ 0 ]
            - dir[ i ][ 1 ] * dir[ i ][ 2 ] * org[ i ][ 2 ];
      b[2] += ( 1.0 - SQR( dir[ i ][ 2 ] )) * org[ i ][ 2 ]
            - dir[ i ][ 0 ] * dir[ i ][ 2 ] * org[ i ][ 0 ]
            - dir[ i ][ 1 ] * dir[ i ][ 2 ] * org[ i ][ 1 ];
    }
    A[1][0] = A[0][1], A[2][0] = A[0][2], A[2][1] = A[1][2];

    mtxInvert( A, 3 ); /* now A^(-1) */

    mtxMultColVec( est, 3, 3, A, b ); /* x = A^(-1) b */

    min[0] = est[0] - 2.5, max[0] = est[0] + 2.5;
    min[1] = est[1] - 2.5, max[1] = est[1] + 2.5;
    min[2] = est[2] - 2.5, max[2] = est[2] + 2.5;

    /* invoke the optimizer */
    optOptimize( est, 3, min, max, 0, bundleError );

    Pt[0] = est[0];
    Pt[1] = est[1];
    Pt[2] = est[2];
    Pt[3] = 1.0;

    return;
}

void
multVecMatrix( Coord r, Matrix m, Coord v ) /* r = m * v */
{
    r[0] = m[0][0] * v[0] + m[0][1] * v[1] + m[0][2] * v[2] + m[0][3] * v[3];
    r[1] = m[1][0] * v[0] + m[1][1] * v[1] + m[1][2] * v[2] + m[1][3] * v[3];
    r[2] = m[2][0] * v[0] + m[2][1] * v[1] + m[2][2] * v[2] + m[2][3] * v[3];
    r[3] = m[3][0] * v[0] + m[3][1] * v[1] + m[3][2] * v[2] + m[3][3] * v[3];

    return;
}

void
invertMatrix( Matrix invertedMatrix, Matrix srcMatrix )
{
    int i, j;
    Matrix copy;  /* temp copy of dest in case src == dest */

    /* invert translation offsets */
    for( i = 0; i < 3; i++ ) {
        copy[ i ][ 3 ] = - srcMatrix[ 0 ][ i ] * srcMatrix[ 0 ][ 3 ]
                         - srcMatrix[ 1 ][ i ] * srcMatrix[ 1 ][ 3 ]
                         - srcMatrix[ 2 ][ i ] * srcMatrix[ 2 ][ 3 ];
    }

    /* transpose rotation part */
    for( i = 0; i < 3; i++ ) {
        for ( j = 0; j < 3; j++ ) {
            copy[ i ][ j ] = srcMatrix[ j ][ i ];
        }
    }

    for( i = 0; i < 3; i++ ) {
        for( j = 0; j < 3; j++ ) {
            invertedMatrix[ i ][ j ] = copy[ i ][ j ];
        }
        invertedMatrix[ i ][ 3 ] = copy[ i ][ 3 ];
    }
    invertedMatrix[ 3 ][ 0 ] = 0.0;
    invertedMatrix[ 3 ][ 1 ] = 0.0;
    invertedMatrix[ 3 ][ 2 ] = 0.0;
    invertedMatrix[ 3 ][ 3 ] = 1.0;

    return;
}

int
main( int argc, char *argv[] )
{
    Matrix     A;              /* constant for each camera location */
    double     k1, k2;         /* radial distortion parameters */
    Matrix    *Rt, *Rt_inv;    /* rotation and translation for each image */
    Coord     **image, **un;   /* distorted, undistorted points in each image */
    int        nPts;           /* number of points */
    /* variables used for quick temporary storage */
    int        i, j;
    Coord      Pt, tmpVec, camPt;
    double     m_x, m_y, m_z, m_u, m_v, m_u_dist, m_v_dist;
    double     err2d_at_Pt, err3d_at_Pt;
    FILE      *fp;

    /* defaults */
    nIm = argc - 2; /* one for program name, one for <camera-file> */
    nPts = 69; /* number in model used for DCS460 calibration for PhotoModeler */

    /* parse arguments */
    for( i = 1; i < argc - 3 && argv[i][0] == '-'; i++ ) {
      if( strcmp( argv[i], "-v" ) == 0 ) {
        verbosity = 1;
        nIm--;
      } else if( strncmp( argv[i], "-n", 2 ) == 0 ) {
        sscanf( argv[i], "-n%d", &nPts );
        nIm--;
      } else {
        fprintf( stderr, "Usage: %s [ -v ] [ -nPTS ] <camera-file>", argv[0] );
        fprintf( stderr, "  <image-file1> ...  <image-fileN>\n" );
        exit( -1 );
      }
    }

    /* Read in the correspondence data and undistort the points */
    image = (Coord **) malloc( nIm * sizeof( Coord * ));
    un = (Coord **) malloc( nIm * sizeof( Coord * ));
    Rt = (Matrix *) malloc( nIm * sizeof( Matrix ));
    Rt_inv = (Matrix *) malloc( nIm * sizeof( Matrix ));

    /* Read in the calibration data */
    fp = fopen( argv[ argc - nIm - 1 ], "r" );
    readCamera( fp, A, Rt, &k1, &k2 );
    fclose( fp );

    /* undistort all coordinates in all images */
    for( i = 0; i < nIm; i++ ) {
      invertMatrix( Rt_inv[i], Rt[i] );

      image[i] = (Coord *) malloc( nPts * sizeof( Coord ));
      un[i] = (Coord *) malloc( nPts * sizeof( Coord ));
      readImageData( argv[ argc - nIm + i ], image[i], nPts );
      UndistortPixelCoords( un[i], image[i], nPts, &A, k1, k2 );

      if( verbosity > 1 ) {
        for( j = 0; j < nPts; j++ ) {
          printf( "raw coordinates = < %lf, %lf >\n",
		  image[i][j][0], image[i][j][1] );
          printf( "undistorted     = < %lf, %lf >\n",
		  un[i][j][0], un[i][j][1] );
          DistortPixel( &m_u, &m_v, un[i][j][0], un[i][j][1], k1, k2, A );
          printf( "re-distorted    = < %lf, %lf >\n", m_u, m_v );
        }
      } /* end verbosity */
    } /* end for each image */

    /* Global variables to avoid fixed parameters to the optimizer */
    l_org = (Coord *) malloc( nIm * sizeof( Coord ));
    l_dir = (Coord *) malloc( nIm * sizeof( Coord ));

    /*
    ** For each point
    **   For each image
    **     Create the ray in world coordinates
    **   Find the nearest point of intersection of those rays
    **   Measure the total distance between that point and each ray
    */
    stxInit( 2 );
    for( j = 0; j < nPts; j++ ) {
      for( i = 0; i < nIm; i++ ) {
        /* by using inverse, we get ray (camera) origin in world space */
        l_org[i][0] = Rt_inv[i][0][3];
        l_org[i][1] = Rt_inv[i][1][3];
        l_org[i][2] = Rt_inv[i][2][3];

        /* direction in camera coords */
        l_dir[i][1] = ( un[i][j][1] - A[1][2] ) / A[1][1];
        l_dir[i][0] = ( un[i][j][0] - A[0][2] - A[0][1] * l_dir[i][1] ) /
                      A[0][0];
        l_dir[i][2] = 1.0; /* above two are normalized coords */
        Normalize( l_dir[i] );

        /* convert direction from camera space to world space */
        multVecMatrix( tmpVec, Rt_inv[i], l_dir[i] );
        memcpy( /* to */ l_dir[i], /* from */ tmpVec, sizeof( Coord ));
      } /* end for( i ) */

      /* compute the nearest point of intersection and the residual error */
      findNearestIntersection( Pt, l_org, l_dir, nIm );
      err3d_at_Pt = bundleError( Pt );
      if( verbosity ) {
        printf( "3D point %2d at %10.6lf %10.6lf %10.6lf\n",
                j, Pt[0], Pt[1], Pt[2] );
        printf( "3D error at point %2d = %lf\n", j, err3d_at_Pt );
      }
      stxUpdate( 0, err3d_at_Pt );

      /* reproject point into each of the images and measure 2D error */
      for( i = 0; i < nIm; i++ ) {
        multVecMatrix( camPt, Rt[i], Pt ); /* world-to-camera */

        /* embed A * ( [ R t ] * M ) inline using structure of A */
        m_x = A[0][0] * camPt[0] + A[0][1] * camPt[1] + A[0][2] * camPt[2];
        m_y = /* 0 * camPt[0] + */ A[1][1] * camPt[1] + A[1][2] * camPt[2];
        m_z = /* 0 * camPt[0] + 0 * camPt[1] +  1.0 *   */        camPt[2];
        /* compute undistorted pixel coords */
        m_u = m_x / m_z;
        m_v = m_y / m_z;
#ifdef MEASURE_IN_DISTORTED_SPACE
        /* distort them to match image[i][j] */
        DistortPixel( &m_u_dist, &m_v_dist, m_u, m_v, k1, k2, A );

        err2d_at_Pt = sqrt( SQR( image[i][j][0] - m_u_dist ) +
                            SQR( image[i][j][1] - m_v_dist ));
        if( verbosity ) {
          printf( "raw point %2d in image %d:         < %lf, %lf >\n",
                  j, i, image[i][j][0], image[i][j][1] );
          printf( "projected point %2d in image %d:   < %lf, %lf >\n",
                  j, i, m_u_dist, m_v_dist );
          printf( "2D error at point %2d, image %d   = %lf\n\n\n\n",
		  j, i, err2d_at_Pt );
        }
        stxUpdate( 1, err2d_at_Pt );
#else
        err2d_at_Pt = sqrt( SQR( un[i][j][0] - m_u ) +
			    SQR( un[i][j][1] - m_v ));
        if( verbosity ) {
	  printf( "raw point %2d in image %d:  < %lf, %lf >\n",
		  j, i, image[i][j][0], image[i][j][1] );
          printf( "undistorted point %2d in image %d:  < %lf, %lf >\n",
                  j, i, un[i][j][0], un[i][j][1] );
          printf( "projected point %2d in image %d:    < %lf, %lf >\n",
                  j, i, m_u, m_v );
          printf( "2D error at point %2d, image %d   = %lf\n\n\n\n",
		  j, i, err2d_at_Pt );
        }
        stxUpdate( 1, err2d_at_Pt );

#endif
      } /* end for( i ) */
    } /* end for( j ) */

    stxPrint( 0,   "", "3D error         " );
    stxPrint( 1, NULL, "2D error         " );
    return 0;
}