
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <float.h>

#include <rjmcmc/regression.h>

#include <rjmcmc/engine.h>
#include <rjmcmc/resultset1d.h>
#include <rjmcmc/part1d_regression_rj.h>
#include <rjmcmc/dataset1d.h>
#include <rjmcmc/rjmcmc_util.h>
#include <rjmcmc/rjmcmc_debug.h>

static double part1d_init(void *arg);
static int part1d_select(void *arg);
static void *part1d_perturb(void *arg, int proc);
static double part1d_misfit(void *arg, void *state);
static int part1d_accept(void *arg, 
			   double current,
			   double proposed);
static int part1d_sample(void *arg,
			   int i);

struct part1d {

  resultset1d_t *results;

  part1d_regression_rj_t *current;
  double current_like;

  part1d_regression_rj_t *proposed;
  double proposed_like;

  int nprocesses;

  int out;
  int accepted;
  int process;

  double birth_prob;
  double death_prob;
  double move_prob;
  double value_prob;
  double lambda_prob;

  
  const dataset1d_t *dataset;
  double dy;
  rjmcmc_uniform_rand_t random;
  rjmcmc_normal_rand_t normal;

  int xsamples;
  double *v;

  double dk;

  regression1d_cb_t user_callback;
  void *user_arg;

  double *partitions;
};

static void
do_user_callback(struct part1d *p)
{
  int i;
  
  if (p->user_callback != NULL) {
    for (i = 0; i < part1d_regression_rj_partitions(p->current); i ++) {
      p->partitions[i] = part1d_regression_rj_partition_position(p->current, i);
    }
    
    (p->user_callback)(p->current,
		       p->partitions,
		       part1d_regression_rj_partitions(p->current),
		       (regression1d_value_at_t)part1d_regression_rj_value_at,
		       part1d_regression_rj_lambda(p->current, 0),
		       p->user_arg);
  }
}

resultset1d_t *
part1d_regression(const dataset1d_t *dataset,
		  int burnin,
		  int total,
		  int min_part,
		  int max_part,
		  int max_order,
		  int xsamples,
		  int ysamples,
		  double credible_interval,
		  double pd,
		  rjmcmc_uniform_rand_t random,
		  rjmcmc_normal_rand_t normal,
		  int results,
		  regression1d_cb_t callback,
		  void *user_arg)
{
  struct part1d s;
  rjmcmc_engine_cb_t cb;

  /*
   * Allocate the results and state information
   */
  if (dataset->lambdastd == 0.0) {
    s.nprocesses = 4;
  } else {
    s.nprocesses = 5;
  }

  s.results = resultset1d_create(burnin,
				 total,
				 xsamples,
				 ysamples,
				 s.nprocesses,
				 max_part,
				 max_order,
				 dataset->xmin,
				 dataset->xmax,
				 dataset->ymin,
				 dataset->ymax,
				 credible_interval,
				 results);
  if (s.results == NULL) {
    rjmcmc_error("part1d_regression: failed to create results\n");
    return NULL;
  }

  s.current = part1d_regression_rj_create(min_part,
					  max_part,
					  max_order,
					  1,
					  dataset->xmin,
					  dataset->xmax,
					  pd);
  if (s.current == NULL) {
    rjmcmc_error("part1d_regression: failed to create current state\n");
    return NULL;
  }

  s.proposed = part1d_regression_rj_create(min_part,
					   max_part,
					   max_order,
					   1,
					   dataset->xmin,
					   dataset->xmax,
					   pd);
  if (s.proposed == NULL) {
    rjmcmc_error("part1d_regression: failed to create proposed state\n");
    return NULL;
  }

  s.dataset = dataset;
  s.dy = dataset->ymax - dataset->ymin;
  s.random = random;
  s.normal = normal;
  
  s.xsamples = xsamples;
  s.v = rjmcmc_create_array_1d(xsamples);
  if (s.v == NULL) {
    rjmcmc_error("part1d_regression: failed to create value array\n");
    return NULL;
  }

  s.dk = (double)(max_order + 1);

  s.user_callback = callback;
  s.user_arg = user_arg;
  
  s.partitions = rjmcmc_create_array_1d(max_part);
  
  /*
   * Set the engine callbacks 
   */
  cb.initialize_state = part1d_init;
  cb.select_process = part1d_select;
  cb.perturb_process = part1d_perturb;
  cb.compute_misfit = part1d_misfit;
  cb.accept = part1d_accept;
  cb.sample = part1d_sample;
  cb.arg = (void*)&s;

  if (rjmcmc_engine_run(&cb,
			burnin,
			total,
			1) < 0) {
    return NULL;
  }


  resultset1d_assemble_results(s.results);

  rjmcmc_destroy_array_1d(s.v);
  part1d_regression_rj_destroy(s.current);
  part1d_regression_rj_destroy(s.proposed);

  rjmcmc_destroy_array_1d(s.partitions);

  return s.results;
}

