/*=========================================================================

Program:   Insight Segmentation & Registration Toolkit
Module:    $RCSfile: itkMaskedSpatialObjectToImageFilter.cxx,v $
Language:  C++
Date:      $Date: 2006/07/19 20:06:17 $
Version:   $Revision: 1.00 $

Copyright (c) Insight Software Consortium. All rights reserved.
See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

This software is distributed WITHOUT ANY WARRANTY; without even 
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef __itkMaskedSpatialObjectToImageFilter_cxx
#define __itkMaskedSpatialObjectToImageFilter_cxx

#include "itkMaskedSpatialObjectToImageFilter.h"
#include "itkSpatialObjectToImageFilter.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkNumericTraits.h"
#include "itkImageRegionConstIterator.h"
#include "itkProgressReporter.h"
#include "itkBinaryThresholdImageFilter.h"
#include "itkBinaryDilateImageFilter.h"
#include "itkBinaryBallStructuringElement.h"
#include "itkResampleImageFilter.h"
#include "itkCenteredAffineTransform.h"
#include "itkNearestNeighborInterpolateImageFunction.h"
#include "itkBinaryThresholdImageFilter.h"

namespace itk
{

/** Constructor */
template <class TInputSpatialObject, class TOutputImage>
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::MaskedSpatialObjectToImageFilter()
{
  this->SetNumberOfRequiredInputs(1);
  m_ChildrenDepth = 1;
  m_Size.Fill(0);
  
  for (unsigned int i = 0; i < OutputImageDimension; i++)
    {
    m_Spacing[i] = 1.0;
    m_Origin[i] = 0.0;
    }

  m_InsideValue = 0;
  m_OutsideValue = 0;
  m_UseObjectValue = false;

  m_UserSuppliedMask = false;
  m_MaskResampleFactor = 4.0;
  m_MaskDilationSize = 2;
}


/** Destructor */
template <class TInputSpatialObject, class TOutputImage>
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::~MaskedSpatialObjectToImageFilter()
{
}
  

/** Set the Input SpatialObject */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::SetInput(const InputSpatialObjectType *input)
{
  // Process object is not const-correct so the const_cast is required here
  this->ProcessObject::SetNthInput(0, 
                                   const_cast< InputSpatialObjectType * >( input ) );
}


/** Set the Input Mask image */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::SetMask(const MaskImageType *mask)
{
  if (mask)
  {
    m_UserSuppliedMask = true;

    // Const-correct the mask image
    this->ProcessObject::SetNthInput(1, 
                                     const_cast< MaskImageType * >( mask ) );
  }
}


/** Get the input Spatial Object */
template <class TInputSpatialObject, class TOutputImage>
const typename MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>::InputSpatialObjectType *
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::GetInput(void) 
{
  if (this->GetNumberOfInputs() < 1)
    {
    return 0;
    }
  
  return static_cast<const TInputSpatialObject * >
    (this->ProcessObject::GetInput(0) );
}


/** Get the input Mask image */
template <class TInputSpatialObject, class TOutputImage>
const typename MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>::MaskImageType *
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::GetMask(void)
{
  if (this->GetNumberOfInputs() < 2)
    {
    return 0;
    }
  
  return static_cast<const MaskImageType * >
    (this->ProcessObject::GetInput(1) );
}


/** GenerateData */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::GenerateData(void)
{
    this->SetupOutputImage();
    
    if (!m_UserSuppliedMask)
        //We must generate the mask ourselves
        this->GenerateMask();
    
    this->GenerateDataUsingMask();
}


