/*********************************************************************************
RunAssessment2d.cxx


This program will run the local accuracy assessment on a set 
of two grayscale images. It will generate the estimated registration error image
based on local HD, local grayscale HD, or robust local grayscale HD metrics.

Input:  prefix, a configuration file
      
 Authors: Eric Billet, Andriy Fedorov 
 PI: Nikos Chrisochoides


Copyright (c) 2008 College of William and Mary
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

   * Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
   * Redistributions in binary form must reproduce the above
copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the
distribution.
   * Neither the name of the College of William and Mary nor the
names of its contributors may be used to endorse or promote products
derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
**********************************************************************************/
#include "itkImage.h"
#include "itkCastImageFilter.h"
#include "itkCurvatureAnisotropicDiffusionImageFilter.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkAdaptiveHistogramEqualizationImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"
#include "itkCannyEdgeDetectionImageFilter.h"
#include <iostream>
#include "itkPasteImageFilter.h"
#include "itkCastImageFilter.h"
#include "itkExtractImageFilter.h"
#include "itkMaskImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"
#include "itkLocalDistanceMapImageFilter.h"
#include "itkLocalDistanceMapSmoothingImageFilter.h"
#include "itkLocalGrayscaleDistanceMapImageFilter.h"
#include "itkSignedMaurerDistanceMapImageFilter.h"
#include "errorStatistics.h"

//typedefs for simplifying declarations 
typedef float InputPixelType;
const unsigned int Dimension = 2;
typedef itk::Image<InputPixelType,Dimension>InputImageType;
typedef   float           FloatPixelType;
typedef itk::Image< FloatPixelType, Dimension >  FloatImageType;
typedef FloatImageType MaskImageType;


typedef float        MiddlePixelType;
typedef float        OutputPixelType;
typedef float                CannyPixelType;

typedef unsigned char UcharPixelType;
typedef itk::Image< UcharPixelType, Dimension > UcharImageType;

typedef itk::CurvatureAnisotropicDiffusionImageFilter< InputImageType, InputImageType >
    CurvatureFlowImageFilterType;
typedef itk::CastImageFilter< FloatImageType, UcharImageType >
    Float2UcharCasterType;
typedef itk::AdaptiveHistogramEqualizationImageFilter<InputImageType> HistoEqualizerType;

typedef itk::ImageFileWriter< UcharImageType >  CharWriterType;
typedef itk::ImageRegionConstIterator< UcharImageType > UcharConstIteratorType;
typedef itk::ImageRegionConstIterator< MaskImageType > MaskConstIteratorType;
typedef itk::ImageRegionIterator< UcharImageType > UcharIteratorType;
typedef itk::RescaleIntensityImageFilter<InputImageType, InputImageType > RescaleFilter;
typedef itk::RescaleIntensityImageFilter<UcharImageType, UcharImageType > UcharRescaleFilter;
typedef itk::MaskImageFilter<UcharImageType,MaskImageType,UcharImageType> MaskFilter;

typedef itk::LocalDistanceMapImageFilter<FloatImageType, FloatImageType, FloatImageType> LocalDistanceFilterType;
typedef itk::LocalDistanceMapSmoothingImageFilter<UcharImageType, UcharImageType, FloatImageType, FloatImageType> LocalDistanceMapSmoothingImageFilterType;
typedef itk::LocalGrayscaleDistanceMapImageFilter<UcharImageType, UcharImageType, FloatImageType> LocalGrayscaleDistanceFilterType;
typedef itk::SignedMaurerDistanceMapImageFilter<UcharImageType,FloatImageType> DistanceFilterType;


typedef itk::CastImageFilter<UcharImageType,FloatImageType>
    CastFilterType;
typedef itk::CastImageFilter<FloatImageType,UcharImageType>
    ReverseCastFilterType;
typedef itk::CannyEdgeDetectionImageFilter<FloatImageType,FloatImageType>
    CannyFilterType;

//GLOBAL SETUP INFORMATION
char image1[255], image2[255], mask1[255],mask2[255], output_name[255], edge1name[255],edge2name[255];
int metric_choice;
double percent, RG_percent;
unsigned contrast_enh;
unsigned numThreads=1;
unsigned min_elements, max_deform, tolerance, G_radius, RG_window;
char metricName[10];


