/*                          Matlab XMeans                                 


 This code was written and is copyright by Rudolph van der Merwe,
 for use in the ReBEL toolkit ( http://choosh.ece.ogi.edu/rebel/ ).

 It was kindly donated by him to the x-means project.

 Matlab MEX wrapper for the XMeans algorithm developed by Dan Pelleg    
 and Andrew Moore at CMU's Auton group. XMeans is a fast implementation 
 of the classic K-means clustering algorithm utilizing a MRKD-tree data 
 structure. It also does BIC based model selection to automatically     
 determine the underlying number of cluster centroids that best models  
 the data.                                                              
                                                                        
 The Matlab interface to this function is:                              
                                                                        
 [CPos, Ndat, BIC] = xmeans(X, k, max_leaf_size, min_box_width, cutoff_factor,    
                            max_iter, num_splits, max_ctrs)                       
                                                                        
 INPUTS

    X             : xdim-by-N Matlab matrix of datapoints. Each column
                    is a datapoint.
    k             : Number of centers (centroids) to start with. [1]
    max_leaf_size : MRKD-tree construction parameter             [40]
    min_box_width : MRKD-tree construction parameter             [0.05]
    cutoff_factor : MRKD-tree construction parameter             [0.5]
    max_iter      : Maximum number of iterations                 [200]
    num_splits    : Number of center splitting attempts          [6]
    max_ctrs      : Maximum number of centers                    []

 OUTPUTS

    CPos          : xdim-by-M Matlab matrix of M best center (centroid) 
                    coordinates.
    Ndat          : 1-by-M Matlab vector of counts of data points belonging to 
                    each of the M centers. Corresponds with CPos.
    BIC           : Final BIC score
    
*/


#include <assert.h>
#include <string.h>
#include <stdlib.h>

#include "simple_kdtree.h"
#include "kmeans.h"
#include "amgr.h"
#include "debug.h"
#include "incremental.h"
#include "traversal.h"
#include "split.h"
#include <math.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <unistd.h>


#include "mex.h"

#define ASSERT(c,s) { if(!(c)) { mexErrMsgTxt(s); return; }}
#define ASSERTF(c,s) { if(!(c)) { mexErrMsgTxt(s); return 0; }}


static int mx2dyv_array(const mxArray *m, dyv_array *x)
{
  int M,N,i;
  double *p;
  
  /* We assume that the N M-dimensional data points are stored 
     in a M-by-N Matlab matrix. */
 
  M = mxGetM(m);   /* get number of rows ... data dimension*/
  N = mxGetN(m);   /* get number of columns ... number of data points */

  p = mxGetPr(m);      /* Get pointer to data storage area */
  
  for(i=0; i < N; i++){   
    add_to_dyv_array(x, (const dyv *) mk_dyv_from_farr(p, M));  
    p += M;  
  }
 
  return 1;
}


int count_points_in_kdtree(const simple_kdtree *kd)
{
  if(is_kd_leaf(kd)) {
	return kd->num_points;
  } else {
	return count_points_in_kdtree(kd->left) + count_points_in_kdtree(kd->right);
  }
}

void nodestats_count_bulk(const simple_kdtree *node, int closest_ctr, traversal_stats *stats)
{
  vctr *centers = stats->centers;
  const int pidx = vctr_pidx(centers, closest_ctr);
  int k = count_points_in_kdtree(node);
  ivec_increment(stats->stat.nums, pidx, k);
}

void nodestats_count_point(int point_id, const dyv *point,
                           int closest_ctr, traversal_stats *stats)
{
  vctr *centers = stats->centers;
  const int pidx = vctr_pidx(centers, closest_ctr);
  ivec_increment(stats->stat.nums, pidx, 1);
}

