/*
   File:        simple_kdtree.c
   Author:      Andrew W. Moore, Dan Pelleg
   Created:     Sun Apr  5 11:49:13 EDT 1998
   Description: Kd-tree
   $Id: simple_kdtree.c,v 1.36 2002/03/04 14:23:11 awm Exp $

   Copyright 1998, Carnegie Mellon University
*/

#include "simple_kdtree.h"
#include "amgr.h"
#include "debug.h"

static const char *const CVSID = "$Id: simple_kdtree.c,v 1.36 2002/03/04 14:23:11 awm Exp $";

void free_simple_kdtree(simple_kdtree *kd)
{
  if (kd->left != NULL)
    free_simple_kdtree(kd->left);
  if (kd->right != NULL)
    free_simple_kdtree(kd->right);
  free_dyv(kd->sum);
  free_hrect(kd->hr);
  if(kd->points != NULL)
	free_dyv_array(kd->points);
  if(kd->point_ids != NULL)
	free_ivec(kd->point_ids);
  AM_FREE(kd, simple_kdtree);
}

unsigned long simple_kdtree_size(simple_kdtree *node)
{
  unsigned long sz = 0;
  if (node->left != NULL )
    sz += simple_kdtree_size(node->left);
  if (node->right != NULL )
    sz += simple_kdtree_size(node->right);
  if(node->points != NULL)
	sz += dyv_array_size(node->points)*sizeof(int);
  if(node->point_ids != NULL)
    sz += ivec_size(node->point_ids)*sizeof(int);
  return sz + sizeof(*node) + dyv_size(node->sum)*sizeof(double) + 
	+ 2*dyv_size(node->hr->lo)*sizeof(double)
	+ 3*sizeof(dyv);
}

void fprint_simple_kdtree(FILE *s, const char *m1, const char *m2, const simple_kdtree *kd)
{
  char buff[1000];
  sprintf(buff, "%s -> %s", m1, m2);
  fprint_int(s, buff, "split_dim", kd->split_dim);
  fprint_double(s, buff, "split_val", kd->split_val);
  fprint_double(s, buff, "rel_width", kd->rel_width);
  fprint_int(s, buff, "num_points", kd->num_points);
  fprint_dyv(s, buff, "sum", kd->sum);
  fprint_hrect(s, buff, "hr", kd->hr);
}

void draw_simple_kdtree_basic(const simple_kdtree *kd, int depth, int min_depth, int max_depth)
{
  if ( kd != NULL && depth <= max_depth )
  {
    draw_simple_kdtree_basic(kd->left,depth+1,min_depth,max_depth);
    draw_simple_kdtree_basic(kd->right,depth+1,min_depth,max_depth);

    if ( depth >= min_depth )
    {
	  /*	  int cols[] = {AG_RED, AG_GREEN, AG_BLUE, AG_PURPLE, AG_BLACK};
	  int num_cols = sizeof(cols) / sizeof(cols[0]);
      ag_set_pen_color(depth % num_cols);
	  */
      ag_set_pen_color(AG_BLUE);
      ag_hrect(kd->hr);
    }
  }
}

void draw_simple_kdtree(const simple_kdtree *kd, int min_depth, int max_depth)
{
  int col = ag_pen_color();
  draw_simple_kdtree_basic(kd,0,min_depth,max_depth);
  ag_set_pen_color(col);
}

void make_simple_left_and_right_points(const dyv_array *points,
                                       const ivec *point_ids,
                                       int split_dim,
									   double split_val,
									   dyv_array **left, dyv_array **right,
                                       ivec **left_ids, ivec **right_ids)
{
  int i;
  *left = mk_empty_dyv_array();
  *right = mk_empty_dyv_array();
  *left_ids = mk_ivec(0);
  *right_ids = mk_ivec(0);

  for ( i = 0 ; i < dyv_array_size(points) ; i++ )
  {
    const dyv *point = dyv_array_ref(points, i);
    double val = dyv_ref(point, split_dim);
    if ( val <= split_val ) {
      add_to_dyv_array(*left, point);
      add_to_ivec(*left_ids, ivec_ref(point_ids, i));
    } else {
      add_to_dyv_array(*right, point);
      add_to_ivec(*right_ids, ivec_ref(point_ids, i));
    }
  }
}

