/*
   File:        kdtr.c
   Author:      Andrew W. Moore
   Created:     Sat Sep 12 16:05:16 EDT 1992
   Description: Basic Kdtrees

   Copyright (C) 1992, Andrew W. Moore
*/

#include <stdio.h>
#include <math.h>
#include "ambs.h"      /* Very basic operations */
#include "amma.h"      /* Fast, non-fragmenting, memory management */
#include "amar.h"      /* Obvious operations on 1-d arrays */
#include "maxdim.h"    /* The MAX_DIM declaration */
#include "gpro.h"      /* Graphics projections kd->2d space */
#include "hype.h"      /* Hyper-rectangles from ../kdtr */
#include "kdtr.h"      /* kd-trees from ../kdtr */

kdnode *make_leaf_node(parent)
kdnode *parent;
{
  kdnode *res = AM_MALLOC(kdnode);
  res -> split = -1;
  res -> pivot = 777.777;
  res -> left = NULL;
  res -> right = NULL;
  res -> parent = parent;
  res -> data = NULL;
  return(res);
}

void init_kdtree(kd,dim,middle,scales)
kdtree *kd;
int dim;
float *middle,*scales;
/*
   middle and scales are not copied, so don't alter them while kd lives
*/
{
  kd -> dim = dim;
  kd -> middle = middle;
  kd -> scales = scales;
  kd -> root = make_leaf_node((kdnode *) NULL);
}

bool is_leaf(kdn)
kdnode *kdn;
{
  return(kdn->split < 0);
}

kdnode *child(kdn,heading)
kdnode *kdn;
int heading;
/*
  Heading must be LO or HI (declared in hype.h)
  kdn must be a non-null, non-leaf.
*/
{
  if ( kdn == NULL || is_leaf(kdn) || (heading != LO && heading != HI) )
    my_error("kdtr.c, child() pissed");
  else
    return( (heading==LO) ? kdn->left : kdn->right );
}

void free_kdnode_and_below(kdn)
kdnode *kdn;
{
  if ( !is_leaf(kdn) )
  {
    free_kdnode_and_below(kdn->left);
    free_kdnode_and_below(kdn->right);
  }
  am_free((char *)kdn,sizeof(kdnode));
}

void free_kdtree_contents(kd)
kdtree *kd;
{
  free_kdnode_and_below(kd->root);
}

void free_descendants(kdn)
kdnode *kdn;
/* Frees everything below (but not including) kdn, and
   turns kdn into a leaf node.
*/
{
  free_kdnode_and_below(kdn->left);
  kdn -> left = NULL;
  free_kdnode_and_below(kdn->right);
  kdn -> right = NULL;
  kdn -> split = -1;
  kdn -> pivot = -77.77;
}

kdnode *leaf_search(kd,floats)
kdtree *kd;
float *floats;
{
  kdnode *kdn = kd -> root;
  while ( !is_leaf(kdn) )
    kdn = (floats[kdn->split] <= kdn->pivot) ? kdn->left : kdn->right;
  return(kdn);
}

kdnode *leaf_search_to_level(kd,floats,level)
kdtree *kd;
float *floats;
int level;
/* 
   level = 0 => get root node
           1 => get appropriate root child etc.
       
   if level too high, returns the greatest-depth node containing floats
*/
{
  kdnode *kdn = kd -> root;
  while ( !is_leaf(kdn) && level > 0 )
  {
    kdn = (floats[kdn->split] <= kdn->pivot) ? kdn->left : kdn->right;
    level--;
  }
  return(kdn);
}

void split_node(kdn,split,pivot)
kdnode *kdn;
int split;
float pivot;
{
  if ( !is_leaf(kdn) )
    my_error("split_node() : can't split leaf");
  else
  {
    kdn -> split = split;
    kdn -> pivot = pivot;
    kdn -> left = make_leaf_node(kdn);
    kdn -> right = make_leaf_node(kdn);
  }
}

void fpkd2(s,kdn,depth)
FILE *s;
kdnode *kdn;
int depth;
{
  if ( !is_leaf(kdn) )
  {
    int i;
    fpkd2(s,kdn->left,depth+1);
    for ( i = 0 ; i < depth ; i++ )
      fprintf(s,".");
    fprintf(s,"Split %d, Pivot %g\n",kdn->split,kdn->pivot);
    fpkd2(s,kdn->right,depth+1);
  }
}

void fprint_kdtree(s,kd)
FILE *s;
kdtree *kd;
{
  if ( is_leaf(kd->root) )
    fprintf(s,"One big leaf\n");
  else
    fpkd2(s,kd->root,0);
}

