#include "UnsupervisedClustering.h"
#include "KMeansPlusPlus.h"
#include "EMGaussianMixtures.h"
#include "GaussianMixtureModel.h"
#include "SNAPImageData.h"
#include "ImageWrapper.h"
#include "ImageWrapperTraits.h"

#include <itkImageRandomConstIteratorWithIndex.h>
#include <algorithm>

UnsupervisedClustering::UnsupervisedClustering()
{
  m_ClusteringEM = NULL;
  m_NumberOfClusters = 3;
  m_NumberOfSamples = 0;
}

UnsupervisedClustering::~UnsupervisedClustering()
{
  if(m_ClusteringEM)
    {
    delete m_ClusteringEM;
    delete m_ClusteringInitializer;
    }

}


void UnsupervisedClustering::SetDataSource(SNAPImageData *imageData)
{
  if(m_DataSource != imageData)
    {
    m_DataSource = imageData;
    m_SamplesDirty = true;

    int nvox = m_DataSource->GetMain()->GetNumberOfVoxels();
    m_NumberOfSamples = (nvox > 10000) ? 10000 : nvox;
    }
}

void UnsupervisedClustering::SetMixtureModel(GaussianMixtureModel *model)
{
  m_ClusteringEM->SetGaussianMixtureModel(model);
  m_MixtureModel = m_ClusteringEM->GetGaussianMixtureModel();
}

void UnsupervisedClustering::SampleDataSource()
{
  assert(m_DataSource->IsSpeedLoaded());

  // Figure out the number of data components
  unsigned int nComp = 0;
  for(LayerIterator lit = m_DataSource->GetLayers(
        MAIN_ROLE | OVERLAY_ROLE);
      !lit.IsAtEnd(); ++lit)
    {
    nComp += lit.GetLayer()->GetNumberOfComponents();
    }

  // Size the data array
  int nvox = m_DataSource->GetMain()->GetNumberOfVoxels();
  int nsam = (m_NumberOfSamples == 0) ? nvox : m_NumberOfSamples;

  // Create data structure for the EM code
  m_Samples.set_size(nsam, nComp);

  // Create a random walk through the speed image, which should be initialized
  // at this point. We iterate over the speed image because we can easily access
  // its internal image (it's always a scalar image)
  typedef SpeedImageWrapper::ImageType SpeedImage;
  const SpeedImage *speed = m_DataSource->GetSpeed()->GetImage();
  typedef itk::ImageRandomConstIteratorWithIndex<SpeedImage> RandomIter;
  RandomIter itRand(speed, speed->GetBufferedRegion());

  // The number of samples is set to a larger number to account for NaNs
  itRand.SetNumberOfSamples(speed->GetBufferedRegion().GetNumberOfPixels());

  // Initialize the 'central' samples list
  m_CenterSamples.clear();
  m_CenterSamples.reserve(400);

  // Define the center region
  itk::ImageRegion<3> rcenter = speed->GetBufferedRegion();
  rcenter.ShrinkByRadius(to_itkSize(Vector3d(rcenter.GetSize()) * 0.2));

  // Do the random walk
  int pVoxel = 0;
  for (; !itRand.IsAtEnd() && pVoxel < nsam; ++itRand)
  {
    // Go to the next sample location
    itk::Index<3> idx = itRand.GetIndex();

    // TODO: this is a really slow way to collect samples!
    int iOffset = 0;
    bool is_finite = true;
    for (LayerIterator lit = m_DataSource->GetLayers(MAIN_ROLE | OVERLAY_ROLE); !lit.IsAtEnd(); ++lit)
    {
      ImageWrapperBase  *iw = lit.GetLayer();
      vnl_vector<double> svec(iw->GetNumberOfComponents());
      iw->SampleIntensityAtReferenceIndex(idx, iw->GetTimePointIndex(), false, svec);
      for (unsigned int k = 0; k < iw->GetNumberOfComponents(); k++)
      {
        if(!std::isfinite(svec[k]))
        {
          std::cout << "Skipping non-finite sample at row " << pVoxel << std::endl;
          is_finite = false;
        }
        m_Samples(pVoxel, iOffset++) = svec[k];
      }
    }

    // Did we get a successful sample?
    if(is_finite)
    {
      // Store as a 'central' sample if in the central 60% of the image
      if (m_CenterSamples.size() < 400 && rcenter.IsInside(idx))
        m_CenterSamples.push_back(pVoxel);

      // Go to the next row
      pVoxel++;
    }
  }

  // Trim the matrix if needed
  if(pVoxel < nsam)
    m_Samples = m_Samples.extract(pVoxel, nComp);

  m_NumberOfVoxels = pVoxel;
  m_NumberOfComponents = nComp;
  m_SamplesDirty = false;
}