lrsize make_left_and_right_points(const dyv_array *points,
                                  const ivec *point_ids, 
                                  int npoints_in_points,
								  int split_dim,
                                  double split_val,
                                  dyv_array **left, dyv_array **right,
                                  ivec **left_ids, ivec **right_ids,
                                  int prealloc_size)
{
  int i, node_chunk;
  ivec *lefties = mk_ivec(0);
  dyv *zero = mk_zero_dyv(dyv_size(dyv_array_ref(points, 0)));
  lrsize ret;

  /* find out the left and right points.... */
  for (i=0 ; i<npoints_in_points ; i++)
  {
    const dyv *point = dyv_array_ref(points, i);
    double val = dyv_ref(point, split_dim);
    if ( val <= split_val )
      add_to_sivec(lefties, i);
    else
	  /* skip */;
  }

  /* ... allocate left vector, populate... */
  *left = mk_empty_dyv_array();
  *left_ids = mk_ivec(0);
  for(i=0; i<sivec_size(lefties); i++) {
    const int point_idx = ivec_ref(lefties, i);
    add_to_dyv_array(*left, dyv_array_ref(points, point_idx));
    add_to_ivec(*left_ids, ivec_ref(point_ids, point_idx));
  }
  /* fill and pre-alloc rest */
  /* we want the next multiple of prealloc_size which is at least dyv_array_ref(*left),
     and greater than zero
     what I written here is an attempt to compute this */
  node_chunk = (dyv_array_size(*left) + (prealloc_size - 1))/prealloc_size*prealloc_size; 
  if(node_chunk == 0)
    node_chunk = prealloc_size;
  for(i=dyv_array_size(*left); i<node_chunk; i++) {
    add_to_dyv_array(*left, zero);
    add_to_ivec(*left_ids, KDTREE_NO_POINT);
  }

  /* ... allocate right vector, populate... */
  *right = mk_empty_dyv_array();
  *right_ids = mk_ivec(0);
  for(i=0; i<npoints_in_points; i++) {
	if(!is_in_sivec(lefties, i)) {
	  add_to_dyv_array(*right, dyv_array_ref(points, i));
      add_to_ivec(*right_ids, ivec_ref(point_ids, i));
    }
  }
  /* fill and pre-alloc rest */
  node_chunk = (dyv_array_size(*right) + (prealloc_size - 1))/prealloc_size*prealloc_size; 
  if(node_chunk == 0) node_chunk = prealloc_size;
  for(i=dyv_array_size(*right); i<node_chunk; i++) {
    add_to_dyv_array(*right, zero);
    add_to_ivec(*right_ids, KDTREE_NO_POINT);
  }

  free_dyv(zero);
  ret.left = ivec_size(lefties);
  ret.right = npoints_in_points - ret.left;
  free_ivec(lefties);
  return ret;
}

lrsize get_left_and_right_sizes(const dyv_array *points, int npoints_in_points,
                                int split_dim, double split_val)
{
  int i;
  lrsize ret = {0, 0};

  /* find out the left and right points.... */
  for (i=0 ; i<npoints_in_points ; i++)
  {
    const dyv *point = dyv_array_ref(points, i);
    double val = dyv_ref(point, split_dim);
    if ( val <= split_val )
      ret.left++;
    else
      ret.right++;
  }
  return ret;
}

