/**
   Toboggan.cpp

   Quick test to see if tobogganed regions would be useful for image effects.
   May not be entirely correct, and could be done more efficiently on higher
   levels of the pyramid, if the averages and variances for regions were merged
   when the regions are merged instead of recomputing from the images.

   TODO: merge average, variances, counts, and labeled regions into a data
   structure for the regions.  Speed-up (see note above).
   
   Clean-up interface for toboggan (primarily output), and remove the saving
   that occurs in that function.

   Neil Birkbeck
   Copyright (c) 2009
*/
#include "glwindow.h"
#include <GL/gl.h>

#include <stdio.h>
#include <vector>
#include <set>

#include <nimage/image.h>
#include <nmath/vec3.h>
#include <nmath/vec2.h>

using namespace nacb;

/**
   \brief Compute the averages and variances from the input image, provided labels,
          and number of labels.
*/
void getAverageVar(const nacb::Imagef & image,
		   const nacb::Image32 & label,
		   int nlabel,
		   std::vector<Vec3f> & averages,
		   std::vector<double> & var,
		   std::vector<int> & counts) {

  averages = std::vector<Vec3f>(nlabel, Vec3f(0, 0, 0));
  var = std::vector<double>(nlabel, 0.0);
  counts = std::vector<int>(nlabel, 0);

  for(int y=0; y<image.h; y++){
    for(int x=0; x<image.w; x++){
      Vec3f color(image(x, y, 0),
		  image(x, y, 1),
		  image(x, y, 2));

      averages[label(x, y)] += color;
      var[label(x, y)] += color.x*color.x + 
	                  color.y*color.y + 
	                  color.z*color.z;
      counts[label(x, y)]++;
    }
  }

  for(int i=0; i<averages.size(); i++){
    averages[i] /= counts[i];
    
    double agray = (averages[i].x + averages[i].y + averages[i].z)/3;
    var[i] = var[i] / (3*counts[i]) - agray * agray;
  }
}


/**
   \brief Compute the graph for a set of labels.
*/
std::vector<std::set<int> > getGraph(nacb::Image32 & label, int nlabel){

  // Compute a graph.
  std::vector<std::set<int> > edges(nlabel);
  for(int y=0; y<label.h-1; y++){
    for(int x=0; x<label.w-1; x++){

      if(label(x, y) != label(x+1, y)){
	edges[label(x, y)].insert(label(x+1, y));
	edges[label(x+1, y)].insert(label(x, y));
      }

      if(label(x, y) != label(x, y+1)){
	edges[label(x, y)].insert(label(x, y+1));
	edges[label(x, y+1)].insert(label(x, y));
      }
    }
  }
  return edges;
}


/**
   \brief Map a label image to the average color for the label.
 */
nacb::Imagef mapOutput(const nacb::Image32 & label,
		       const std::vector<Vec3f> & averages) {
  nacb::Imagef output(label.w, label.h, 3);

  for (int y=0; y<label.h; y++) {
    for (int x=0; x<label.w; x++) {
      int lab = label(x, y);
      output(x, y, 0) = averages[lab].x;      
      output(x, y, 1) = averages[lab].y;      
      output(x, y, 2) = averages[lab].z;      
    }
  }

  return output;
}


/**
   \brief Merge a set of labels using the parent label (e.g., the next level
          of the segmentation).
*/
nacb::Image32 mergeLabels(nacb::Image32 & labels,
			  std::vector<int> & glabel) {
  nacb::Image32 labelsMerged(labels.w, labels.h, 1);
  
  for(int y=0; y<labels.h; y++){
    for(int x=0; x<labels.w; x++){
      labelsMerged(x, y) = glabel[labels(x, y)];
    }
  }
  
  return labelsMerged;
}