static double part1d_init(void *arg)
{
  struct part1d *s = (struct part1d *)arg;

  if (part1d_regression_rj_initialize(s->current,
				      &(s->dataset),
				      1,
				      s->random,
				      s->normal) < 0) {
    rjmcmc_error("part1d_init: failed to initialize\n");
    return -1.0;
  }

  s->current_like = part1d_regression_rj_misfit(s->current,
						&(s->dataset),
						1);
  return s->current_like;
}

static int part1d_select(void *arg)
{
  struct part1d *s = (struct part1d *)arg;

  return rjmcmc_random_choose_int(0, s->nprocesses - 1, s->random);
}

static void* part1d_perturb(void *arg, int proc)
{
  struct part1d *s = (struct part1d *)arg;

  s->process = proc;

  resultset1d_propose(s->results, s->process);

  switch (proc) {
  case 0: /* Birth */
    s->out = part1d_regression_rj_propose_birth(s->current,
						s->proposed,
						&(s->dataset),
						1,
						s->random,
						s->normal,
						&(s->birth_prob));
    break;

  case 1: /* Death */
    s->out = part1d_regression_rj_propose_death(s->current,
						s->proposed,
						&(s->dataset),
						1,
						s->random,
						s->normal,
						&(s->death_prob));
    break;

  case 2: /* Move */
    s->out = part1d_regression_rj_propose_move(s->current,
					       s->proposed,
					       &(s->dataset),
					       1,
					       s->random,
					       s->normal,
					       &(s->move_prob));
    break;

  case 3: /* Value */
    s->out = part1d_regression_rj_propose_value(s->current,
						s->proposed,
						&(s->dataset),
						1,
						s->random,
						s->normal,
						&(s->value_prob));
    break;

  case 4: /* Sigma */
    s->out = part1d_regression_rj_propose_lambda(s->current,
						 s->proposed,
						 &(s->dataset),
						 1,
						 s->random,
						 s->normal,
						 &(s->lambda_prob));
    break;

  default:
    return NULL;
  }

  return s->proposed;
}

static double part1d_misfit(void *arg, void *state)
{
  struct part1d *s = (struct part1d *)arg;
  
  if (s->out == 0) {
    s->proposed_like = FLT_MAX;
  } else {
    s->proposed_like = part1d_regression_rj_misfit(s->proposed,
						   &(s->dataset),
						   1);
  }

  return s->proposed_like;
}

static int part1d_accept(void *arg, 
			 double current_like,
			 double proposed_like)
{
  struct part1d *s = (struct part1d *)arg;
  double u;
  
  if (s->out == 0) {
    return 0;
  }

  u = log(s->random());

  switch (s->process) {
  case 0: /* Birth */
    /* printf("b %g\n", s->birth_prob); */
    s->accepted = u < (log(s->birth_prob/s->dk) + current_like - proposed_like);
    break;

  case 1: /* Death */
    /* printf("d %g\n", s->death_prob); */
    s->accepted = u < (log(s->death_prob*s->dk) + current_like - proposed_like);
    break;

  case 2: /* Move */
    /* printf("m %g\n", s->move_prob); */
    s->accepted = u < (log(s->move_prob) + current_like - proposed_like);
    break;

  case 3: /* Value */
    /* printf("v %g\n", s->value_prob); */
    s->accepted = u < (log(s->value_prob) + current_like - proposed_like);
    break;

  case 4: /* Lambda */
    s->accepted = u < (log(s->lambda_prob) + current_like - proposed_like);
    break;

  default:
    return 0;
  }

  if (s->accepted) {
    part1d_regression_rj_clone(s->proposed, s->current);
    s->current_like = s->proposed_like;

    resultset1d_accept(s->results, s->process);
  }

  return s->accepted;
}

static int part1d_sample(void *arg,
			 int i)
{
  int j;
  struct part1d *s = (struct part1d *)arg;

  int npart;

  if (part1d_regression_rj_evaluate(s->current,
				    0,
				    s->dataset->xmin,
				    s->dataset->xmax,
				    s->xsamples,
				    s->v) < 0) {
    rjmcmc_error("part1d_sample: failed to evaluate current state\n");
    return -1;
  }

  resultset1d_sample(s->results,
		     i,
		     s->v);

  resultset1d_sample_misfit(s->results,
			    i,
			    s->current_like);

  resultset1d_sample_order(s->results,
			   i,
			   part1d_regression_rj_order(s->current, 0));
  
  resultset1d_sample_lambda(s->results,
			   i,
			   part1d_regression_rj_lambda(s->current, 0));

  npart = part1d_regression_rj_partitions(s->current);
  resultset1d_sample_npartitions(s->results,
				 i,
				 npart);
  
  for (j = 2; j < npart; j ++) {
    resultset1d_sample_partition_x(s->results,
				   i,
				   part1d_regression_rj_partition_position(s->current,
									   j));
  }

  do_user_callback(s);

  return 0;
}