//FUNCTION HEADERS

/**
*Function: ParseSetup(char *)
* This function will parse the input file for all setup information.
*/
unsigned ParseSetup(char* prefixDirName,char* setupName);

/**
*Function: CountEdges(image, mask)
*This function will count the edges that are in the mask that have been detected 
* and returns the percentage of edges
* percentage = edges/points_in_mask
*/
int CountEdges(UcharImageType::Pointer inputImage,MaskImageType::Pointer mask );

/**
*Function: DetectEdges(image, edge, var, uthresh, lthresh)
*This function will perform 2d edge detection on a 3d image. it extract 2d slices, 
*detects edges in the slice, then pastes all of the slices back together. 3D edge
*detection could be used instead but accuracy improves slightly when 2d edges were
* used.
*/
//void DetectEdges(UcharImageType::Pointer inputImage, UcharImageType::Pointer &edgeImage, float variance, float uthreshold, float lthreshold);
//3D edge detection
void DetectEdges(UcharImageType::Pointer inputImage, UcharImageType::Pointer &outputImage, float var, float uthresh, float lthresh);


/**
*Function: PreprocessImages()
*This function will preprocess the input image and mask.  The preprocess steps include:
*smoothing, contrast enhancement, and edge detection. Contrast enhancment can be skipped
*by manipulating the contrast_enh flag.   
*/
unsigned PreprocessImages(InputImageType::Pointer input, MaskImageType::Pointer mask, UcharImageType::Pointer &edge_image, double percent);