void drkd2(gp,kdn,xlo,ylo,xhi,yhi)
gproj *gp;
kdnode *kdn;
float xlo,ylo,xhi,yhi;
{
  if ( is_leaf(kdn) )
    return;
  else if ( kdn->split == gp->x.comp )
  {
    drkd2(gp,kdn->left,xlo,ylo,kdn->pivot,yhi);
    gp_draw_line(gp,kdn->pivot,ylo,kdn->pivot,yhi);
    drkd2(gp,kdn->right,kdn->pivot,ylo,xhi,yhi);
  }
  else if ( kdn->split == gp->y.comp )
  {
    drkd2(gp,kdn->left,xlo,ylo,xhi,kdn->pivot);
    gp_draw_line(gp,xlo,kdn->pivot,xhi,kdn->pivot);
    drkd2(gp,kdn->right,xlo,kdn->pivot,xhi,yhi);
  }
  else if ( gp->base[kdn->split] <= kdn->pivot )
    drkd2(gp,kdn->left,xlo,ylo,xhi,yhi);
  else
    drkd2(gp,kdn->right,xlo,ylo,xhi,yhi);
}

void gp_draw_kdtree(gp,kd)
gproj *gp;
kdtree *kd;
{
  drkd2(gp,kd->root,gp->x.lo,gp->y.lo,gp->x.hi,gp->y.hi);
}

void gproj_from_kdtree(kd,gp)
kdtree *kd;
gproj *gp;
/*
   Assumes component 0 of tree will be x-axis
   Assumes component 1 of tree will be y-axis
*/
{
  gproj_from_m_and_s(gp,kd->middle,kd->scales,kd->dim,0,1);
}

/* KDNODE_LIST operations */

kdnode_list *add_to_kdnode_list(kdn,ks)
kdnode *kdn;
kdnode_list *ks;
{
  kdnode_list *res = AM_MALLOC(kdnode_list);
  res -> kdnode = kdn;
  res -> next = ks;
  return(res);
}

void free_kdnode_list(ks)
kdnode_list *ks;
{
  while ( ks != NULL )
  {
    kdnode_list *next = ks -> next;
    am_free((char *)ks,sizeof(kdnode_list));
    ks = next;
  }
}

/* DIR data structure. For use by adjacent node search */

typedef struct dir_struct
{
  int comp;
  int heading; /* Must be LO or HI */
} dir;

typedef struct dirset_struct
{
  int dim;
  int set_size;
  bool lo_used[MAX_DIM];
  bool hi_used[MAX_DIM];
} dirset;

void direction_to_sibling(kdn,d)
kdnode *kdn;
dir *d;
{
  kdnode *parent = kdn->parent;
  if ( parent == NULL )
    my_error("dir2sib");
  else
  {
    bool am_left_child = EQ_PTR(kdn,parent->left);
    bool am_right_child = EQ_PTR(kdn,parent->right);

    d->comp = parent -> split;

    if ( am_left_child )
      d->heading = HI;
    else if ( am_right_child )
      d->heading = LO;
    else
      my_error("dir2sub oh bugger");
  }
}

void opposite_dir(d,opp_dir)
dir *d;
dir *opp_dir;
{
  opp_dir -> comp = d -> comp;
  opp_dir -> heading = (d->heading == LO) ? HI : LO;
}

void fill_dirset(dset,kddim)
dirset *dset;
int kddim;
{
  dset -> dim = kddim;
  dset -> set_size = 2 * kddim;
  set_bools_constant(dset->lo_used,kddim,TRUE);
  set_bools_constant(dset->hi_used,kddim,TRUE);
}

bool dirset_empty(dset)
dirset *dset;
{
  return(dset->set_size == 0);
}

bool in_dirset(d,dset)
dir *d;
dirset *dset;
{
  bool *used = (d->heading == LO) ? dset->lo_used : dset->hi_used;
  return(used[d->comp]);
}

void remove_from_dirset(dset,d)
dirset *dset;
dir *d;
{
  bool *used = (d->heading == LO) ? dset->lo_used : dset->hi_used;
  if ( used[d->comp] )
  {
    used[d->comp] = FALSE;
    dset -> set_size -= 1;
  }
}

/* Finding adjancent nodes is a slightly tricky algorithm. See the author's
   little red book, or else check out the following comments.

   Remember what hyper-rectangles do. They can have infinite sides as well
   as fixed limits if we want them to.

   Remember what a dir is. It's a direction pointing in kd-space. It is
   axis-aligned, which means all we need to know about it is which
   axis of kd-space it's aligned with, and whether it's pointing LO or
   HI direction.

   Well, the first function we define finds some extreme leaf nodes subject
   to the above two structures.
*/