simple_kdtree *mk_simple_kdtree_basic(const dyv_array *points, const ivec *point_ids,
                                      const dyv *universe, 
									  double min_box_width_rel, int max_leaf_size)
{
  simple_kdtree *kd = AM_MALLOC(simple_kdtree);
  const int num_dims = dyv_size(universe);

  kd->hr = mk_hrect_bounding_dyv_array(points);
  kd->split_dim = hrect_widest_dim(kd->hr);
  kd->split_val = hrect_middle_ref(kd->hr, kd->split_dim);
  kd->rel_width = hrect_width_ref(kd->hr, kd->split_dim) / dyv_ref(universe, kd->split_dim);
  kd->num_points = dyv_array_size(points);
  kd->left = NULL;
  kd->right = NULL;
  kd->cached = NULL;

  /* calculate bias: npoints . sum_{i dims}(hr.width_i/universe.width_i)^2 */
  {
	int i;
	kd->bias = 0;
	for(i=0; i<num_dims; i++) {
	  double tmp = hrect_width_ref(kd->hr, i)/dyv_ref(universe, i);
	  tmp *= tmp;
	  kd->bias += tmp;
	}
	kd->bias *= kd->num_points;
  }

  if (kd->rel_width <= min_box_width_rel
	  && kd->num_points <= max_leaf_size) { /* thin enough -- make a leaf */
    kd->sum = mk_sum_of_dyv_array(points);
	kd->sum_mags = sum_mags_of_dyv_array(points);
	kd->points = mk_copy_dyv_array(points);
    kd->point_ids = mk_copy_ivec(point_ids);
  } else {
    dyv_array *left_rows;
    dyv_array *right_rows;
    ivec *left_ids;
    ivec *right_ids;
	kd->points = NULL;
	kd->point_ids = NULL;

    make_simple_left_and_right_points(points,
                                      point_ids,
                                      kd->split_dim, kd->split_val,
                                      &left_rows, &right_rows,
                                      &left_ids, &right_ids);
 
	/* if for any reason the split did not work (maybe all hr is singular point)
	   make a leaf */
	if(dyv_array_size(left_rows) == 0 || dyv_array_size(right_rows) == 0) {
	  int i;
	  kd->sum = mk_sum_of_dyv_array(points);
	  kd->sum_mags = sum_mags_of_dyv_array(points);
	  kd->points = mk_dyv_array_of_zeroed_dyvs(kd->num_points, num_dims);
	  for(i=0; i<kd->num_points; i++) {
		const dyv *x_i = dyv_array_ref(points, i);
		dyv_array_set(kd->points, i, x_i);
	  }
      kd->point_ids = mk_copy_ivec(point_ids);
	  free_dyv_array(left_rows);
      free_ivec(left_ids);
	  free_dyv_array(right_rows);
      free_ivec(right_ids);
	  return kd;
	}
	kd->left = mk_simple_kdtree_basic(left_rows, left_ids, universe, min_box_width_rel, max_leaf_size);
    free_dyv_array(left_rows);
	left_rows = NULL;
    free_ivec(left_ids);
    left_ids = NULL;

	kd->right = mk_simple_kdtree_basic(right_rows, right_ids, universe, min_box_width_rel, max_leaf_size);
    free_dyv_array(right_rows);
	right_rows = NULL;
    free_ivec(right_ids);
    right_ids = NULL;

    kd->sum = (kd->left==NULL) ? mk_copy_dyv(kd->right->sum) :
	                (kd->right==NULL) ? mk_copy_dyv(kd->left->sum) :
                    mk_dyv_plus(kd->left->sum,kd->right->sum);
    kd->sum_mags = (kd->left==NULL) ? kd->right->sum_mags :
              (kd->right==NULL) ? kd->left->sum_mags :
              kd->left->sum_mags + kd->right->sum_mags;
  }

  return kd;
}

simple_kdtree *mk_simple_kdtree(const dyv_array *x, double min_box_width_rel, int max_leaf_size)
{
  simple_kdtree *kd;
  const int num_dims = dyv_size(dyv_array_ref(x, 0));
  const int num_rows = dyv_array_size(x);
  dyv *universe = mk_dyv(num_dims);
  ivec *identity = mk_identity_ivec(num_rows);

  /* see what is the width of each dim, make up universe struct */
  min_box_width_rel = fabs(min_box_width_rel);
  {
	hrect *hr = mk_hrect_bounding_dyv_array(x);
	int j;
	for(j=0; j<num_dims; j++)
	  dyv_set(universe, j, hrect_width_ref(hr, j));
	free_hrect(hr);
  }

  DBG_PRINTF(D_TRACE, ("Building simple_kdtree from %d records and %d attributes...\n", 
					   num_rows, num_dims));

  kd = mk_simple_kdtree_basic(x, identity, universe, min_box_width_rel, max_leaf_size);
  DBG_PRINTF(D_TRACE, ("...finished making tree\n"));
  free_dyv(universe);
  free_ivec(identity);
  return kd;
}

int kd_num_nodes(simple_kdtree *kd)
{
  int result;
  if ( kd == NULL )
    result = 0;
  else
    result = 1 + kd_num_nodes(kd->left) + kd_num_nodes(kd->right);
  return result;
}

simple_kdtree *find_simple_leaf(simple_kdtree *kd,dyv *x)
{
  simple_kdtree *leaf = NULL;
  while ( kd != NULL )
  {
    leaf = kd;
    if ( dyv_ref(x,kd->split_dim) <= kd->split_val )
      kd = kd->left;
    else
      kd = kd->right;
  }
  if ( leaf == NULL ) my_error("find_leaf...NULL kd passed in");
  return leaf;
}

double dsqd_hrect_point(const hrect *hr, const dyv *point, dyv *work)
{
  double result;
  copy_dyv(point, work);
  clip_to_inside_hrect(hr, work);
  result = dyv_dsqd(work, point);
  return result;
}

dyv_array *mk_dyv_array_of_zeroes(int array_size,int dyv_size)
{
  dyv_array *result = mk_empty_dyv_array();
  dyv *zeroes = mk_zero_dyv(dyv_size);
  int j;
  for ( j = 0 ; j < array_size ; j++ )
    add_to_dyv_array(result,zeroes);
  free_dyv(zeroes);
  return result;
}