/**
   \brief Compute sobel edges for the input image (using gray-scale)
*/
nacb::Imagef sobel(const nacb::Imagef & input) {
  nacb::Imagef fx(3, 3, 1);
  nacb::Imagef fy(3, 3, 1);

  fx = 0;
  fy = 0;

  fx(0, 0) = -1;
  fx(0, 1) = -2;
  fx(0, 2) = -1;
  
  fx(2, 0) = 1;
  fx(2, 1) = 2;
  fx(2, 2) = 1;

  fy(0, 0) = -1;
  fy(1, 0) = -2;
  fy(2, 0) = -1;
  
  fy(0, 2) = 1;
  fy(1, 2) = 2;
  fy(2, 2) = 1;

  nacb::Imagef gray = (input.getChannel(0) + 
		       input.getChannel(1) + 
		       input.getChannel(2))*0.3333;
  nacb::Imagef gx = gray.convolve(fx);
  nacb::Imagef gy = gray.convolve(fy);
  nacb::Imagef edges(gx.w, gx.h, 1);

  for(int i=0; i<edges.w*edges.h; i++)
    edges.data[i] = sqrt(gx.data[i]*gx.data[i] + gy.data[i]*gy.data[i]);

  return edges;
}


/**
   \brief Merge regions using the second region score (utilizes averages
          and variances); required for all levels but the first.

   \param edges     the graph connectivity.
   \param averages  the average colors of the components in the graph.
   \param var       the variances of graph components.
   \param counts    the membership count of graph components.
   \param glabel    output labels (one for each component).

   \returns the number of labels in the output
*/
int  mergeRegions(const std::vector<std::set<int> > & edges,
		  const std::vector<Vec3f> & averages,
		  const std::vector<double> & var,
		  const std::vector<int> & counts,
		  std::vector<int> & glabel) {
  int nlabel = averages.size();
  glabel = std::vector<int>(nlabel, -1);
  int nlabels2 = 0;

  printf("Merging graphed regions.\n");
  // Merge graphed regions.
  for (int i=0; i<nlabel; i++){
    if (glabel[i] >= 0)
      continue;
    
    int cur = i;
    std::vector<int> visited;

    while (glabel[cur] == -1) {
      std::set<int>::iterator f, e = edges[cur].end();
      float lowest = 1e10;
      int  bestIndex = -1;

      glabel[cur] = nlabels2;
      visited.push_back(cur);

      for (f = std::set<int>::iterator(edges[cur].begin()); f != e; ++f) {
	// Cache score?
	int n = *f;
	Vec3f cdiff = (averages[cur] - averages[n]);

	double t2 = ((counts[cur] * counts[n]) * cdiff.dot(cdiff)) /
	  (counts[cur] * var[n] + counts[n] * var[cur]);
	
	if (t2 < lowest) {
	  lowest = t2;
	  bestIndex = n;
	}
      }
      // This means there was no neighbors...shouldn't happen.
      if (bestIndex == -1 || lowest > 2000.0) {
	break;
      }
      cur = bestIndex;
    }

    if(glabel[cur] != nlabels2) {
      // Fix up the labels.
      for (int i=0; i<visited.size(); i++) {
	glabel[visited[i]] = glabel[cur];
      }
    }
    else {
      nlabels2++;
    }
  }
  return nlabels2;
}


