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

  Program:   Insight Segmentation & Registration Toolkit
  Module:    $RCSfile: GridFeatureParser.cxx,v $
  Language:  C++
  Date:      $Date: 2003/05/23 21:21:32 $
  Version:   $Revision: 1.15 $

  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.


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

#include <iostream>
#include <memory>

#include <fstream>
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageRegionConstIterator.h"
#include "itkImageRegionIterator.h"
#include "vnl/vnl_matrix.h"

#include "itkKullbackLeiblerNMF.h"
#include "itkExceptionObject.h"


int main( int argc, char *argv[] )
{

 bool ok = true;
 if( argc < 6 )
    {
    std::cerr << "Missing Parameters " << std::endl;
	  std::cerr << "Usage: KLNMFApp mhdFile outFile numClasses xsize ysize  (hard wired for 2 dims )" << std::endl;
	  std::cerr << "Usage: KLNMFApp testimageKullbackLeiblerNMF.mhd outimage.mhd 4 20 40" << std::endl;

    std::cerr << "" << std::endl;
    return 1;
    }


  //Input/output file name; 
  std::string inFile = (std::string) argv[1];
  std::string outFile = (std::string) argv[2];
  unsigned int numClasses = atoi(argv[3]);

  unsigned int dims = 2;
  std::vector<unsigned int> dimsize;
  dimsize.resize(dims);
  for(unsigned int i=0; i<dims; i++)
    {
    dimsize[i] = atoi(argv[4+i]);
    }
  
  typedef itk::Image< float, 2> MatrixImageType;
  typedef itk::ImageFileReader< MatrixImageType > MatrixImageReaderType;
  typedef MatrixImageType::SizeType MatrixImageSizeType;  

    // Volume meta information for each resolution level
  MatrixImageReaderType::Pointer ereader = MatrixImageReaderType::New();
  ereader->SetFileName( inFile.c_str() );
  try
    {
    ereader->Update();
    }
  catch( itk::ExceptionObject & err ) 
    { 
    std::cout << err << std::endl; 
    ok = false;
    }
  catch( ... )
    {
    std::cout << "Err: caught unknown exception" << std::endl;
    ok = false;
    }

  //Get the input data matrix size.
  MatrixImageSizeType esize = ereader->GetOutput()->GetBufferedRegion().GetSize();
  unsigned int nrows = esize.GetElement(0);
  unsigned int ncols = esize.GetElement(1);   
 
  //Set up the image iterator
  typedef itk::ImageRegionConstIterator< MatrixImageType > InputImageConstIterator;
  InputImageConstIterator
    imgIt( ereader->GetOutput(), 
             ereader->GetOutput()->GetLargestPossibleRegion() );

  //The data in image container is stored in column major format
  //the data in the matrix is stored in row major form
  //Copy the image data to the input matrix in the correct format
  typedef vnl_matrix< double > InputMatrixType;
  InputMatrixType mat(nrows,ncols);
  imgIt.GoToBegin();

  for(unsigned int r=0; r<mat.rows(); r++)
    {
    for(unsigned int c=0; c<mat.columns(); c++)
      {
      mat.put(r,c, (double) imgIt.Get());
      ++imgIt;
      }
    }

  //Instantiate the calculator
  itk::KullbackLeiblerNMF< double > nmf;

  nmf.SetInput( mat );
  nmf.SetNumberOfClasses(numClasses);
  nmf.SetVerbose(true);
  nmf.Initialize();

  try
    {
    nmf.Compute();
    
    //Write the decomposed H matrix in an image
    typedef itk::Image< unsigned char, 2> OutImageType;
    typedef OutImageType::SizeType OutImageSizeType;
    typedef OutImageType::IndexType OutImageIndexType;
    typedef OutImageType::RegionType OutImageRegionType;
    OutImageSizeType size;
    size[0] = dimsize[0];
    size[1] = dimsize[1];
     
    OutImageIndexType index;
    index.Fill(0);

    OutImageRegionType region;
    region.SetSize(size);
    region.SetIndex(index);

    OutImageType::Pointer outImg = OutImageType::New();
    outImg->SetLargestPossibleRegion(region);
    outImg->SetBufferedRegion(region);
    outImg->Allocate();

    typedef itk::ImageRegionIterator< OutImageType > Iterator;
    Iterator it( outImg, outImg->GetBufferedRegion() );

    //check if the vector size matches with the user provided size parameter
    unsigned int numpix=1;
    for(unsigned int i=0; i< dimsize.size(); i++)
      numpix *= dimsize[i];

    vnl_vector<unsigned int> *vec = nmf.GetMLELabel(nmf.GetHMatrix(), 2);

    if(numpix != vec->size())
      std::cout << "The out image size does not match with nmf decomposed matrix." << std::cout;

    it.GoToBegin();
    vnl_vector<unsigned int>::iterator vecIt = vec->begin();
    while(!it.IsAtEnd())
      {
      it.Set(*vecIt);
      ++it;
      ++vecIt;
      }

    typedef itk::ImageFileWriter< OutImageType > WriterType;
    WriterType::Pointer writer = WriterType::New();

    writer->SetFileName( outFile.c_str() );  
    writer->SetInput(outImg);
    writer->Update();

    }
  catch (const itk::ExceptionObject &e)
    {
    std::cout << e.what() << std::endl;
    }
  catch (const std::exception &e)
    {
    std::cout << e.what() << std::endl;
    exit(-1);
    }
  catch (...)
    {
    std::cout << "Module failed for unknown reason. Investigation needed." << std::endl;
    return -1;
    }  
  

  return 0;
}