#ifdef THIS_FUNCTION_ALREADY_DEFINED_IN_DSUT_C
/* Makes an ivec (vector of integers containing
   a subset of the numbers { 0 , 1 , ... n-1 }).
   This subset is of size "size" and has no duplicates.
   Oh yeah, and its random. */
/* WARNING IMPLEMENTED STUPIDLY SO THAT IT COSTS O(n).
   THERE SHOULD BE AN O(size) IMPLEMENTATION IF THIS
   PROGRAMMER WEREN'T A MORON */
ivec *mk_random_subset_of_ints(int n,int size)
{
  ivec *iv = mk_identity_ivec(n);
  shuffle_ivec(iv);
  while ( ivec_size(iv) > size )
    ivec_remove_last_element(iv);
  return iv;
}
#endif

bool is_kd_leaf(const simple_kdtree *node)
{
  return !node->left || !node->right;
}

/* rows may be NULL dneoting "use all rows" */
double sum_mags_of_dym_rows(dym *dm, ivec *rows)
{
  int num_rows = (rows == NULL) ? dym_rows(dm) : ivec_size(rows);
  double ret = 0.0;
  int j;
  for ( j = 0 ; j < num_rows ; j++ )
  {
    int row = (rows==NULL) ? j : ivec_ref(rows,j);
    dyv *this_row = mk_dyv_from_dym_row(dm,row);
	ret += dyv_scalar_product(this_row, this_row);
    free_dyv(this_row);
  }
  return ret;
}

double sum_mags_of_dyv_array(const dyv_array *da)
{
  const int num_rows = dyv_array_size(da);
  double ret = 0.0;
  int j;
  for ( j = 0 ; j < num_rows ; j++ )
  {
    const dyv *this_row = dyv_array_ref(da, j);
	ret += dyv_scalar_product(this_row, this_row);
  }
  return ret;
}

/* Draws the points in x. If more than 6000, draws a sample of 6000 of them */
void draw_simple_x(dym *x)
{
  int max_dots = 6000;
  ivec *rows = mk_identity_ivec(dym_rows(x));
  if ( dym_rows(x) > max_dots )
  {
    ivec *actual_rows = mk_ivec(max_dots);
    int i;
    shuffle_ivec(rows);
    for ( i = 0 ; i < max_dots ; i++ )
      ivec_set(actual_rows,i,ivec_ref(rows,i));
    free_ivec(rows);
    rows = actual_rows;
  }
 
  draw_x_simple_sample(x,rows);
  free_ivec(rows);
}

/* Draws all the rows in x mentioned in "rows" but no others */
void draw_x_simple_sample(dym *x,ivec *rows)
{
  bool dots = ivec_size(rows) < 1000;
  int i;
  for ( i = 0 ; i < ivec_size(rows) ; i++ )
  {
    int row = ivec_ref(rows,i);
    double x0 = dym_ref(x,row,0);
    double x1 = (dym_cols(x)<2) ? x0 : dym_ref(x,row,1);
    if ( dots )
      ag_dot(x0,x1);
    else
      ag_pixel(x0,x1);
  }
}

#ifdef NEVER
void ds_realize_all(datset *ds)
{
  int attnum;
  for ( attnum = 0 ; attnum < ds_num_atts(ds) ; attnum++ )
  {
    if ( ds_attnum_to_symbolicp(ds,attnum) )
    {
      if ( ds_all_symbolic_values_are_numbers(ds,attnum) )
        ds_realize(ds,attnum);
      else
      {
        my_error(mk_printf("I cannot convert attribute %s to numeric\n",
                           ds_attnum_to_name(ds,attnum)));
      }
    }
  }
}
#endif

int kd_num_dims(const simple_kdtree *kd) {
  /* traverse the tree to find a point, return its dimensionality */
  while(!is_kd_leaf(kd))
	kd = kd->left;
  return dyv_size(dyv_array_ref(kd->points, 0));
}

const char *const DEBUG_FLAG_NAMES_SIMPLE_KDTREE[] =
 {"D_SHOWCOUNTERS", "D_TRACE", "D_INTERACTIVE",
  "D_SHOW_START_CENTERS", "D_SHOW_END_CENTERS",
  "D_SHOW_RUNNING_CENTERS", "D_SHOWDATA", "D_DRAWPOINTS",
  "D_DRAW_END_CENTERS", "D_PASSES",
  "D_SHOW_LKLIHOOD", "D_VCTR", "D_KDCACHE", "D_SHOW_BVALUE", "D_EXTRA",
  "D_PRINTCLUSTERS",
  "D_DELETE_CTRS",
  ""};