/****************************************
MAIN PROGRAM
****************************************/
int main( int argc, char ** argv )
{
// Verify the number of parameters in the command line
  if( argc < 3 )
    {
    std::cerr << "Usage: " << std::endl;
    std::cerr << argv[0] << "[prefix directory name][setupfile name]" << std::endl;
    return EXIT_FAILURE;
    }

  //get the input infor
  ParseSetup(argv[1], argv[2]);
 

  //read in input images
  typedef  itk::ImageFileReader< InputImageType > ReaderType;
  typedef  itk::ImageFileReader< MaskImageType > MaskReaderType;
  typedef  itk::ImageFileReader< UcharImageType > CharReaderType;
  ReaderType::Pointer reader = ReaderType::New();
  reader->SetFileName( image1 );
  ReaderType::Pointer reader2 = ReaderType::New();
  reader2->SetFileName( image2 );

  MaskReaderType::Pointer mask_reader=MaskReaderType::New();
  mask_reader->SetFileName(mask1);
  MaskReaderType::Pointer mask_reader2=MaskReaderType::New();
  mask_reader2->SetFileName(mask2);  //may be the same as first mask

  mask_reader->Update();
  mask_reader2->Update();

  //output edge images
  UcharImageType::Pointer edges1 = UcharImageType::New();
  UcharImageType::Pointer edges2 = UcharImageType::New();

  //call preprocessing on each set(image, mask)
  std::cout<<std::endl<<"Preprocessing 1st image:"<<std::endl;
  PreprocessImages(reader->GetOutput(),mask_reader->GetOutput(), edges1, percent);
  std::cout<<std::endl<<"Preprocessing 2nd image:"<<std::endl;
  PreprocessImages(reader2->GetOutput(),mask_reader2->GetOutput(), edges2, percent);


  //write out the edges 
  CharWriterType::Pointer writer = CharWriterType::New();
  writer->SetFileName( edge1name );	
  writer->SetInput(edges1 );
  writer->Update();
  CharWriterType::Pointer writer2 = CharWriterType::New();
  writer2->SetFileName( edge2name );	
  writer2->SetInput(edges2);
  writer2->Update();


  //Now compute the distance map for error estimation

  FloatImageType::Pointer LDImage = FloatImageType::New();
  std::cout<<std::endl<<"Finding ldmap...";
  
  if(!strcmp(&metricName[0],"HD")){
   
   	DistanceFilterType::Pointer dtFilter = DistanceFilterType::New();
          dtFilter->SetInput(edges1);
          dtFilter->SquaredDistanceOff();
	  dtFilter->SetNumberOfThreads(numThreads);
          dtFilter->Update();
          
	DistanceFilterType::Pointer dtFilter2 = DistanceFilterType::New();
          dtFilter2->SetInput(edges2);
          dtFilter2->SquaredDistanceOff();
	  dtFilter2->SetNumberOfThreads(numThreads);
          dtFilter2->Update();


       LocalDistanceFilterType::Pointer ld_filter = LocalDistanceFilterType::New();
	 ld_filter->SetInput1(dtFilter->GetOutput());
	 ld_filter->SetInput2(dtFilter2->GetOutput());
         ld_filter->Update();
	 ld_filter->SetNumberOfThreads(numThreads);
         LDImage = ld_filter->GetOutput();
  }
  else{

        LocalGrayscaleDistanceFilterType::Pointer gld_filter = LocalGrayscaleDistanceFilterType::New();
	 gld_filter->SetInput1(edges1);
	 gld_filter->SetInput2(edges2); 
	 gld_filter->SetMaxDef(max_deform);
	 gld_filter->SetTol(tolerance);
         gld_filter->SetRadius(G_radius);
   
	 gld_filter->SetNumberOfThreads(numThreads);
	 gld_filter->Update();
     
     FloatImageType::Pointer ld_map = gld_filter->GetOutput();
     LDImage = gld_filter->GetOutput();
   
    

    if(!strcmp(&metricName[0],"RGHD")){
      LocalDistanceMapSmoothingImageFilterType::Pointer rgld_filter = LocalDistanceMapSmoothingImageFilterType::New();
	 rgld_filter->SetInput1(edges1); 
	 rgld_filter->SetInput2(edges2); 
         rgld_filter->SetInput3(ld_map);
	 rgld_filter->SetRadius(RG_window);
	 rgld_filter->SetMinElements(min_elements);
         rgld_filter->SetPercent(RG_percent);
	 rgld_filter->SetNumberOfThreads(numThreads);
	 rgld_filter->Update();
         LDImage = rgld_filter->GetOutput();


    }

  }

  //write out the distance map
  typedef itk::ImageFileWriter< FloatImageType  >  WriterType;
  WriterType::Pointer ld_writer = WriterType::New();
  ld_writer->SetFileName(output_name);
  ld_writer->SetInput( LDImage );
  ld_writer->Update();

  //Compute Error Estimation Statistics
  //make a combined edge image

   UcharImageType::Pointer combinedEdges = UcharImageType::New();
   combinedEdges->SetRegions( edges1->GetLargestPossibleRegion() );
   combinedEdges->SetSpacing(edges1->GetSpacing());
   combinedEdges->Allocate();
   UcharConstIteratorType F_edgeIt(edges1, edges1->GetLargestPossibleRegion()); 
   UcharConstIteratorType G_edgeIt(edges2, edges2->GetLargestPossibleRegion()); 
   UcharIteratorType edgeIt(combinedEdges, combinedEdges->GetLargestPossibleRegion()); 
 
   for(F_edgeIt.GoToBegin(),G_edgeIt.GoToBegin(), edgeIt.GoToBegin(); !F_edgeIt.IsAtEnd(); ++F_edgeIt,++G_edgeIt, ++edgeIt){
		if(F_edgeIt.Get() || G_edgeIt.Get())
			edgeIt.Set(1);	      
 	        else
                       edgeIt.Set(0);
		
   }
      
      

        errorStatistics<FloatImageType, UcharImageType> errLD;
	errLD.SetImage1(LDImage);
	errLD.SetEdgeImage1(combinedEdges);
	errLD.SetImage1Region(LDImage->GetLargestPossibleRegion());
	errLD.SetEdgeImage1Region(combinedEdges->GetLargestPossibleRegion());
	errLD.SetTrim_Bottom(1);
	errLD.SetPercentage(.90);
	errLD.ComputeValues();

  std::cout<<std::endl<<"Estimated Error Statistics:"<<std::endl;
  std::cout<<"---------------------------"<<std::endl;
  std::cout<<"Max:"<<errLD.GetMax()<<std::endl<<"RMS:"<<errLD.GetRMS()<<std::endl<<"PMax(90%): "<<errLD.GetPmax()<<" "
           <<std::endl<<"Mean: "<<errLD.GetMean()<<std::endl<<"LTS(90%) "<<errLD.GetLTS()<<std::endl<<"LMS: "
           <<errLD.GetLMS();

  return EXIT_SUCCESS;

}//end main()