static int genMatlabOutput(mxArray **m0, mxArray **m1, mxArray **m2, simple_kdtree *kd, vctr *centers, 
                           const kmeans_parm *orig_parm, double score, int xdim, int number_of_centers)
{

  double *p0; 
  double *p1;
  double *p2; 

  int i,j;

  kmeans_parm noncaching_parm;
  ivec *all_centers = mk_live_ctrs(centers);

  kmeans_state kmstate;
  traversal_stats stats;

  *m0 = mxCreateDoubleMatrix(xdim, number_of_centers, mxREAL);
  *m1 = mxCreateDoubleMatrix(1, number_of_centers, mxREAL);
  *m2 = mxCreateDoubleMatrix(1, 1, mxREAL);

  p0 = mxGetPr(*m0);
  p1 = mxGetPr(*m1);
  p2 = mxGetPr(*m2);

  *p2 = score;   /* Assign final score */

  /* make a parm version that never uses cache */
  memcpy(&noncaching_parm, orig_parm, sizeof(noncaching_parm));
  noncaching_parm.usecache = FALSE;

  kmstate.kdtree = kd;
  kmstate.current_centers = centers;
  kmstate.b_cutoff = 0.0;		/* unused */
  kmstate.phase = PHS_UNPHASED;
  kmstate.phase1_ctrs = 0;		/* unused */

  stats.update_point = nodestats_count_point;
  stats.update_bulk = nodestats_count_bulk;
  stats.prune = distortion_prune;

  /* those value are probably not needed for the print funcs */
  stats.mk_stat_snapshot = nodestats_mk_snapshot;
  stats.mk_cache_entry = nodestats_mk_cache_entry;
  stats.apply_cache = nodestats_apply_cache;
  stats.free_cached_stats = nodestats_free_cache;
  stats.centers = centers;
  /* no need to keep storage for stats */
  stats.stat.distortion = NULL;
  stats.stat.nums = NULL;
  stats.stat.sums = NULL;
  stats.aux = NULL;

  stats.stat = nodestats_mk_stats(1+vctr_max_pidx(centers), xdim);

  cached_bl_traversal(&kmstate, all_centers, &stats, &noncaching_parm);

  /* Copy all current live centers to Matlab matrix */
  for(j=get_first_live_vctr(centers), i=0; is_valid_vctr(centers, j); j=get_next_live_vctr(centers, j), i++) {
    copy_dyv_to_farr(vctr_ref(centers, j), p0);
    p0 += xdim;
    *p1++ = (double) ivec_ref(stats.stat.nums, vctr_pidx(centers, j));
  }

  free_ivec(all_centers);

  return 1;
  
}


/* Matlab gateway function */

