#include "RegionSegmentationMethod.h"

#include "core/image/CrossT.h"
#include "core/image/LineT.h"
#include "core/image/CircleT.h"
#include "core/imagedisplay/ImageDisplay.h"

using namespace OBJREC;
using namespace NICE;
using namespace std;

RegionSegmentationMethod::RegionSegmentationMethod()
{

}

RegionSegmentationMethod::RegionSegmentationMethod ( const NICE::Config * _conf )
{
//   conf = c;
  this->initFromConfig ( _conf );
}

RegionSegmentationMethod::~RegionSegmentationMethod()
{

}

void RegionSegmentationMethod::initFromConfig ( const NICE::Config * _conf, const std::string & _confSection )
{
  //nothing to do here
}

//überprüfe, ob Elterpixel neugelabelt wurde, wenn ja dann label auch Kind neu
inline int rootOf ( int *l, int j )
{
  while ( l[j] != j ) {
    // short path compression
    l[j] = l[l[j]];
    j    = l[j];
  }
  return j;
}

///////////////////// ///////////////////// /////////////////////
//                      SEGMENTATION STUFF
///////////////////// ///////////////////// ////////////////// 

int RegionSegmentationMethod::transformSegmentedImg ( const NICE::ColorImage & img, NICE::Matrix & mask ) const
{
  int xsize = img.width();
  int ysize = img.height();
  int n = xsize * ysize;
  int *l = new int[n];
  int v[3], vo[3];

  int j = 0;
  for ( int y = 0 ; y < ysize ; y++ ) // horizontal
  {
    l[j] = j;
    vo[0] = img.getPixel ( 0, y, 2 );
    vo[1] = img.getPixel ( 0, y, 0 );
    vo[2] = img.getPixel ( 0, y, 1 );
    j++;
    for ( int x = 1 ; x < xsize ; x++, j++ )
    {
      v[0] = img.getPixel ( x, y, 2 );
      v[1] = img.getPixel ( x, y, 0 );
      v[2] = img.getPixel ( x, y, 1 );
      if ( v[0] == vo[0] && v[1] == vo[1] && v[2] == vo[2] )
      {
        l[j] = l[j-1];
      }
      else
      {
        l[j] = j;
      }
      vo[0] = v[0];
      vo[1] = v[1];
      vo[2] = v[2];
    }
  }

  j = 0;
  for ( int x = 0 ; x < xsize ; x++ ) //vertikal
  {
    vo[0] = img.getPixel ( x, 0, 2 );
    vo[1] = img.getPixel ( x, 0, 0 );
    vo[2] = img.getPixel ( x, 0, 1 );
    for ( int y = 1 ; y < ysize ; y++ )
    {
      j = y * ( xsize ) + x;
      v[0] = img.getPixel ( x, y, 2 );
      v[1] = img.getPixel ( x, y, 0 );
      v[2] = img.getPixel ( x, y, 1 );
      if ( v[0] == vo[0] && v[1] == vo[1] && v[2] == vo[2] ) {
        int rj = rootOf ( l, j );
        int rn = rootOf ( l, j - xsize );
        if ( rj != rn ) {
          l[rj] = rn;
        }
      }
      vo[0] = v[0];
      vo[1] = v[1];
      vo[2] = v[2];
    }
  }

  j = 0;
  for ( int x = 1 ; x < xsize ; x++ ) //links oben
  {
    vo[0] = img.getPixel ( x - 1, 0, 2 );
    vo[1] = img.getPixel ( x - 1, 0, 0 );
    vo[2] = img.getPixel ( x - 1, 0, 1 );
    for ( int y = 1 ; y < ysize ; y++ )
    {
      j = y * ( xsize ) + x;
      v[0] = img.getPixel ( x, y, 2 );
      v[1] = img.getPixel ( x, y, 0 );
      v[2] = img.getPixel ( x, y, 1 );
      if ( v[0] == vo[0] && v[1] == vo[1] && v[2] == vo[2] ) {
        int rj = rootOf ( l, j );
        int rn = rootOf ( l, j - xsize - 1 );
        if ( rj != rn ) {
          l[rj] = rn;
        }
      }
      vo[0] = img.getPixel ( x - 1, y, 2 );
      vo[1] = img.getPixel ( x - 1, y, 0 );
      vo[2] = img.getPixel ( x - 1, y, 1 );
    }
  }

  j = 0;
  for ( int x = xsize - 2 ; x >= 0 ; x-- ) //rechts oben
  {
    vo[0] = img.getPixel ( x + 1, 0, 2 );
    vo[1] = img.getPixel ( x + 1, 0, 0 );
    vo[2] = img.getPixel ( x + 1, 0, 1 );
    for ( int y = 1 ; y < ysize ; y++ )
    {
      j = y * ( xsize ) + x;
      v[0] = img.getPixel ( x, y, 2 );
      v[1] = img.getPixel ( x, y, 0 );
      v[2] = img.getPixel ( x, y, 1 );
      if ( v[0] == vo[0] && v[1] == vo[1] && v[2] == vo[2] ) {
        int rj = rootOf ( l, j );
        int rn = rootOf ( l, j - xsize + 1 );
        if ( rj != rn ) {
          l[rj] = rn;
        }
      }
      vo[0] = img.getPixel ( x + 1, y, 2 );
      vo[1] = img.getPixel ( x + 1, y, 0 );
      vo[2] = img.getPixel ( x + 1, y, 1 );
    }
  }

  int size = 0;

  // relabel
  std::map<int, int> bijection;
  for ( j = 0 ; j < n ; j++ )
  {
    if ( l[j] == j ) {
      bijection[j] = size;
      size++;
    }
  }

  if ( size <= 0 ) {
    fprintf ( stderr, "RegionSegmentationMethod:: no components of this color found !!\n" );
    return -1;
  }

  mask.resize ( xsize, ysize );

  for ( j = 0 ; j < n ; j++ )
  {
    int label = rootOf ( l, j );
    int x = j % ( xsize );
    int y = j / ( xsize );

    mask ( x, y ) = bijection[label];
  }
  delete [] l;

  return size;
}