/*********************************************
Parse Setup file
**********************************************/
unsigned ParseSetup(char* prefixDirName, char* setupName){

  int prefixLen = strlen(prefixDirName);
  std::ifstream setupFile(setupName);
  char comment[255];
  setupFile.getline(&comment[0], 255);
 
  strcpy(&image1[0], prefixDirName);
  strcpy(&image2[0], prefixDirName);
  setupFile >> &image1[prefixLen] >> &image2[prefixLen];
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);
  
  strcpy(&mask1[0], prefixDirName);
  strcpy(&mask2[0], prefixDirName);
  setupFile >> &mask1[prefixLen] >> &mask2[prefixLen];
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >> percent;
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  strcpy(&edge1name[0], prefixDirName);
  strcpy(&edge2name[0], prefixDirName);
  setupFile>>&edge1name[prefixLen] >> &edge2name[prefixLen];
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  strcpy(&output_name[0], prefixDirName);
  setupFile >> &output_name[prefixLen];
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >>contrast_enh;
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >>numThreads;
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >> &metricName[0];
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);
  
 //if local HD chosen, no more info to collect
  if(!strcmp(&metricName[0],"HD")){
     return 0;
  }

  
  setupFile.getline(&comment[0], 255);

  setupFile >>G_radius ;
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >>tolerance; 
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);
  
  setupFile >>max_deform; 
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  if(!strcmp(&metricName[0],"GHD")){
     return 0;
  }
  
  setupFile.getline(&comment[0], 255);

  setupFile >>RG_window ;
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

  setupFile >>RG_percent; 
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);
  
  setupFile >>min_elements; 
  setupFile.getline(&comment[0], 255);
  setupFile.getline(&comment[0], 255);

}

