// $Id: Accuracy.cpp,v 1.2 2006/08/02 03:39:28 pkaz Exp $
//
// This program registers the specified data sets and computes the resulting
// Fiducial Registration Error (FRE) and the mean, standard deviation and
// maximum Target Registration Error (TRE), if any targets exist.
//
// The program operation depends on the number of input files specified:
//
//   - for 1 input file, the program performs a registration between all marker
//     positions in the input file and the corresponding marker positions
//     in the Phantom CNC data. Because all markers are used for registration,
//     there are no targets and therefore TRE is not computed. Instead,
//     the program estimates the Fiducial Localization Error (FLE) from the
//     computed FRE.
//
//   - for 2 input files, the program computes a registration between the
//     data in each input file, using the registration markers identified
//     on the Phantom (Phantom::IsRegPoint). The remaining markers are
//     considered to be targets and are used for computing the TRE.

#include <iostream>

// cisst includes
#include <cisstNumerical/nmrSVD.h>

#include "DataSet.h"
#include "Phantom.h"

using namespace std;


// AccuracyDataSet is derived from DataSet.  It adds the methods for
// computing the registration and related results (FRE, TRE, etc.).
class AccuracyDataSet: public DataSet {
protected:
    vctFrm3 transform;
    vctDouble3 *deltas;
    double *dists;
    bool register_all;
    double fre;
    double tre_mean, tre_stddev, tre_max;
public:
    AccuracyDataSet(int maxpts, bool reg_all);
    ~AccuracyDataSet();
    bool ComputeRegistration();
    double ComputeResiduals();
    double EstimateFLE();
    void PrintResults();
};

// AccuracyDataSet constructor: allocates arrays.
AccuracyDataSet::AccuracyDataSet(int maxpts, bool reg_all) : DataSet(maxpts),
                                 register_all(reg_all)
{
    deltas = new vctDouble3[maxpts];
    dists = new double[maxpts];
}

// AccuracyDataSet destructor: frees memory
AccuracyDataSet::~AccuracyDataSet()
{
    delete [] deltas;
    delete [] dists;
}

// ComputeRegistration:  computes registration using method proposed by
// Arun(1987) and modified by Umeyama(1991). Returns true if registration successful.
bool AccuracyDataSet::ComputeRegistration()
{
    if (register_all)
        cout << "Computing registration with all " << nrpoints << " points." << endl;
    else
      cout << "Computing registration with " << Phantom::GetNumRegPts() <<
           " registration points." << endl;

    int i;
    int num = 0;
    // Compute averages
    vctDouble3 avg[2];
    avg[0].SetAll(0.0);
    avg[1].SetAll(0.0);
    for (i = 0; i < nrpoints; i++) {
        if (register_all || Phantom::IsRegpoint(nameIndex[i])) {
            num++;
            avg[0].Add(data[0][i]);
            avg[1].Add(data[1][i]);
        }
    }
    avg[0] /= num;
    avg[1] /= num;
    // Now set P[0], P[1] to data[0], data[1] minus mean (average) value.
    // Then, compute the sum of the outer products.
    vctDouble3 *P[2];
    P[0] = new vctDouble3[nrpoints];
    P[1] = new vctDouble3[nrpoints];
    vctDouble3x3 H, sumH;
    sumH.SetAll(0.0);
    for (i = 0; i < nrpoints; i++) {
        if (register_all || Phantom::IsRegpoint(nameIndex[i])) {
            P[0][i].DifferenceOf(data[0][i], avg[0]);
            P[1][i].DifferenceOf(data[1][i], avg[1]);
            H.OuterProductOf(P[0][i], P[1][i]);
            sumH.Add(H);
        }
    }
    // Now, compute SVD of sumH
    vctDouble3x3 U, Vt;
    vctDouble3 S;
    nmrSVD(sumH, U, S, Vt);
    // Compute X = V*U' = (U*V')'
    vctDouble3x3 X = (U*Vt).Transpose();
    double det = vctDeterminant<3>::Compute(X);
    // If determinant is not 1, apply correction from Umeyama(1991)
    if (fabs(det-1.0) > 1e-6) {
        vctDouble3x3 Fix(0.0);
        Fix.Diagonal() = vct3(1.0, 1.0, -1.0);
        X = (U*Fix*Vt).Transpose();
        det = vctDeterminant<3>::Compute(X);
        if (fabs(det-1.0) > 1e-6) {
            cout << "Registration failed!!" << endl;
            return false;
        }
    }
    vctMatRot3 R;
    R.Assign(X);
    transform = vctFrm3(R, avg[1]-R*avg[0]);
    return true;
}