///////////////////// ///////////////////// /////////////////////
//                         GET / SET
///////////////////// ///////////////////// /////////////////////  

void RegionSegmentationMethod::getGraphRepresentation ( const NICE::ColorImage & cimg, NICE::Matrix & mask, RegionGraph & rg )
{
  int regioncount = segRegions ( cimg, mask );
  rg.computeGraph ( mask, regioncount );
}

int RegionSegmentationMethod::segRegions ( const NICE::ColorImage & cimg, NICE::Matrix & mask ) const
{
  Image img ( cimg.width(), cimg.height() );
  for ( int y = 0; y < cimg.height(); y++ )
  {
    for ( int x = 0; x < cimg.width(); x++ )
    {
      int val = 0;
      for ( int i = 0; i < 3; i++ )
        val += cimg.getPixel ( x, y, i );
      val /= 3;
      img.setPixel ( x, y, val );
    }
  }
  cerr << "no color segmentation method, using grayimage segmentation instead" << endl;
  return segRegions ( img, mask );
}

int RegionSegmentationMethod::segRegions (
    const NICE::MultiChannelImage3DT<double> & img,
    NICE::MultiChannelImageT<int> & mask,
    const int isGray
    ) const
{
	const int xsize = img.width();
	const int ysize = img.height();
	const int zsize = img.depth();

	int amountRegions = -1;
	mask.reInit( xsize, ysize, zsize );
	for (int z = 0; z < zsize; z++)
	{
		NICE::Matrix reg;
		int aR;
    if ( isGray )
		{
      NICE::Image slice = img.getChannel(z);
      aR = segRegions( slice, reg );
		} else {
      NICE::ColorImage slice = img.getColor(z);
      aR = segRegions( slice, reg );
		}

		if ( aR > amountRegions ) amountRegions = aR;

		for (int y = 0; y < ysize; y++)
		{
			for (int x = 0; x < xsize; x++)
			{
				mask.set( x, y, reg(x,y), (uint)z );
			}
		}
	}

  return amountRegions;
}

void RegionSegmentationMethod::markContours ( const NICE::ColorImage & cimg, NICE::Matrix & mask, std::vector<int> &color, NICE::ColorImage &marked )
{
  // mark contours
  for ( int y = 1; y < cimg.height() - 1; y++ )
  {
    for ( int x = 1; x < cimg.width() - 1; x++ )
    {
      bool diff = false;
      for ( int i = -1; i < 2; i++ )
      {
        for ( int j = -1; j < 2; j++ )
        {
          if ( mask ( x, y ) != mask ( x + i, y + j ) )
            diff = true;
        }
        if ( diff )
          break;
      }
      if ( diff )
      {
        for(int c = 0; c < 3; c++)
          marked.setPixel ( x, y, c, color[c] );
      }
      else
      {
        for(int c = 0; c < 3; c++)
          marked.setPixel ( x, y, c, cimg.getPixel(x,y,c) );
      }
    }
  }
}

void RegionSegmentationMethod::visualizeGraphRepresentation ( const NICE::ColorImage & cimg, NICE::Matrix & mask )
{
  RegionGraph rg;
  getGraphRepresentation ( cimg, mask, rg );

  vector<Node*> g;
  rg.get ( g );

  Image overlay ( cimg.width(), cimg.height() );
  overlay.set ( 0 );

  // mark contours
  for ( int y = 1; y < cimg.height() - 1; y++ )
  {
    for ( int x = 1; x < cimg.width() - 1; x++ )
    {
      bool diff = false;
      for ( int i = -1; i < 2; i++ )
      {
        for ( int j = -1; j < 2; j++ )
        {
          if ( mask ( x, y ) != mask ( x + i, y + j ) )
            diff = true;
        }
        if ( diff )
          break;
      }
      if ( diff )
        overlay.setPixel ( x, y, 3 );
    }
  }

  // mark graph
  for ( int i = 0 ; i < ( int ) g.size() ; i++ )
  {
    int x, y;
    g[i]->getCentroid ( x, y );
    Cross cross ( Coord ( x, y ) , 7 );
    overlay.draw ( cross, 1 );
    vector<Node*> nb;
    g[i]->getNeighbors ( nb );
    for ( int j = 0; j < ( int ) nb.size(); j++ )
    {
      int xn, yn;
      nb[j]->getCentroid ( xn, yn );
      Line line ( Coord ( x, y ), Coord ( xn, yn ) );
      overlay.draw ( line, 2 );
    }
  }

  showImageOverlay ( cimg, overlay, "Regions" );
}