kdnode_list *fexts(kdn,hy,d,lsofar)
kdnode *kdn;
hype *hy;
dir *d;
kdnode_list *lsofar;
/*
   Finds all leaf nodes at or below kdn which also satisfy the following
   constraints, adds them on to lsofar, and returns the result.

   All added nodes must have a non-zero intersection with hyperrectangle hy.

   All added nodes must be extreme in direction "d". IE, the path from
   kdn to a successful leaf node may not include a turn in the opposite
   direction to direction d
*/
{
  if ( is_leaf(kdn) )
    lsofar = add_to_kdnode_list(kdn,lsofar);
  else if ( kdn -> split == d -> comp )
    lsofar = fexts(child(kdn,d->heading),hy,d,lsofar);
  else
  {
    if ( hype_low_limit_is_below(hy,kdn->split,kdn->pivot) )
      lsofar = fexts(kdn->left,hy,d,lsofar);
    if ( hype_high_limit_is_above(hy,kdn->split,kdn->pivot) )
      lsofar = fexts(kdn->right,hy,d,lsofar);
  }

  return(lsofar);
}

kdnode_list *fadjs(kdn,hy,dset,lsofar)
kdnode *kdn;
hype *hy;
dirset *dset;
kdnode_list *lsofar;
/*
   Finds all leaf nodes adjacent to kdn which also satisfy the following
   constraints, adds them on to lsofar, and returns the result.

   All added nodes must have a non-zero intersection with hyperrectangle hy.

   Any adjacent node has, of course, exactly one plane with which they touch
   kdn. The direction from kdn to the neighbor, travelling perpendicular to
   this plane must be in dset.
*/
{
  kdnode *parent = kdn -> parent;
  if ( parent != NULL )
  {
    dir d;
    direction_to_sibling(kdn,&d);
    if ( in_dirset(&d,dset) )
    {
      dir opp_dir;
      opposite_dir(&d,&opp_dir);
      lsofar = fexts(child(parent,d.heading),hy,&opp_dir,lsofar);
      remove_from_dirset(dset,&d);
      hype_update(hy,d.comp,d.heading,parent->pivot);
    }
    if ( !dirset_empty(dset) )
      lsofar = fadjs(parent,hy,dset,lsofar);
  }
  return(lsofar);
}

kdnode_list *adjacent_leaves(kdn,kddim)
kdnode *kdn;
{
  dirset dset;
  hype *hy = create_hype(kddim);
  kdnode_list *result;

  fill_dirset(&dset,kddim);
  result = fadjs(kdn,hy,&dset,(kdnode_list *) NULL);

  free_hype(hy);
  return(result);
}

void hype_of_kdnode(kdn,hy)
kdnode *kdn;
hype *hy;
/*
   PRE: hype size must match kdtree dim
*/
{
  dirset dset;
  set_hype_infinite(hy);
  fill_dirset(&dset,hy->size);
  
  while ( !dirset_empty(&dset) && kdn -> parent != NULL )
  {
    dir d;
    direction_to_sibling(kdn,&d);
    if ( in_dirset(&d,&dset) )
    {
      hype_update(hy,d.comp,d.heading,kdn->parent->pivot);
      remove_from_dirset(&dset,&d);
    }
    kdn = kdn -> parent;
  }
}

int kdnode_depth(kdn)
kdnode *kdn;
{
  int result = 0;
  for ( ; kdn ->parent != NULL ; kdn = kdn -> parent )
    result++;
  return(result);
}

int number_nodes_below(kdn)
kdnode *kdn;
{
  if ( is_leaf(kdn) )
    return(1);
  else
    return(number_nodes_below(kdn->left) + number_nodes_below(kdn->right));
}

int kdtree_size(kd)
kdtree *kd;
{
  return(number_nodes_below(kd->root));
}

#define MAX_DEPTH 100

void find_pdh(kdn,depth,count)
kdnode *kdn;
int depth;
int *count;
{
  if ( depth >= MAX_DEPTH ) my_error("flnfbgstrqdw");
  if ( is_leaf(kdn) )
    count[depth] += 1;
  else
  {
    find_pdh(kdn->left,depth+1,count);
    find_pdh(kdn->right,depth+1,count);
  }
}

int partition_depth_histogram(kd)
kdtree *kd;
{
  int count[MAX_DEPTH];
  int i;
  set_ints_constant(count,MAX_DEPTH,0);
  find_pdh(kd->root,0,count);
  
  for ( i = 0 ; i < MAX_DEPTH ; i++ )
    if ( count[i] != 0 )
      printf("D%d=%d ",i,count[i]);
  printf("\n");
}

int max_depth_below(kdn)
kdnode *kdn;
{
  int result;

  if ( is_leaf(kdn) )
    result = 0;
  else
  {
    int d_left = max_depth_below(kdn->left);
    int d_right = max_depth_below(kdn->right);
    result = 1 + int_max(d_left,d_right);
  }

  return(result);
}

int kdtree_depth(kd)
kdtree *kd;
{
  return(max_depth_below(kd->root));
}