/**
   \brief Compute tobogganed segmentation for the input image.
   
   Calling interface not complete.
   Returns a debug image.  
*/
nacb::Imagef toboggan(const nacb::Imagef & image){
  nacb::Imagef gray = (image.getChannel(0) + 
		       image.getChannel(1) + 
		       image.getChannel(2))*0.33333;
  nacb::Imagef g = gray.gradient(0);
  nacb::Image32 label(image.w, image.h, 1);
  nacb::Imagef gx = g.getChannel(0);
  nacb::Imagef gy = g.getChannel(1);
  nacb::Imagef gmag = (gx * gx) + (gy * gy);

  label = -1;
  int nlabel = 0;

  for(int y=0; y<g.h; y++){
    for(int x=0; x<g.w; x++){
      int neigh[4][2] = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
      
      int xuse = x;
      int yuse = y;
      
      if(label(x, y) != -1)
	continue;

      // Continue this trail until a label is found, or 
      // until there is no more labels, in which case (the parent is itself).
      std::vector<Vec2<int> > accessed;

      while(label(xuse, yuse) == -1) {
	bool found = false;
	float smallest = gmag(xuse, yuse);
	int xs = 0, ys = 0;	

	accessed.push_back(Vec2<int>(xuse, yuse));
	label(xuse, yuse) = nlabel;

	for(int k=0; k<4; k++) {
	  int x2 = neigh[k][0] + xuse;
	  int y2 = neigh[k][1] + yuse;
	  
	  if( (x2 >=0 ) && (y2 >= 0) &&
	      (x2 < g.w) && (y2 < g.h) ) {
	    
	    if(gmag(x2, y2) <= smallest) {
	      smallest = gmag(x2, y2);
	      found = true;
	      xs = x2;
	      ys = y2;
	    }
	  }
	}
	if(!found) {
	  break;
	}
	xuse = xs;
	yuse = ys;
      }
      int foundLabel = label(xuse, yuse);
      if(foundLabel == nlabel)
	nlabel++;
      else {
	for(int i=0; i<accessed.size(); i++){
	  Vec2<int> pt = accessed[i];
	  label(pt.x, pt.y) = foundLabel;
	}
      }
    }
  }

  /**
     Now recursively compute the upper levels of the segmentation
     hierarchy.
  */
  std::vector<Vec3f> averages;
  std::vector<double> var;
  std::vector<int> counts;
  std::vector<int> glabel;

  nacb::Imagef sobelEdges = sobel(image);
  nacb::Imagef edgeMap(image.w, image.h, 1);
  
  edgeMap = 0;
  
  nacb::Imagef output;
  
  char fname[1024];

  for(int i=0; i<6; i++){
    std::vector<std::set<int> > edges = getGraph(label, nlabel);
    getAverageVar(image, label, nlabel, averages, var, counts);

    if(i == 0)
      output = mapOutput(label, averages);
    
    int nlabel2 = mergeRegions(edges, averages, var, counts, glabel);
    label = mergeLabels(label, glabel);
    nlabel = nlabel2;

    if(i>=4){
      for(int y=0; y<edgeMap.h-1; y++){
	for(int x=0; x<edgeMap.w-1; x++){
	  if( (label(x, y) != label(x + 1, y)) ||
	      (label(x, y) != label(x, y + 1)) ) {
	    edgeMap(x, y) = pow(2, i - 4);
	  }
	}
      }
    }

    snprintf(fname, 1024, "/tmp/level-%d.jpg", i);
    getAverageVar(image, label, nlabel, averages, var, counts);
    mapOutput(label, averages).save(fname);

    printf("%d\n", nlabel);
  }
  
  edgeMap = edgeMap * sobelEdges;
  edgeMap = edgeMap.boxFilter(1).boxFilter(1).boxFilter(1).boxFilter(1);
  edgeMap = edgeMap > 0.025;

  for(int y=0; y<output.h; y++) {
    for(int x=0; x<output.w; x++){
      float e = std::max(0.0f, 1.0f - (edgeMap(x, y))*15.f);

      for(int k=0; k<3; k++)
	output(x, y, k) *= e;
    }
  }
  return output;
}


class Window: public GLWindow 
{
  GLuint tex;
  nacb::Imagef image; 

public:
  Window(const nacb::Imagef & _image) : GLWindow(640, 640) {
    image = _image;
    
    glGenTextures(1, &tex);
    glBindTexture(GL_TEXTURE_2D, tex);

    image = image.boxFilter(2).boxFilter(2);
    nacb::Imagef output = toboggan(image);    
    output.write("/tmp/out.jpg");
    output.initTexture();
  }

  void drawScene() {
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glOrtho(0, image.w, 0, image.h, -1, 1);

    glMatrixMode(GL_MODELVIEW);
    glLoadIdentity();

    glEnable(GL_TEXTURE_2D);
    glBindTexture(GL_TEXTURE_2D, tex);

    glBegin(GL_QUADS);
    glTexCoord2f(0, 1);
    glVertex2f(0, 0);

    glTexCoord2f(1, 1);
    glVertex2f(image.w, 0);

    glTexCoord2f(1, 0);
    glVertex2f(image.w, image.h);
    
    glTexCoord2f(0, 0);
    glVertex2f(0, image.h);
    
    glEnd();
  }
};


int main(int ac, char * av[]){
  nacb::Imagef image(av[1]);
  image = image.resize(image.w/2, image.h/2);

  Window * window = new Window(image);
  window->loop(1);
  return 0;
}