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

  Program:   Insight Segmentation & Registration Toolkit
  Module:    $RCSfile: itkSTAPLEImageFilter.txx,v $
  Language:  C++
  Date:      $Date: 2004-03-02 01:19:34 $
  Version:   $Revision: 1.3 $

  Copyright (c) 2002 Insight 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 _itkSTAPLEImageFilter_txx
#define _itkSTAPLEImageFilter_txx
#include "itkSTAPLEImageFilter.h"

#include "itkImageRegionIterator.h"
#include "itkImageRegionConstIterator.h"

namespace itk
{

template <typename TInputImage, typename TOutputImage>
void
STAPLEImageFilter<TInputImage, TOutputImage>
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os,indent);
  os << indent << "m_MaximumIterations = " << m_MaximumIterations << std::endl;
  os << indent << "m_ForegroundValue = " << m_ForegroundValue << std::endl;
  os << indent << "m_ConfidenceWeight = " << m_ConfidenceWeight << std::endl;
  os << indent << "m_ElapsedIterations = " << m_ElapsedIterations << std::endl;
}

template< typename TInputImage, typename TOutputImage >
void
STAPLEImageFilter< TInputImage, TOutputImage >
::GenerateData()
{
  const double epsilon = 1.0e-10;
  
  typedef ImageRegionConstIterator< TInputImage > IteratorType;
  typedef ImageRegionIterator< TOutputImage> FuzzyIteratorType;
  
  const double min_rms_error = 1.0e-14; // 7 digits of precision
  
  unsigned int i, iter, number_of_input_files;
  
  // Allocate the output "fuzzy" image.
  this->GetOutput()->SetBufferedRegion(this->GetOutput()->GetRequestedRegion());
  this->GetOutput()->Allocate();
  typename TOutputImage::Pointer W = this->GetOutput();
  
  // Initialize the output to all 0's
  FuzzyIteratorType out = FuzzyIteratorType(W, W->GetRequestedRegion());
  for (out.GoToBegin(); !out.IsAtEnd(); ++out)
    {
    out.Set(0.0);
    }
  
  // Record the number of input files.
  number_of_input_files = this->GetNumberOfInputs();
  
  IteratorType *D_it = new IteratorType[number_of_input_files];
  
  double *p = new double[number_of_input_files];  // sensitivity
  double *q = new double[number_of_input_files];  // specificity
  
  double *last_q = new double[number_of_input_files];
  double *last_p = new double[number_of_input_files];
  
  for (i=0; i < number_of_input_files; ++i)
    {
    last_p[i] = -10.0;
    last_q[i] = -10.0;
    }
  
  // Come up with an initial Wi which is simply the average of
  // all the segmentations.
  IteratorType in;
  for (i = 0; i < number_of_input_files; ++i)
    {
    if ( this->GetInput(i)->GetRequestedRegion() != W->GetRequestedRegion() )
      {
      itkExceptionMacro(<<"One or more input images do not contain matching RequestedRegions");
      }
    
    in  = IteratorType( this->GetInput(i), W->GetRequestedRegion() );
    out = FuzzyIteratorType( W, W->GetRequestedRegion() );
    
    for (in.GoToBegin(), out.GoToBegin(); !in.IsAtEnd(); ++in, ++out)
      {
      if (in.Get() > m_ForegroundValue - epsilon && in.Get()
          < m_ForegroundValue + epsilon)
        {
        out.Set(out.Get() + 1.0);
        }
      } 
    }  // end for
  
  // Divide sum by num of files, calculate the estimate of g_t
  double N = 0.0;
  double g_t = 0.0;
  for (out.GoToBegin(); !out.IsAtEnd(); ++out)
    {
    out.Set( out.Get() / static_cast<double>(number_of_input_files) );
    g_t += out.Get();
    N = N +1.0;
    }
  g_t = (g_t / N) * m_ConfidenceWeight;
  
  double p_num, p_denom, q_num, q_denom;
  
  for (iter = 0; iter < m_MaximumIterations; ++iter)
    {
    // Now iterate on estimating specificity and sensitivity
    for (i=0; i < number_of_input_files; ++i)
      {
      in = IteratorType(this->GetInput(i), W->GetRequestedRegion());
      out = FuzzyIteratorType(W, W->GetRequestedRegion());
      
      p_num = p_denom = q_num = q_denom = 0.0;
      
      // Sensitivity and specificity of this user
      for (in.GoToBegin(), out.GoToBegin(); !in.IsAtEnd(); ++in, ++out)
        {
        if (in.Get() > m_ForegroundValue - epsilon
            && in.Get() < m_ForegroundValue + epsilon) // Dij == 1
          {
          p_num += out.Get(); // out.Get() := Wi
          }
        else //        if (in.Get() != m_ForegroundValue) // Dij == 0
          {
          q_num += (1.0 - out.Get()); // out.Get() := Wi
          }
        p_denom += out.Get();  
        q_denom += (1.0 - out.Get());  
        }
      
      p[i] = p_num / p_denom;
      q[i] = q_num / q_denom;        
      }
    
    // Now recreate W using the new p's and q's
    // Need an iterator on each D
    // const double g_t = 0.1;  // prior likelihood that a pixel is incl.in segmentation
    double alpha1, beta1;
    
    for (i = 0; i < number_of_input_files; ++i)
      {
      D_it[i] = IteratorType(this->GetInput(i), W->GetRequestedRegion());
      }
    
    out = FuzzyIteratorType(W, W->GetRequestedRegion());
    
    for (out.GoToBegin(); !out.IsAtEnd(); ++out)
      {
      alpha1 = beta1 = 1.0;
      for (i =0; i < number_of_input_files; ++i)
        {
        if (D_it[i].Get() > m_ForegroundValue - epsilon && D_it[i].Get() < m_ForegroundValue + epsilon) // Dij == 1
          {
          alpha1 = alpha1 * p[i];
          beta1  = beta1 * (1.0 - q[i]);
          }
        else //Dij == 0
          {
          alpha1 = alpha1 * (1.0 -p[i]);
          beta1  = beta1  * q[i];
          }
        ++D_it[i];
        }
      out.Set(g_t * alpha1 /
              ( g_t * alpha1  + (1.0 - g_t) * beta1) );
      }

    this->InvokeEvent( IterationEvent() );
    
    // Check for convergence
    bool flag = false;
    if (iter !=0 )  // not on the first iteration
      {
      flag = true;
      for (i=0; i < number_of_input_files; ++i)
        {
        if ( ( (p[i] - last_p[i]) * (p[i] - last_p[i]) ) > min_rms_error )
          {
          flag = false;
          break;
          }
        if ( ( (q[i] - last_q[i]) * (q[i] - last_q[i]) ) > min_rms_error )
          {
          flag = false;
          break;
          }
        }
      }
    for (i = 0; i < number_of_input_files; ++i)
      {
      last_p[i] = p[i];
      last_q[i] = q[i];
      }

    if( this->GetAbortGenerateData() )
      {
      this->ResetPipeline(); 
      flag = true;
      }

    if (flag == true)
      {
      break;
      }
    }
  
  // Copy p's, q's, etc. to member variables

  m_Sensitivity.clear();
  m_Specificity.clear();
  for (i = 0; i< number_of_input_files; i++)
    {
    m_Sensitivity.push_back( p[i] );
    m_Specificity.push_back( q[i] );
    }
  m_ElapsedIterations = iter;

  delete[] q;
  delete[] p;
  delete[] last_q;
  delete[] last_p;
  delete[] D_it;  
}

} // end namespace itk

#endif