void UnsupervisedClustering::SetNumberOfClusters(int nClusters)
{
  if(m_NumberOfClusters != nClusters)
    {
    m_NumberOfClusters = nClusters;
    }
}

void UnsupervisedClustering::SetNumberOfSamples(int nSamples)
{
  if(m_NumberOfSamples != nSamples)
    {
    m_NumberOfSamples = nSamples;
    m_SamplesDirty = true;
    }
}

void UnsupervisedClustering::InitializeClusters()
{
  this->InitializeEM();
}

void UnsupervisedClustering::InitializeEM()
{
  // Make sure the data source is specified
  assert(m_DataSource);

  // Make sure samples exist
  if(m_SamplesDirty || m_Samples.empty())
    this->SampleDataSource();

  if(m_ClusteringEM)
    {
    delete m_ClusteringEM;
    delete m_ClusteringInitializer;
    }

  // Allocate the EM algorithm
    m_ClusteringEM = new EMGaussianMixtures(m_Samples, m_NumberOfClusters);

    // Allocate the K means ++
    m_ClusteringInitializer = new KMeansPlusPlus(m_Samples, m_NumberOfClusters);

    m_ClusteringInitializer->Initialize();

    m_ClusteringEM->SetGaussianMixtureModel(m_ClusteringInitializer->GetGaussianMixtureModel());

    m_ClusteringEM->SetMaxIteration(10);

    // Get the GMM
    m_MixtureModel = m_ClusteringEM->GetGaussianMixtureModel();

    // Sort the clusters based on center samples
    SortClustersByRelevance();
}

void UnsupervisedClustering::SortClustersByRelevance()
{
  int ng = m_MixtureModel->GetNumberOfGaussians();
  vnl_vector<double> log_pdf(ng), log_w(ng), w(ng);

  // the array to sort
  typedef std::pair<double, int> RelevancePair;
  std::vector<RelevancePair> rel(ng);

  for(int k = 0; k < ng; k++)
    {
    log_w[k] = log(m_MixtureModel->GetWeight(k));
    w[k] = m_MixtureModel->GetWeight(k);
    rel[k].first = 0.0;
    rel[k].second = k;
    }

  for(unsigned int i = 0; i < m_CenterSamples.size(); i++)
    {
    int s = m_CenterSamples[i];
    for(int k = 0; k < ng; k++)
      {
      log_pdf[k] = m_MixtureModel->EvaluateLogPDF(k, m_Samples[s]);
      }

    for(int k = 0; k < ng; k++)
      {
      // Accumulate the posterior
      rel[k].first -= EMGaussianMixtures::ComputePosterior(
            ng, log_pdf.data_block(), w.data_block(), log_w.data_block(), k);
      }
    }

  // Sort the relevance array
  std::sort(rel.begin(), rel.end());

  // Create a new mixture model
  SmartPtr<GaussianMixtureModel> gmmtemp = GaussianMixtureModel::New();
  gmmtemp->Initialize(m_MixtureModel->GetNumberOfComponents(), ng);
  for(int k = 0; k < ng; k++)
    {
    int j = rel[k].second;
    gmmtemp->SetWeight(k, m_MixtureModel->GetWeight(j));
    gmmtemp->SetGaussian(k, m_MixtureModel->GetMean(j), m_MixtureModel->GetCovariance(j));
    }

  // Assign it to the EM
  m_ClusteringEM->SetGaussianMixtureModel(gmmtemp);
  m_MixtureModel = m_ClusteringEM->GetGaussianMixtureModel();
}


void UnsupervisedClustering::Iterate()
{
  m_ClusteringEM->UpdateOnce();
  m_MixtureModel->PrintParameters();
}