/** SetupOutputImage */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::SetupOutputImage(void)
{
  // Get the input, and output pointers
  const InputSpatialObjectType * InputObject = this->GetInput();
  OutputImagePointer OutputImage = this->GetOutput();

  // Generate the image
  SizeType size;

  InputObject->ComputeBoundingBox();
  for(unsigned int i=0;i<ObjectDimension;i++)
    {
    size[i] = (long unsigned int)(InputObject->GetBoundingBox()->GetMaximum()[i]
                                  - InputObject->GetBoundingBox()->GetMinimum()[i]);
    }
  
  typename OutputImageType::IndexType index;
  index.Fill(0);
  typename OutputImageType::RegionType region;
  
  // If the size of the output has been explicitly specified, the filter
  // will set the output size to the explicit size, otherwise the size from the spatial
  // object's bounding box will be used as default.

  bool specified = false;
  for (unsigned int i = 0; i < OutputImageDimension; i++)
    {
    if (m_Size[i] != 0)
      {
      specified = true;
      break;
      }
    }

  if (specified)
    {
    region.SetSize( m_Size );
    }
  else
    {
    region.SetSize( size );
    }
  region.SetIndex( index );

  OutputImage->SetLargestPossibleRegion( region);     // 
  OutputImage->SetBufferedRegion( region );           // set the region 
  OutputImage->SetRequestedRegion( region );          //                                                                       
  // If the spacing has been explicitly specified, the filter
  // will set the output spacing to that explicit spacing, otherwise the spacing from
  // the spatial object is used as default.
  
  specified = false;
  for (unsigned int i = 0; i < OutputImageDimension; i++)
    {
    if (m_Spacing[i] != 0)
      {
      specified = true;
      break;
      }
    }
  
  if (specified)
    {
    OutputImage->SetSpacing(this->m_Spacing);         // set spacing
    }
  else
    {
    OutputImage->SetSpacing(InputObject->GetIndexToObjectTransform()->GetScaleComponent());   // set spacing
    }
  OutputImage->SetOrigin(m_Origin);                 //   and origin
  OutputImage->Allocate();   // allocate the image
}

/** Generate Mask */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::GenerateMask(void)
{
    // NOTE: The mask is store in Inputs(1)
    MaskImageType::Pointer maskImage = NULL;

    //Convert SpatialObject to low-resolution mask
    typedef itk::SpatialObjectToImageFilter<TInputSpatialObject, MaskImageType>
       SpatialObjectToImageFilterType;
    SpatialObjectToImageFilterType::Pointer filterConvertMask = 
        SpatialObjectToImageFilterType::New();
    filterConvertMask->SetInput(this->GetInput());
    SpatialObjectToImageFilterType::SizeType maskSize;          //Set Size
    SpatialObjectToImageFilterType::SpacingType maskSpacing;    //Set Spacing
    for (unsigned int dim=0; dim < OutputImageDimension; dim++)
    {
        maskSize[dim] = m_Size[dim] / m_MaskResampleFactor;
        maskSpacing[dim] = m_Spacing[dim] * ((double)m_Size[dim]/(double)maskSize[dim]);
    }
	filterConvertMask->SetOrigin(m_Origin);
    filterConvertMask->SetSize(maskSize);
    filterConvertMask->SetSpacing(maskSpacing);
    filterConvertMask->SetInsideValue(itk::NumericTraits<MaskPixelType>::max());
    filterConvertMask->SetOutsideValue(itk::NumericTraits<MaskPixelType>::Zero);
    filterConvertMask->SetUseObjectValue(false);
    filterConvertMask->Update();

	//Binary threshold images to ensure SpatialObjects which are greyscale
	//ie. GaussianSpatialObject or ImageSpatialObject.
	//See Insight Journal comment from 07-24-2006.
	typedef itk::BinaryThresholdImageFilter<MaskImageType, MaskImageType>
		BinaryThresholdImageFilterType;
	BinaryThresholdImageFilterType::Pointer filterThresholdMask =
		BinaryThresholdImageFilterType::New();
	filterThresholdMask->SetLowerThreshold(1);
	filterThresholdMask->SetInput(filterConvertMask->GetOutput());
	filterThresholdMask->Update();

    //Create structuring element
    typedef itk::BinaryBallStructuringElement<MaskPixelType, OutputImageDimension>
        BinaryBallElementType;
    BinaryBallElementType kernel;
    kernel.SetRadius(m_MaskDilationSize);
    kernel.CreateStructuringElement();

    //Binary dilate the resampled Mask image
    typedef itk::BinaryDilateImageFilter<MaskImageType, MaskImageType, BinaryBallElementType>
        BinaryDilateImageFilterType;
    BinaryDilateImageFilterType::Pointer filterDilateMask = BinaryDilateImageFilterType::New();
    filterDilateMask->SetInput(filterThresholdMask->GetOutput());
    filterDilateMask->SetKernel(kernel);
    filterDilateMask->SetDilateValue(itk::NumericTraits<MaskPixelType>::max());
    filterDilateMask->Update();

    //Resample mask to full size
    typedef itk::ResampleImageFilter<MaskImageType, MaskImageType, MaskCoordRep>
        ResampleFilterType;
    typedef itk::CenteredAffineTransform<MaskCoordRep, OutputImageDimension>
        TransformType;
    typedef itk::NearestNeighborInterpolateImageFunction<MaskImageType, MaskCoordRep>
        InterpolatorType;

    InterpolatorType::Pointer interpolator = InterpolatorType::New();
    TransformType::Pointer transform = TransformType::New();
    transform->SetIdentity();
    ResampleFilterType::Pointer filterMaskResample = ResampleFilterType::New();
    filterMaskResample->SetTransform(transform);
    filterMaskResample->SetInterpolator(interpolator);
    filterMaskResample->SetInput(filterDilateMask->GetOutput());
    double finalMaskOrigin[OutputImageDimension];       //Set Origin
    ResampleFilterType::SizeType finalMaskSize;         //Set Size
    ResampleFilterType::SpacingType finalMaskSpacing;   //Set Spacing
    for (unsigned int dim=0; dim < OutputImageDimension; dim++)
    {
       finalMaskOrigin[dim] = m_Origin[dim]; finalMaskOrigin[dim] = 0.0;
       finalMaskSize[dim] = m_Size[dim];
       finalMaskSpacing[dim] = m_Spacing[dim];
    }

    filterMaskResample->SetOutputOrigin(finalMaskOrigin);
    filterMaskResample->SetSize(finalMaskSize);
    filterMaskResample->SetOutputSpacing(finalMaskSpacing);
    filterMaskResample->Update();
    
    //Set mask image
    maskImage = filterMaskResample->GetOutput();
    this->SetMask(maskImage);
}


