/*=========================================================================
    CMVision.cc
  -------------------------------------------------------------------------
    Implementation of the CMVision real time Color Machine Vision library
  -------------------------------------------------------------------------
    Copyright 1999, 2000         #### ### ### ## ## ## #### ##  ###  ##  ##
    James R. Bruce              ##    ####### ## ## ## ##   ## ## ## ######
    School of Computer Science  ##    ## # ## ## ## ##  ### ## ## ## ## ###
    Carnegie Mellon University   #### ##   ##  ###  ## #### ##  ###  ##  ##
  -------------------------------------------------------------------------
    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
  -------------------------------------------------------------------------
    Revision History:
      1999-11-18:  Initial release version (JRB)
      2000-05-20:  Added Bugfixes from Peter,
                   fixed bounding box bug (JRB)
      2000-06-04:  Some other minor fixes (JRB)
  =========================================================================*/

#include "cmvision.h"

//==== Utility Functions ===========================================//
// These could be coded as macros, but the inline versions seem to
// optimize better, and tend to have cleaner definitions

// sum of integers over range [x,x+w)
inline int range_sum(int x,int w)
{
  return(w*(2*x + w-1) / 2);
}

// returns maximal value of two parameters
template <class num>
inline num max(num a,num b)
{
  return((a > b)? a : b);
}

// returns minimal value of two parameters
template <class num>
inline num min(num a,num b)
{
  return((a < b)? a : b);
}

// returns index of least significant set bit
template <class num>
num bottom_bit(num n)
{
  int i = 0;
  if(!n) return(0);
  while(!(n&(1<<i))) i++;
  return(i + 1);
}

// returns index of most significant set bit
template <class num>
num top_bit(num n)
{
  int i = 1;
  if(!n) return(0);
  while(n>>i) i++;
  return(i);
}


//==== Class Implementation ========================================//

void CMVision::classifyFrame(yuv422 *img,unsigned *map)
// Classifies an image passed in as img, saving bits in the entries
// of map representing which thresholds that pixel satisfies.
{
  int i,m,s;
  yuv422 p;

  unsigned *uclas = uclass; // Ahh, the joys of a compiler that
  unsigned *vclas = vclass; //   has to consider pointer aliasing
  unsigned *yclas = yclass;

  s = width * height;
  for(i=0; i<s; i+=2){
    p = img[i/2];
    m = uclas[p.u] & vclas[p.v];
    map[i + 0] = m & yclas[p.y1];
    map[i + 1] = m & yclas[p.y2];
  }
}

int CMVision::encodeRuns(rle *out,unsigned *map)
// Changes the flat array version of the threshold satisfaction map
// into a run length encoded version, which speeds up later processing
// since we only have to look at the points where values change.
{
  int x,y,j,l;
  unsigned m;
  int size;
  unsigned *row;
  rle r;

  size = width * height;

  j = 0;
  for(y=0; y<height; y++){
    row = &map[y * width];
    row[width - 1] = CMV_NONE;

    x = 0;
    while(x < width){
      m = row[x];
      l = x;
      while(row[x] == m) x++;
      x += (row[x] == CMV_NONE);

      r.color  = m;
      r.length = x - l;
      r.parent = j;
      out[j++] = r;
      if(j >= CMV_MAX_RUNS) return(0);
    }
  }

  return(j);
}

void CMVision::connectComponents(rle *map,int num)
// Connect components using four-connecteness so that the runs each
// identify the global parent of the connected region they are a part
// of.  It does this by scanning adjacent rows and merging where similar
// colors overlap.  Used to be union by rank w/ path compression, but now
// is just uses path compression as the global parent index seems to be
// a faster approximation of rank in practice.
// WARNING: This code is *extremely* complicated and twitchy.  It appears
//   to work for all inputs, but minor changes can easily cause big
//   problems.  Read the papers on this library and have a good
//   understanding of tree-based union find before you touch it
{
  int x1,x2;
  int l1,l2;
  rle r1,r2;
  int i,p,s,n;

  l1 = l2 = 0;
  x1 = x2 = 0;

  // Lower scan begins on second line, so skip over first
  while(x1 < width){
    x1 += map[l1++].length;
  }
  x1 = 0;

  // Do rest in lock step
  r1 = map[l1];
  r2 = map[l2];
  s = l1;
  while(l1 < num){
    if(r1.color==r2.color && r1.color){
      if((x1>=x2 && x1<x2+r2.length) || (x2>=x1 && x2<x1+r1.length)){
        if(s != l1){
          map[l1].parent = r1.parent = r2.parent;
          s = l1;
        }else{
          // find terminal roots of each path
          n = r1.parent;
          while(n != map[n].parent) n = map[n].parent;
          p = r2.parent;
          while(p != map[p].parent) p = map[p].parent;

          // must use smaller of two to preserve DAGness!
          if(n < p){
            map[p].parent = n;
          }else{
            map[n].parent = p;
          }
        }
      }
    }

    // Move to next point where values may change
    if(x1+r1.length < x2+r2.length){
      x1 += r1.length;
      r1 = map[++l1];
    }else{
      x2 += r2.length;
      r2 = map[++l2];
    }
  }

  // Now we need to compress all parent paths
  for(i=0; i<num; i++){
    p = map[i].parent;
    if(p > i){
      while(p != map[p].parent) p = map[p].parent;
      map[i].parent = p;
    }else{
      map[i].parent = map[p].parent;
    }
  }

  // Ouch, my brain hurts.
}