/*********************
Begin Preprocessing
**********************/
unsigned PreprocessImages(InputImageType::Pointer input, MaskImageType::Pointer mask, UcharImageType::Pointer &edge_image, double percent){
  //smooth
  unsigned iterations = 15;
  float timestep = 0.0625;
  float conductance = 1.0;

  std::cout<<"First Smoothing...";
  CurvatureFlowImageFilterType::Pointer smoothing = 
                         CurvatureFlowImageFilterType::New();


  smoothing->SetInput( input );
  smoothing->SetNumberOfIterations( iterations );
  smoothing->SetTimeStep( timestep );
  smoothing->SetConductanceParameter(conductance);
  smoothing->Update();
  InputImageType::Pointer tempimage =smoothing->GetOutput();
  std::cout<<"done";

  if(contrast_enh){
    //contrast enhancement 
    HistoEqualizerType::Pointer histe = HistoEqualizerType::New();
	histe->SetInput(smoothing->GetOutput());
	histe->SetAlpha(.25);
	histe->SetBeta(.25);
	histe->Update();
        tempimage = histe->GetOutput();
  


    std::cout<<"Smoothing...";
    //smooth again
    CurvatureFlowImageFilterType::Pointer smoothing2 = 
                         CurvatureFlowImageFilterType::New();

    smoothing2->SetInput(tempimage);
    smoothing2->SetNumberOfIterations( iterations );
    smoothing2->SetTimeStep( timestep );
    smoothing2->SetConductanceParameter(conductance);
    smoothing2->Update();
    tempimage = smoothing2->GetOutput();
    std::cout<<"done";
  }
  //rescale image

  RescaleFilter::Pointer rescale = RescaleFilter::New();
	
  rescale->SetOutputMinimum(   0 );
  rescale->SetOutputMaximum( 255);
  rescale->SetInput( tempimage );
	

  Float2UcharCasterType::Pointer f2uc_caster = Float2UcharCasterType::New();
  f2uc_caster->SetInput(rescale->GetOutput());
  f2uc_caster->Update();
  
  std::cout<<std::endl<<"Detecting edges...";
  //detect edges

  UcharImageType::Pointer image =f2uc_caster->GetOutput();


double uthresh = 0.0;	
double ustep=4; //step to change threshold
double lowest=.25; //lowest possible step
int pot_ct; //potential edge count
short found = 0;
double mid;
do{
   do{
      uthresh += ustep;
     
      DetectEdges(image, edge_image, 1, uthresh, uthresh);//detect edges
  
      pot_ct = CountEdges(edge_image, mask);//count edges
     
	if(pot_ct==percent){
		found=1;
		mid=uthresh;
	}
      }while(pot_ct>percent&&!found);
      if(!found){
        DetectEdges(image,edge_image, 1, uthresh, lowest);//detect edges
        pot_ct = CountEdges(edge_image, mask);
 
      if(pot_ct<percent){
        uthresh-=ustep;
        ustep/=2;
      }
      if(pot_ct==percent){
	found=1;
	mid=lowest;
      }
   }
}while(pot_ct<percent&& ustep >lowest && !found);


/* At this point, we have found the upper threshold
* so now, we need to find the lower threshold.  We do
*this by performing a binary search on range [lowest..uthresh]
*/

double low = lowest;
double high = uthresh;


if(ustep > lowest){
while (low <= high&& !found) {
           mid = (low + high) / 2;
           //Detect edges at mid
           
           DetectEdges(image, edge_image, 1, uthresh, mid);//detect edges
           pot_ct = CountEdges(edge_image, mask);//count edges
 
           if (pot_ct < percent)
               high = mid - lowest;
           else if (pot_ct > percent)
               low = mid + lowest;
           else
               found = 1; // parameters found
}
}
//if found = 0, the parameters could not be found
if(!found){
 std::cerr<<"ERROR:  Edge parameters could not be found for that percentage";
  return 0;
}
  

	MaskFilter::Pointer masker = MaskFilter::New();
  	 masker->SetInput1(edge_image);
  	 masker->SetInput2(mask);
  	 masker->Update();
         edge_image = masker->GetOutput();

        UcharRescaleFilter::Pointer rescaler = UcharRescaleFilter::New();
	//The output of an edge filter is 0 or 1
  	rescaler->SetOutputMinimum(   0 );
  	rescaler->SetOutputMaximum( 1 );
	rescaler->SetInput( edge_image );
        rescaler->Update();
	edge_image = rescaler->GetOutput();
        
std::cout<<"done";
return 1; 
}
/*******************
End of the Preprocessing Function
******************/

/*********************************************
 CountEdges() finds the percentage of edges
**********************************************/
int CountEdges(UcharImageType::Pointer inputImage,MaskImageType::Pointer mask ){
    long edge_counter = 0;
    long counter=0;

   
    UcharConstIteratorType inputIt( inputImage, inputImage->GetLargestPossibleRegion() );
    MaskConstIteratorType maskIt( mask, mask->GetLargestPossibleRegion() );

   for ( inputIt.GoToBegin(), maskIt.GoToBegin(); !inputIt.IsAtEnd(); ++inputIt, ++maskIt)
    {
	if(maskIt.Value()){
          counter++;
          if( inputIt.Value()!=0 )
            edge_counter++;
        }
 
    }

    //round the result
   return int((double(edge_counter)/double(counter))*100.0 + .5);
}



void DetectEdges(UcharImageType::Pointer inputImage, UcharImageType::Pointer &outputImage, float var, float uthresh, float lthresh)
{

    CastFilterType::Pointer caster = CastFilterType::New();
    caster->SetInput(inputImage);
    caster->Update();

          CannyFilterType::Pointer edgeDetector2 = CannyFilterType::New();
          edgeDetector2->SetUpperThreshold(uthresh);
          edgeDetector2->SetLowerThreshold(lthresh);
          edgeDetector2->SetVariance(var);
          edgeDetector2->SetInput(caster->GetOutput());
          edgeDetector2->Update();
          
    ReverseCastFilterType::Pointer revCaster = ReverseCastFilterType::New();
    revCaster->SetInput(edgeDetector2->GetOutput());
    revCaster->Update();

    outputImage = revCaster->GetOutput();

}