/** GenerateData (using Mask) */
template <class TInputSpatialObject, class TOutputImage>
void
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::GenerateDataUsingMask(void)
{
  itkDebugMacro(<< "MaskedSpatialObjectToImageFilter::Update() called");

  // Get the input, mask and output pointers
  const InputSpatialObjectType * InputObject = this->GetInput();
  OutputImagePointer OutputImage = this->GetOutput();
  const MaskImageType * MaskImage = this->GetMask();

  //Get output region
  typename OutputImageType::RegionType region =
      OutputImage->GetLargestPossibleRegion();

  //Setup Progress reporter
  ProgressReporter progress(this, 0, region.GetNumberOfPixels());

  typedef itk::ImageRegionIteratorWithIndex<OutputImageType> OutputIteratorType;
  typedef itk::ImageRegionConstIterator<MaskImageType> MaskIteratorType;

  OutputIteratorType itOutput(OutputImage, region);
  MaskIteratorType itMask(MaskImage, region);

  itk::Point<double,ObjectDimension> point;

  while(!itOutput.IsAtEnd())
    {
    //Check if the mask is "ON" for this point
    if (itMask.Get() > NumericTraits<MaskPixelType>::Zero)
    {
        // ValueAt requires the point to be in physical coordinate i.e
        for(unsigned int i=0;i<ObjectDimension;i++)
        {
            point[i]=(itOutput.GetIndex()[i]*m_Spacing[i])+m_Origin[i];
        }
        double val = 0.0;
        InputObject->ValueAt(point,val,99999);
        
        if( m_InsideValue != 0 || m_OutsideValue != 0 )
        {
            if (val)
            {
                if(m_UseObjectValue)
                {
                    itOutput.Set(static_cast<ValueType>(val));
                }
                else
                {
                    itOutput.Set(m_InsideValue);
                }
            }
            else
            {
                itOutput.Set(m_OutsideValue);
            }
        }
        else
        {
            itOutput.Set(static_cast<ValueType>(val));
        }
    }
    else
    {
        //The mask is "OFF" for this region - set as outside
        itOutput.Set(m_OutsideValue);        
    }
    ++itOutput;
    ++itMask;
    progress.CompletedPixel();
    }
  
  itkDebugMacro(<< "MaskedSpatialObjectToImageFilter::Update() finished");

} // end update function  


template<class TInputSpatialObject, class TOutputImage>
void 
MaskedSpatialObjectToImageFilter<TInputSpatialObject,TOutputImage>
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Size : " << m_Size << std::endl;
  os << indent << "Children depth : " << m_ChildrenDepth << std::endl;
  os << indent << "Inside Value : " << m_InsideValue << std::endl;
  os << indent << "Outside Value : " << m_OutsideValue << std::endl;
  if(m_UseObjectValue)
    {
    os << indent << "Using Object Value : ON" << std::endl;
    }
  else
    {
    os << indent << "Using Object Value : OFF" << std::endl;
    }
  if(m_UserSuppliedMask)
    {
    os << indent << "User Supplied Mask : YES" << std::endl;
    }
  else
    {
    os << indent << "User Supplied Mask : NO" << std::endl;
    }
}



} // end namespace itk

#endif