int CMVision::extractRegions(region *reg,rle *rmap,int num)
// Takes the list of runs and formats them into a region table,
// gathering the various statistics we want along the way.
// num is the number of runs in the rmap array, and the number of
// unique regions in reg[] (< CMV_MAX_REGIONS) is returned.
// Implemented as a single pass over the array of runs.
{
  int x,y,i;
  int b,n,a;
  rle r;

  x = y = n = 0;
  for(i=0; i<num; i++){
    r = rmap[i];

    if(r.color){
      if(r.parent == i){
        // Add new region if this run is a root (i.e. self parented)
        rmap[i].parent = b = n;  // renumber to point to region id
        reg[b].color = bottom_bit(r.color) - 1;
        reg[b].area = r.length;
        reg[b].x1 = x;
        reg[b].y1 = y;
        reg[b].x2 = x + r.length;
        reg[b].y2 = y;
        reg[b].sum_x = range_sum(x,r.length);
        reg[b].sum_y = y * r.length;
        n++;
        if(n >= CMV_MAX_REGIONS) return(CMV_MAX_REGIONS);
      }else{
        // Otherwise update region stats incrementally
        b = rmap[r.parent].parent;  // update to point to region id
        reg[b].area += r.length;
        reg[b].x2 = max(x + r.length,reg[b].x2);
        reg[b].x1 = min(x,reg[b].x1);
        reg[b].y2 = y; // last set by lowest run
        reg[b].sum_x += range_sum(x,r.length);
        reg[b].sum_y += y * r.length;
      }
    }

    // step to next location
    x = (x + r.length) % width;
    y += (x == 0);
  }

  // calculate centroids from stored temporaries
  for(i=0; i<n; i++){
    a = reg[i].area;
    reg[i].cen_x = (float)reg[i].sum_x / a;
    reg[i].cen_y = (float)reg[i].sum_y / a;
  }

  return(n);
}

int CMVision::separateRegions(region *reg,int num)
// Splits the various regions in the region table a separate list
// for each color.  The lists are threaded through the table using
// the region's 'next' field.  Returns the maximal area of the
// regions, which we use below to speed up sorting.
{
  region *p;
  int i,l;
  int area,max_area;

  // clear out the region table
  for(i=0; i<CMV_MAX_COLORS; i++){
    region_count[i] = 0;
    region_list[i] = NULL;
  }

  // step over the table, adding successive
  // regions to the front of each list
  max_area = 0;
  for(i=0; i<num; i++){
    p = &reg[i];
    area = p->area;
    if(area >= CMV_MIN_AREA){
      if(area > max_area) max_area = area;
      l = p->color;
      region_count[l]++;
      p->next = region_list[l];
      region_list[l] = p;
    }
  }

  return(max_area);
}

// These are the tweaking values for the radix sort given below
// Feel free to change them, though these values seemed to work well
// in testing.  Don't worry about extra passes to get all 32 bits of
// the area; the implementation only does as many passes as needed to
// touch the most significant set bit (MSB of biggest region's area)
#define CMV_RBITS 6
#define CMV_RADIX (1 << CMV_RBITS)
#define CMV_RMASK (CMV_RADIX-1)

CMVision::region *CMVision::sortRegionListByArea(region *list,int passes)
// Sorts a list of regions by their area field.
// Uses a linked list based radix sort to process the list.
{
  region *tbl[CMV_RADIX],*p,*pn;
  int slot,shift;
  int i,j;

  // Handle trivial cases
  if(!list || !list->next) return(list);

  // Initialize table
  for(j=0; j<CMV_RADIX; j++) tbl[j] = NULL;

  for(i=0; i<passes; i++){
    // split list into buckets
    shift = CMV_RBITS * i;
    p = list;
    while(p){
      pn = p->next;
      slot = ((p->area) >> shift) & CMV_RMASK;
      p->next = tbl[slot];
      tbl[slot] = p;
      p = pn;
    }

    // integrate back into partially ordered list
    list = NULL;
    for(j=0; j<CMV_RADIX; j++){
      p = tbl[j];
      tbl[j] = NULL;  // clear out table for next pass
      while(p){
        pn = p->next;
        p->next = list;
        list = p;
        p = pn;
      }
    }
  }

  return(list);
}