// ComputeResiduals: uses registration result to compute FRE and TRE. 
// Stores distance errors (and delta-XYZ values) for each point for later output.
// Returns FRE.
double AccuracyDataSet::ComputeResiduals()
{
    int i;
    tre_mean = 0.0;
    tre_max = 0.0;
    double err2 = 0.0; // error squared
    double dist2;   // distance squared
    for (i = 0; i < nrpoints; i++) {
        deltas[i] = data[1][i] - transform*data[0][i];
        dist2 = deltas[i].NormSquare();
        dists[i] = sqrt(dist2);
        if (register_all || Phantom::IsRegpoint(nameIndex[i]))
            err2 += dist2;
        else {
            tre_mean += dists[i];
            if (dists[i] > tre_max)
                tre_max = dists[i];
        }
    }
    // Nreg = number of registration points
    // Ntarget = number of target points
    int Nreg = register_all?nrpoints:Phantom::GetNumRegPts();
    int Ntarget = nrpoints-Nreg;
    if (Ntarget > 0) {
        tre_mean /= Ntarget;                // Mean TRE
        tre_stddev = 0.0;
        for (i = 0; i < nrpoints; i++) {
          if (!Phantom::IsRegpoint(nameIndex[i]))
              tre_stddev += (dists[i]-tre_mean)*(dists[i]-tre_mean);
        }
        tre_stddev = sqrt(tre_stddev/(Ntarget-1)); // Standard deviation
    }
    fre = sqrt(err2/Nreg);
    return fre;
}

// EstimateFLE: estimates FLE from FRE using equation from Fitzpatrick(1998).
double AccuracyDataSet::EstimateFLE()
{
    double N = (double)(register_all?nrpoints:Phantom::GetNumRegPts());
    return sqrt((N/(N-2))*fre*fre);
}

// PrintResults: prints the results.
void AccuracyDataSet::PrintResults()
{
    printf("\nPoint  \t   Error \t    X   \t    Y   \t    Z\n");
    for (int i=0; i < nrpoints; i++) {
        printf("%4s   \t%8.3lf\t%8.3lf\t%8.3lf\t%8.3lf\n",
                Phantom::GetName(nameIndex[i]),
                dists[i], deltas[i].X(), deltas[i].Y(), deltas[i].Z());
    }
    if (!register_all)
        cout << endl << "TRE (mean, stddev, max) = " << tre_mean << ", "
             << tre_stddev << ", " << tre_max << endl;
}


int main(int argc, char* argv[])
{
    if (argc < 2) {
        cout << endl << "Syntax:  " << argv[0] << " input_file1 [input_file2]" << endl << endl;
        cout << "Performs registration between marker positions in specified input file(s)"
             << endl << "and/or gold standard (CNC positions)." << endl << endl;
        cout << argv[0] << " input_file1: Registers points in input file to points in CNC"
             << endl << "   coordinates using all points common to both data sets."
             << endl << "   Displays FRE and (estimated) FLE." << endl;
        cout << argv[0] << " input_file1 input_file2: Registers data in input_file1 to data in "
             << endl << "   input_file2 using the four registration points (1,3,10,11), which"
             << endl << "   must exist in both files. Displays FRE and TRE." << endl;
        return 0;
    }

    // Register all points if only one file specified
    bool register_all = (argc == 2);

    AccuracyDataSet DS(Phantom::GetNumPts(), register_all);
    DS.ReadFromFile(0, argv[1]);
    if (argc > 2)
        DS.ReadFromFile(1, argv[2]);
    else
        DS.ReadFromArray(1, Phantom::GetCNC());
    DS.FinalizeInput();

    if (DS.ComputeRegistration()) {
        double fre = DS.ComputeResiduals();
        cout << "FRE = " << fre << endl;
        if (register_all) {
            double fle = DS.EstimateFLE();
            cout << "FLE (est.) = " << fle << endl;
        }
        DS.PrintResults();
    }

    return 0;
}