void mexFunction(int nlhs,      mxArray *plhs[],
                 int nrhs,const mxArray *prhs[]) 
{

  int xdim, number_of_centers;

  double k, k_step, min_box_width, final_score;
  int max_leaf_size, num_splits, del_steps_ratio;
  kmeans_parm parm, caching_parm;
  
  const char *init_ctrs = "";

  const char *method_name = "blacklist";

  simple_kdtree *kd =0;

  ASSERT(nlhs == 3, "Number of LHS arguments is not 3.");
  ASSERT(nrhs == 8, "Number of RHS arguments is not 8.");

  xdim = mxGetM(prhs[0]);  /* Dimension of input data */
 
  k = *mxGetPr(prhs[1]);
  max_leaf_size = (int) mxGetScalar(prhs[2]);
  min_box_width = *mxGetPr(prhs[3]);
  num_splits = (int) mxGetScalar(prhs[6]);
  parm.cutoff_factor = *mxGetPr(prhs[4]);
  parm.max_iter = (int) mxGetScalar(prhs[5]);
  parm.max_ctrs = (int) mxGetScalar(prhs[7]);
  parm.mtd = MTD_BLACKLIST;
  parm.max_secs = (int) parm.max_iter * 30;
  am_srand(1);   /* stir the pot a bit... */

  parm.loop_over_k = FALSE;
  parm.usecache = CACHE_RDWR;
  parm.delete_ctrs = FALSE;
  if(!strcasecmp(method_name, "slow")) {
	parm.mtd = MTD_SLOW;
  } else if(!strcasecmp(method_name, "hplane"))
	parm.mtd = MTD_HPLANE;
  else if(!strcasecmp(method_name, "blacklist")) {
	parm.loop_over_k = FALSE;
	if(k==0) {
	  if(num_splits <= 0)
		ASSERT(0,"Must supply num_splits");
	  k = 1.0*(int)parm.max_ctrs/num_splits;
	  if(k<1.0)
		k=1.0;
	}
	parm.mtd = MTD_BLACKLIST;
  } else if(!strcasecmp(method_name, "BL-delete")) {
	parm.loop_over_k = FALSE;
	parm.delete_ctrs = TRUE;
	if(k==0) {
	  if(num_splits <= 0)
		ASSERT(0,"Must supply num_splits");
	  k = 1.0*(int)parm.max_ctrs/num_splits;
	  if(k<1.0)
		k=1.0;
	}
    if(del_steps_ratio <= 0)
      ASSERT(0,"del_steps_ratio must be positive");
	parm.mtd = MTD_BLACKLIST;
  } else if(!strcasecmp(method_name, "BL-search")) {
	parm.loop_over_k = TRUE;
	parm.mtd = MTD_BLACKLIST;
  }
  else if(!strcasecmp(method_name, "BL-RS"))
	parm.mtd = MTD_BLACKLIST_RANDSTART;
  else if(!strcasecmp(method_name, "BL-aprx"))
	parm.mtd = MTD_BLACKLIST_APPROX;
  else if(!strcasecmp(method_name, "BL-nocache")) {
	parm.usecache = CACHE_NONE;
	parm.mtd = MTD_BLACKLIST_NOCACHE;
  }
  else
	ASSERT(0,"Don't know anything about the indicated method. Legal values are: slow, hplane, blacklist, BL-RS, BL-aprx, BL-nocache, BL-search");

  if(parm.loop_over_k) {
	if(num_splits <= 0) {
	  num_splits = 1;
	  k_step = 0;
	} else {
	  if(k==0) {
		k_step = 1.0*(int)parm.max_ctrs/(num_splits+1); /* how much to increase k each time */
		k = 1.0*(int)parm.max_ctrs/num_splits;
	  } else
		k_step = ((int)parm.max_ctrs - k)/num_splits; /* how much to increase k each time */
	}
	if(k_step < 1.0)
	  k_step = 1.0;
  }

  if(k<=0)
	ASSERT(0,"Must supply k");

  /* make a parm version that always uses cache */
  memcpy(&caching_parm, &parm, sizeof(caching_parm));
  caching_parm.usecache = CACHE_RDWR;

  {  /* Foo block begin */
    /* convert Matlab data structures to Auton data structures */
    dyv_array *x = mk_empty_dyv_array();
    if(!mx2dyv_array(prhs[0], x)) return;  /* First convert dataset */
    /* build simple kd-tree from data */
    kd = mk_simple_kdtree((const dyv_array *) x, min_box_width, max_leaf_size);
      
    /* repeat running kmeans, splitting centers until done */
    {  /* Foo block begin */
      vctr *ctrs = 0;
      vctr *best_ctrs = 0;		/* best estimate for ctrs so far */
      double best_score = -HUGE_VAL;
      const int ndims = kd_num_dims(kd);
      int splitnum = 0;           /* counter for how many splits we've done already */
      bool maybe_ncenters_didnt_change; /* used to mark iteration where #centers may not change
                                           and that's OK with us*/

      if(strcmp(init_ctrs, "")) { /* initialize ctrs from file */
        ctrs = mk_vctr_from_filename(init_ctrs);
      }

      do {
        int num_old_ctrs;
        double score;
        maybe_ncenters_didnt_change = FALSE;

        /* do a single grand-iteration */
        ctrs = test_kmeans(kd, (int) (k + 0.5), &parm, ctrs, ndims);

        num_old_ctrs = vctr_live_size(ctrs);

        /* get the score of the current estimate */
        score = bic_score(kd, ctrs, &caching_parm, NULL);	/* always use cache for this */

        /* update best-scoring estimate */
        if(score > best_score) {
          if(best_ctrs)
            free_vctr(best_ctrs);
          best_ctrs = mk_copy_vctr(ctrs);
          best_score = score;
          final_score = best_score;
        }

        /* quit if landed straight on max #ctrs */
        if(num_old_ctrs == (int)parm.max_ctrs)
          break;

        /* If we're just looping over k, increase k.
           Otherwise, do a splitting/deletion step */
        if(parm.loop_over_k) {
          k += k_step;
          free_vctr(ctrs);		/* discard current centers */
          ctrs = 0;
        } else {
          if(num_splits - splitnum >= 1) {		/* no point splitting in last iteration */
            /* we either split or delete here, depending on the number of the
               iteration we're in */
            /* we don't want the very first step to be a deletion, so we're adding
               1 to splitnum */
            if(parm.delete_ctrs && ((splitnum + 1) % del_steps_ratio == 0)) { /* it's a deletion step */
              maybe_ncenters_didnt_change = TRUE;
              delete_ctr(kd, ctrs, &parm, maybe_ncenters_didnt_change);
            } else {              /* it's a splitting step */
              split_centers(kd, ctrs, &parm);
            }

            /* quit if #centers doesn't change */
            if(!maybe_ncenters_didnt_change && vctr_live_size(ctrs) == num_old_ctrs)
              break;
          }
        }
      } while( ++splitnum < num_splits);

      /* Assign best center coordinates and number of member points to Matlab output matrix */
      number_of_centers = vctr_live_size(best_ctrs); /* number of live centers */
      genMatlabOutput(&plhs[0], &plhs[1], &plhs[2], kd, best_ctrs, &parm, final_score, xdim, number_of_centers);  

      free_vctr(best_ctrs);
      if(ctrs)
        free_vctr(ctrs);
    
    } /* Foo block end */
  } /* Foo block end */

  if(kd) {
    free_cache_entries(kd);
    free_simple_kdtree(kd);
  }
    
}  