void CMVision::sortRegions(int max_area)
// Sorts entire region table by area, using the above
// function to sort each threaded region list.
{
  int i,p;

  // do minimal number of passes sufficient to touch all set bits
  p = top_bit((max_area + CMV_RBITS-1) / CMV_RBITS);

  // sort each list
  for(i=0; i<CMV_MAX_COLORS; i++){
    region_list[i] = sortRegionListByArea(region_list[i],p);
  }
}

//==== Interface/Public Functions ==================================//

#define ZERO(x) memset(x,0,sizeof(x))

void CMVision::clear()
{
  ZERO(yclass);
  ZERO(uclass);
  ZERO(vclass);

  ZERO(region_list);
  ZERO(region_count);

  ZERO(colors);

  map = NULL;
}

bool CMVision::initialize(int nwidth,int nheight)
// Initializes library to work with images of specified size
{
  width = nwidth;
  height = nheight;

  if(map) delete(map);

  map = new unsigned[width * height];

  return(map != NULL);
}

// sets bits in k in array arr[l..r]
template <class num>
void set_bits(num *arr,int len,int l,int r,num k)
{
  int i;

  l = max(l,0);
  r = min(r+1,len);

  for(i=l; i<r; i++) arr[i] |= k;
}

template <class num>
void clear_bits(num *arr,int len,int l,int r,num k)
{
  int i;

  l = max(l,0);
  r = min(r+1,len);

  k = ~k;
  for(i=l; i<r; i++) arr[i] &= k;
}

#define CMV_STATE_SCAN   0
#define CMV_STATE_COLORS 1
#define CMV_STATE_THRESH 2
#define CMV_MAX_BUF 246

bool CMVision::loadOptions(char *filename)
// Loads in options file specifying color names and representative
// rgb triplets.  Also loads in color class threshold values.
{
  char buf[CMV_MAX_BUF],str[CMV_MAX_BUF];
  FILE *in;
  int state,i,n;

  int r,g,b;
  double m;
  color_info *c;

  int y1,y2,u1,u2,v1,v2;
  unsigned k;

  // Open options file
  in = fopen(filename,"rt");
  if(!in) return(false);

  // Clear out previously set options
  for(i=0; i<CMV_COLOR_LEVELS; i++){
    yclass[i] = uclass[i] = vclass[i] = 0;
  }
  for(i=0; i<CMV_MAX_COLORS; i++){
    if(colors[i].name){
      delete(colors[i].name);
      colors[i].name = NULL;
    }
  }

  // Loop ever lines, processing via a simple parser
  state = 0;
  while(fgets(buf,CMV_MAX_BUF,in)){
    switch(state){
      case CMV_STATE_SCAN:
        n = sscanf(buf,"[%s",str);
        if(n == 1){
          if(!strncasecmp(str,"colors]",CMV_MAX_BUF)){
            state = CMV_STATE_COLORS;
            i = 0;
	  }else if(!strncasecmp(str,"thresholds]",CMV_MAX_BUF)){
	    state = CMV_STATE_THRESH;
            i = 0;
	  }else{
            printf("CMVision: Ignoring unknown option header '%s'.\n",str);
          }
	}
        break;
      case CMV_STATE_COLORS:
        n = sscanf(buf,"(%d,%d,%d) %lf %s",&r,&g,&b,&m,str);
        if(n == 5){
          // printf("(%d,%d,%d) %lf '%s'\n",r,g,b,m,str);
          if(i < CMV_MAX_COLORS){
            c = &colors[i];
            c->color.red   = r;
            c->color.green = g;
            c->color.blue  = b;
            c->name  = strdup(str);
            c->merge = m;
            i++;
	  }else{
	    printf("CMVision: Too many colors, ignoring '%s'.\n",str);
	  }
	}else if(n == 0){
          state = CMV_STATE_SCAN;
        }
        break;
      case CMV_STATE_THRESH:
        n = sscanf(buf,"(%d:%d,%d:%d,%d:%d)",&y1,&y2,&u1,&u2,&v1,&v2);
        if(n == 6){
          // printf("(%d:%d,%d:%d,%d:%d)\n",y1,y2,u1,u2,v1,v2);
          if(i < CMV_MAX_COLORS){
            c = &colors[i];
            c->y_low = y1;  c->y_high = y2;
            c->u_low = u1;  c->u_high = u2;
            c->v_low = v1;  c->v_high = v2;

            k = (1 << i);
            set_bits(yclass,CMV_COLOR_LEVELS,y1,y2,k);
            set_bits(uclass,CMV_COLOR_LEVELS,u1,u2,k);
            set_bits(vclass,CMV_COLOR_LEVELS,v1,v2,k);
            i++;
	  }else{
	    printf("CMVision: Too many thresholds.\n");
	  }
	}else if(n == 0){
          state = CMV_STATE_SCAN;
        }
        break;
    }
  }

  /*
  for(i=0; i<CMV_COLOR_LEVELS; i++){
    printf("%08X %08X %08X\n",yclass[i],uclass[i],vclass[i]);
  }
  */

  fclose(in);

  return(true);
}

