/*********************************************************************************
*RunAssessment.cxx
*
*
*
*This program will run the 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:  a configuration file
*      
*author: Eric Billet
*
*creation date: 4/16/2008
*
*********************************************************************************/
#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();

}