bool CMVision::saveOptions(char *filename)
{
  color_info *c;
  FILE *out;
  int i;

  out = fopen(filename,"wt");
  if(!out) return(false);

  fprintf(out,"[Colors]\n");
  i = 0;
  while(colors[i].name){
    c = &colors[i];
    fprintf(out,"(%3d,%3d,%3d) %lf %s\n",
      c->color.red,c->color.green,c->color.blue,
      c->merge,c->name);
    i++;
  }

  fprintf(out,"\n[Thresholds]\n");
  i = 0;
  while(colors[i].name){
    c = &colors[i];
    fprintf(out,"(%3d:%3d,%3d:%3d,%3d:%3d)\n",
      c->y_low,c->y_high,
      c->u_low,c->u_high,
      c->v_low,c->v_high);
    i++;
  }

  fclose(out);

  return(true);
}

void CMVision::close()
{
  if(map) delete(map);
  map = NULL;
}


//==== Vision Testing Functions ====================================//

bool CMVision::testClassify(rgb *out,yuv422 *image)
{
  int i,s;
  rgb black = {0,0,0};

  if(!image || !out) return(false);

  classifyFrame(image,map);

  s = width * height;

  i = 0;
  while(i < s){
    while(i<s && !map[i]){
      out[i] = black;
      i++;
    }
    while(i<s && map[i]){
      out[i] = colors[bottom_bit(map[i])-1].color;
      i++;
    }
  }

  return(true);
}

bool CMVision::getThreshold(int color,
       int &y_low,int &y_high,
       int &u_low,int &u_high,
       int &v_low,int &v_high)
{
  color_info *c;

  if(color<0 || color>=CMV_MAX_COLORS) return(false);

  c = &colors[color];
  y_low = c->y_low;  y_high = c->y_high;
  u_low = c->u_low;  u_high = c->u_high;
  v_low = c->v_low;  v_high = c->v_high;

  return(true);
}

bool CMVision::setThreshold(int color,
       int y_low,int y_high,
       int u_low,int u_high,
       int v_low,int v_high)
{
  color_info *c;
  unsigned k;

  if(color<0 || color>=CMV_MAX_COLORS) return(false);

  c = &colors[color];
  k = 1 << color;

  clear_bits(yclass,CMV_COLOR_LEVELS,c->y_low,c->y_high,k);
  clear_bits(uclass,CMV_COLOR_LEVELS,c->u_low,c->u_high,k);
  clear_bits(vclass,CMV_COLOR_LEVELS,c->v_low,c->v_high,k);

  c->y_low = y_low;  c->y_high = y_high;
  c->u_low = u_low;  c->u_high = u_high;
  c->v_low = v_low;  c->v_high = v_high;

  set_bits(yclass,CMV_COLOR_LEVELS,y_low,y_high,k);
  set_bits(uclass,CMV_COLOR_LEVELS,u_low,u_high,k);
  set_bits(vclass,CMV_COLOR_LEVELS,v_low,v_high,k);

  return(true);
}

//==== Main Vision Functions =======================================//

bool CMVision::processFrame(yuv422 *image)
{
  int runs;
  int regions;
  int max_area;

  if(!image) return(false);

  classifyFrame(image,map);
  runs = encodeRuns(rmap,map);
  connectComponents(rmap,runs);

  regions = extractRegions(region_table,rmap,runs);
  max_area = separateRegions(region_table,regions);
  sortRegions(max_area);

  return(true);
}

int CMVision::numRegions(int color_id)
{
  if(color_id<0 || color_id>=CMV_MAX_COLORS) return(CMV_NONE);
  return(region_count[color_id]);
}

CMVision::region *CMVision::getRegions(int color_id)
{
  if(color_id<0 || color_id>=CMV_MAX_COLORS) return(NULL);
  return(region_list[color_id]);
}