#include <stdio.h>
#include <stdlib.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_statistics_double.h>
#include <math.h>
#include "tz_error.h"
#include "tz_constant.h"
#include "tz_voxel_linked_list.h"
#include "tz_stack_sampling.h"
#include "tz_voxel_graphics.h"
#include "tz_neurotrace.h"
#include "tz_stack_math.h"
#include "tz_stack_utils.h"
#include "tz_objdetect.h"
#include "tz_voxeltrans.h"
#include "tz_stack_stat.h"
#include "tz_stack_draw.h"
#include "tz_geo3d_vector.h"
#include "tz_stack_bwmorph.h"
#include "tz_stack_bwdist.h"

INIT_EXCEPTION_MAIN(e)

int main(int argc, char* argv[])
{
  Stack *seedimg = Read_Stack("../data/fly_neuron_objseed.tif");
  Object_3d_List *objs = Stack_Find_Object(seedimg, 1, 0);

  Stack *stack2 = NULL;
  Object_3d *obj = NULL;
  stack2 = NULL;

  Object_3d_List *objs2 = NULL;

  while (objs != NULL) {
    stack2 = Read_Stack("../data/fly_neuron_bundle.tif");
    obj = Stack_Grow_Object(stack2, 1, objs->data->voxels[0]);
    if (obj->size > 1000) {
      if (objs2 == NULL) {
	objs2 = Object_3d_List_New();
	objs2->data = obj;
      } else {
	Object_3d_List_Add(&objs2, obj);
	Print_Object_3d_Info(obj);
      }
    } else {
      Kill_Object_3d(obj);
    }
    objs = objs->next;
  }

  Print_Object_3d_List_Compact(objs2);

  Stack *sig = Read_Stack("../data/fly_neuron.tif");
  //Stack *sig = Read_Stack("../data/fly_neuron.tif");
  Stack *stack = Translate_Stack(sig, COLOR, 0);
  //Stack *stack = Read_Stack("../data/blobimg2.tif");

  Stack *traced = Make_Stack(GREY, sig->width, sig->height, sig->depth);
  One_Stack(traced);

  Stack *soma = Read_Stack("../data/fly_neuron_soma.tif");
  Stack_Threshold(soma, 1);
  Stack_Binarize(soma);
  Stack_Not(soma, soma);
  Stack_And(traced, soma, traced);
  Kill_Stack(soma);

  Neurochain *chain = New_Neurochain();
  Neurochain *chain2 = NULL;
  while (objs2 != NULL) {
    if (objs2->data->size > 100) {
      Set_Neuroseg(&(chain->seg), 5, 5, 12, -TZ_PI_2, 0.0);
      Set_Position(chain->position, objs2->data->voxels[0][0],
		   objs2->data->voxels[0][1], objs2->data->voxels[0][2]);

#if 1
      Try {
	Fit_Neuroseg(sig, &(chain->seg), chain->position, TRUE);
      } Catch(e) {
	TZ_ERROR(e);
      }
#endif

#if 1
      Try {
	chain2 = Trace_Neuron(sig, chain, BOTH, traced);
      } Catch(e) {
	TZ_ERROR(e);
      }
#endif
      Neurochain_Erase(traced,  chain2);
      chain = Extend_Neurochain(chain2);
    }
    objs2 = objs2->next;
  }

  Neurochain_Label(stack, Neurochain_Head(chain));
  Write_Stack("../data/fly_neuron_trace.tif", stack);

  return 0;
}


/**
 * \file ModellerFilter.h
 */

#ifndef ATK_MODELLING_MODELLERFILTER_H
#define ATK_MODELLING_MODELLERFILTER_H

#include <ATK/Modelling/config.h>
#include <ATK/Modelling/Types.h>

#include <ATK/Core/TypedBaseFilter.h>

#include <gsl/gsl>

#include <Eigen/Eigen>

#include <tuple>
#include <unordered_set>
#include <vector>

namespace ATK
{
    /// The common class for ModellerFilter
  template<typename DataType_>
  class ATK_MODELLING_EXPORT ModellerFilter: public TypedBaseFilter<DataType_>
  {
  public:
    using Parent = TypedBaseFilter<DataType_>;
    using DataType = DataType_;

  public:
    /**
     * The main ModellerFilter constructor
     * @param nb_dynamic_pins is the number of dymanic pins (that have a voltage that may vary with time)
     * @param nb_input_pins is the number of input pins (that will have varying voltage with time)
     */
    ModellerFilter(gsl::index nb_dynamic_pins, gsl::index nb_input_pins);

    virtual Eigen::Matrix<DataType, Eigen::Dynamic, 1> get_static_state() const = 0;

    /// Returns the number of dynamic pins
    virtual gsl::index get_nb_dynamic_pins() const = 0;

    /// Returns the number of static pins
    virtual gsl::index get_nb_static_pins() const = 0;

    /// Returns the number of input pins
    virtual gsl::index get_nb_input_pins() const = 0;

    /// Returns the number of components
    virtual gsl::index get_nb_components() const = 0;

    /// Returns the name of a dynamic pin, usefull to set output
    virtual std::string get_dynamic_pin_name(gsl::index identifier) const = 0;

    /// Returns the name of a input pin, usefull to set input
    virtual std::string
        get_input_pin_name(gsl::index identifier) const = 0;

    /// Returns the name of a static pin, usefull to set static
    virtual std::string get_static_pin_name(gsl::index identifier) const = 0;

    gsl::index find_dynamic_pin(const std::string& name);
    gsl::index find_input_pin(const std::string& name);
    gsl::index find_static_pin(const std::string& name);

    /// Get number of parameters
    virtual gsl::index get_number_parameters() const = 0;

    /// Get the name of a parameter
    virtual std::string get_parameter_name(gsl::index identifier) const = 0;

    /// Get the value of a parameter
    virtual DataType_ get_parameter(gsl::index identifier) const = 0;

    /// Set the value of a parameter
    virtual void set_parameter(gsl::index identifier, DataType_ value) = 0;
  };
}

#endif


#ifndef DOCH
#define DOCH

#include <gsl/gsl_vector.h>
#include <assert.h>
#include "utils.h"
#include "topic.h"
#include "typedefs.h"

#define OFFSET 0
#define MH_GEM_STDEV 0.05
#define MH_GEM_MEAN_STDEV 0.05
#define MH_GEM_STDEV 0.05

/*
 * resample the levels of a document
 *
 */

void doc_sample_levels(doc* d, short do_permute, short do_remove);

/*
 * update a level count
 *
 */

void doc_update_level(doc* d, int l, double update);

/*
 * read corpus from data
 *
 */

void read_corpus(char* filename, corpus* c, int depth);

/*
 * allocate a new corpus
 *
 */

corpus* corpus_new(double mean, double scale);

/*
 * score the corpus
 *
 */

double gem_score(corpus* corp);

/*
 * GEM MH updates
 *
 */

void corpus_mh_update_gem(corpus* corp);
void corpus_mh_update_gem_mean(corpus* corp);
void corpus_mh_update_gem_scale(corpus* corp);

void compute_log_p_level(doc* d, double gem_mean, double gem_scale);

/*
 * write the document clustering to a file
 *
 */

void write_corpus_assignment(corpus* corp, FILE* file);
void write_corpus_levels(corpus* corp, FILE* file);

// free a corpus

void free_corpus(corpus* corp);
void free_doc(doc* d);

#endif


/*
 * BRAINS
 * (B)LR (R)everberation-mapping (A)nalysis (I)n AGNs with (N)ested (S)ampling
 * Yan-Rong Li, liyanrong@ihep.ac.cn
 * Thu, Aug 4, 2016
 */

/*!
 *  \file dnest_line2d.c
 *  \brief run dnest for 2d line analysis.
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <math.h>
#include <float.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>
#include <mpi.h>

#include "brains.h"

DNestFptrSet *fptrset_line2d;

/*!
 * this function does dnest sampling.
 */
int dnest_line2d(int argc, char **argv)
{
  int i;

  set_blr_model2d();

  num_params_blr = num_params_blr_model + num_params_nlr
                 + num_params_res + num_params_linecenter + 1; /* include line sys err */
  num_params_blr_tot = num_params_blr;
  num_params = parset.n_con_recon + num_params_blr_tot + num_params_var;
  idx_resp = num_params_blr_tot + num_params_drw + num_params_trend;
  idx_difftrend = idx_resp + num_params_resp;
  idx_linecenter = num_params_blr_model + num_params_nlr + num_params_res;

  par_fix = (int *) malloc(num_params * sizeof(int));
  par_fix_val = (double *) malloc(num_params * sizeof(double));
  par_range_model = malloc( num_params * sizeof(double *));
  par_prior_gaussian = malloc( num_params * sizeof(double *));
  for(i=0; i<num_params; i++)
  {
    par_range_model[i] = malloc(2*sizeof(double));
    par_prior_gaussian[i] = malloc(2*sizeof(double));
  }
  par_prior_model = malloc( num_params * sizeof(int));

  fptrset_line2d = dnest_malloc_fptrset();
  /* setup functions used for dnest*/
  fptrset_line2d->from_prior = from_prior_line2d;
  fptrset_line2d->print_particle = print_particle_line2d;
  fptrset_line2d->restart_action = restart_action_2d;
  fptrset_line2d->accept_action = accept_action_2d;
  fptrset_line2d->kill_action = kill_action_2d;
  fptrset_line2d->perturb = perturb_line2d;
  fptrset_line2d->read_particle = read_particle_line2d;

  if(parset.flag_exam_prior != 1)
  {
    fptrset_line2d->log_likelihoods_cal_initial = log_likelihoods_cal_initial_line2d;
    fptrset_line2d->log_likelihoods_cal_restart = log_likelihoods_cal_restart_line2d;
    fptrset_line2d->log_likelihoods_cal = log_likelihoods_cal_line2d;
  }
  else
  {
    fptrset_line2d->log_likelihoods_cal_initial = log_likelihoods_cal_line2d_exam;
    fptrset_line2d->log_likelihoods_cal_restart = log_likelihoods_cal_line2d_exam;
    fptrset_line2d->log_likelihoods_cal = log_likelihoods_cal_line2d_exam;
  }

  set_par_range_model2d();
  /* setup the fixed parameters */
  set_par_fix_blrmodel();

  /* setup the remaining paramters */
  for(i=num_params_blr_model; i<num_params; i++)
  {
    par_fix[i] = 0;
    par_fix_val[i] = -DBL_MAX;
  }

  /* cope with narrow line */
  if(parset.flag_narrowline == 2)
  {
    /* flux */
    if(parset.flux_narrowline_err == 0.0)
    {
      par_fix[num_params_blr_model] = 1.0;
      par_fix_val[num_params_blr_model] = 0.0;
    }
  }
  if(parset.flag_narrowline >= 2)
  {
    /* width */
    if(parset.width_narrowline_err == 0.0)
    {
      par_fix[num_params_blr_model+1] = 1.0;
      par_fix_val[num_params_blr_model+1] = 0.0;
    }
    /* shift */
    if(parset.shift_narrowline_err == 0.0)
    {
      par_fix[num_params_blr_model+2] = 1.0;
      par_fix_val[num_params_blr_model+2] = 0.0;
    }
  }

  /* if flag_fixvar is true, fix continuum variation parameter sigma and tau */
  if(parset.flag_fixvar == 1)
  {
    par_fix[num_params_blr + 1] = 1;
    par_fix_val[num_params_blr + 1] = var_param[1];
    par_fix[num_params_blr + 2] = 1;
    par_fix_val[num_params_blr + 2] = var_param[2];
  }

  /* fix systematic error of line */
  if(parset.flag_line_sys_err != 1)
  {
    par_fix[num_params_blr-1] = 1;
    par_fix_val[num_params_blr-1] = log(1.0);
  }

  /* fix systematic error of continuum */
  if(parset.flag_con_sys_err != 1)
  {
    par_fix[num_params_blr] = 1;
    par_fix_val[num_params_blr] = log(1.0);
  }

  /* fix non-linear response */
  if(parset.flag_nonlinear !=1)
  {
    par_fix[idx_resp + 1] = 1;
    par_fix_val[idx_resp + 1] = 0.0;
  }

  print_par_names_model2d();

  force_update = parset.flag_force_update;
  if(parset.flag_para_name != 1)
    logz_line2d = dnest(argc, argv, fptrset_line2d, num_params, dnest_options_file);

  dnest_free_fptrset(fptrset_line2d);
  return 0;
}

/*!
 * this function setups parameter ranges.
 *
 * The order of parameters is:                              \n
 *   I.   blr model.............()                          \n
 *   II.  narrow line...........(if flag_narrowline is true)\n
 *   III. spectral broadening...()                          \n
 *   IV.  line center...........(if flag_linecenter is true)\n
 *   V.   systematic error......()                          \n
 *   VI.  variability...........()                          \n
 *   VII. long-term trend.......()                          \n
 *   VIII.response A and Ag.....()                          \n
 *   IX.  different trend.......(if flag_difftend is ture)  \n
 *   X.   continuum light curve.()                          \n
 */
void set_par_range_model2d()
{
  int i, j, i1, i2;

  /* setup parameter range, BLR parameters first */
  for(i=0; i<num_params_blr_model; i++)
  {
    par_range_model[i][0] = blr_range_model[i][0];
    par_range_model[i][1] = blr_range_model[i][1];

    par_prior_model[i] = UNIFORM;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 0.0;
  }

  /* cope with narrow line */
  for(i=num_params_blr_model; i<num_params_blr_model + num_params_nlr; i++)
  {
    par_range_model[i][0] = nlr_range_model[i - num_params_blr_model][0];
    par_range_model[i][1] = nlr_range_model[i - num_params_blr_model][1];

    par_prior_model[i] = nlr_prior_model[i - num_params_blr_model];
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 1.0; /* note that for logarithm prior of flux, this value is not used, so does not matter */
  }
  /* cope with spectral broadening */
  for(i=num_params_blr_model + num_params_nlr; i<num_params_blr_model + num_params_nlr + num_params_res; i++)
  {
    par_range_model[i][0] = -10.0;
    par_range_model[i][1] =  10.0;

    par_prior_model[i] = GAUSSIAN;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 1.0;
  }
  /* cope with line center */
  for(i=num_params_blr-num_params_linecenter-1; i< num_params_blr-1; i++)
  {
    par_range_model[i][0] = -10.0;
    par_range_model[i][1] =  10.0;

    par_prior_model[i] = GAUSSIAN;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 1.0;
  }

  /* the last is systematic error */
  i = num_params_blr-1;
  par_range_model[i][0] = sys_err_line_range[0];
  par_range_model[i][1] = sys_err_line_range[1];

  par_prior_model[i] = UNIFORM;
  par_prior_gaussian[i][0] = 0.0;
  par_prior_gaussian[i][1] = 0.0;

  /* variability parameters */
   /* first systematic error */
  i = num_params_blr;
  par_range_model[i][0] = var_range_model[0][0];
  par_range_model[i][1] = var_range_model[0][1];

  par_prior_model[i] = UNIFORM;
  par_prior_gaussian[i][0] = 0.0;
  par_prior_gaussian[i][1] = 0.0;

  for(i=num_params_blr+1; i<num_params_drw + num_params_blr; i++)
  {
    if(var_param_std[i-num_params_blr] > 0.0)
    {
      par_range_model[i][0] = var_param[i-num_params_blr] - 5.0 * var_param_std[i-num_params_blr];
      par_range_model[i][1] = var_param[i-num_params_blr] + 5.0 * var_param_std[i-num_params_blr];

      /* make sure that the range lies within the initial range */
      par_range_model[i][0] = fmax(par_range_model[i][0], var_range_model[i-num_params_blr][0]);
      par_range_model[i][1] = fmin(par_range_model[i][1], var_range_model[i-num_params_blr][1]);

      par_prior_model[i] = GAUSSIAN;
      par_prior_gaussian[i][0] = var_param[i-num_params_blr];
      par_prior_gaussian[i][1] = var_param_std[i-num_params_blr];
    }
    else
    {
      par_range_model[i][0] = var_range_model[i-num_params_blr][0];
      par_range_model[i][1] = var_range_model[i-num_params_blr][1];

      par_prior_model[i] = UNIFORM;
      par_prior_gaussian[i][0] = 0.0;
      par_prior_gaussian[i][1] = 0.0;
    }
  }
  /* long-term trend of continuum */
  for(i=num_params_drw + num_params_blr; i< num_params_drw + num_params_trend + num_params_blr; i++)
  {
    par_range_model[i][0] = var_range_model[3][0];
    par_range_model[i][1] = var_range_model[3][1];

    par_prior_model[i] = GAUSSIAN;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 1.0;
  }

  /* response A and Ag */
  j = 0;
  i1 = idx_resp;
  i2 = idx_resp + num_params_resp;
  for(i=i1; i<i2; i++)
  {
    par_range_model[i][0] = resp_range[j][0];
    par_range_model[i][1] = resp_range[j][1];

    par_prior_model[i] = UNIFORM;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 0.0;

    j++;
  }

  /* different trend in continuum and line */
  j = 0;
  i1 = idx_difftrend;
  i2 = idx_difftrend + num_params_difftrend;
  for(i=i1; i< i2; i++)
  {
    par_range_model[i][0] = var_range_model[4 + j][0];
    par_range_model[i][1] = var_range_model[4 + j][1];

    par_prior_model[i] = UNIFORM;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 0.0;

    j++;
  }

  /* continuum ligth curve values */
  for(i=num_params_blr+num_params_var; i<num_params; i++)
  {
    par_range_model[i][0] = var_range_model[4+num_params_difftrend][0];
    par_range_model[i][1] = var_range_model[4+num_params_difftrend][1];

    par_prior_model[i] = GAUSSIAN;
    par_prior_gaussian[i][0] = 0.0;
    par_prior_gaussian[i][1] = 1.0;
  }

  return;
}

/*!
 *  print names and prior ranges for parameters
 *
 */
void print_par_names_model2d()
{
  if(thistask != roottask)
    return;

  int i, j;
  FILE *fp;
  char fname[BRAINS_MAX_STR_LENGTH], str_fmt[BRAINS_MAX_STR_LENGTH];

  sprintf(fname, "%s/%s", parset.file_dir, "data/para_names_model2d.txt");
  fp = fopen(fname, "w");
  if(fp == NULL)
  {
    fprintf(stderr, "# Error: Cannot open file %s.\n", fname);
    exit(0);
  }

  strcpy(str_fmt, "%4d %-15s %10.6f %10.6f %4d %4d %15.6e\n");

  printf("# Print parameter name in %s\n", fname);

  fprintf(fp, "#*************************************************\n");
  fprint_version(fp);
  fprintf(fp, "#*************************************************\n");

  fprintf(fp, "%4s %-15s %10s %10s %4s %4s %15s\n", "#", "Par", "Min", "Max", "Prior", "Fix", "Val");

  i=-1;
  for(j=0; j<num_params_blr_model; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "BLR model", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }

  for(j=0; j<num_params_nlr; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "narrow line", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }
  for(j=0; j<num_params_res; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "line broaden", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }
  for(j=0; j<num_params_linecenter; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "line center", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }

  i++;
  fprintf(fp, str_fmt, i, "sys_err_line", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);

  i++;
  fprintf(fp, str_fmt, i, "sys_err_con", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  i++;
  fprintf(fp, str_fmt, i, "sigmad", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  i++;
  fprintf(fp, str_fmt, i, "taud", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);

  for(j=0; j<num_params_trend; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "trend", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }

  i++;
  fprintf(fp, str_fmt, i, "A", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);

  i++;
  fprintf(fp, str_fmt, i, "Ag", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);

  for(j=0; j<num_params_difftrend; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "diff trend", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }

  for(j=0; j<parset.n_con_recon; j++)
  {
    i++;
    fprintf(fp, str_fmt, i, "time series", par_range_model[i][0], par_range_model[i][1], par_prior_model[i],
                            par_fix[i], par_fix_val[i]);
  }

  fclose(fp);
}

/*!
 * this function generates a sample from prior.
 */
void from_prior_line2d(void *model)
{
  int i;
  double *pm = (double *)model;

  for(i=0; i<num_params; i++)
  {
    if(par_prior_model[i] == GAUSSIAN)
    {
      pm[i] = dnest_randn()*par_prior_gaussian[i][1] + par_prior_gaussian[i][0];
      dnest_wrap(&pm[i], par_range_model[i][0], par_range_model[i][1]);
    }
    else
    {
      pm[i] = par_range_model[i][0] + dnest_rand() * (par_range_model[i][1] - par_range_model[i][0]);
    }
  }

  /* cope with fixed parameters */
  for(i=0; i<num_params_blr + num_params_var; i++)
  {
    if(par_fix[i] == 1)
      pm[i] = par_fix_val[i];
  }

  which_parameter_update = -1;

  return;
}

/*!
 * this function calculates likelihood at initial step.
 */
double log_likelihoods_cal_initial_line2d(const void *model)
{
  double logL;
  logL = prob_initial_line2d(model);
  return logL;
}

/*!
 * this function calculates likelihood at initial step.
 */
double log_likelihoods_cal_restart_line2d(const void *model)
{
  double logL;
  logL = prob_restart_line2d(model);
  return logL;
}

/*!
 * this function calculates likelihood.
 */
double log_likelihoods_cal_line2d(const void *model)
{
  double logL;
  logL = prob_line2d(model);
  return logL;
}

/*!
 * this function prints out parameters.
 */
void print_particle_line2d(FILE *fp, const void *model)
{
  int i;
  double *pm = (double *)model;

  for(i=0; i<num_params; i++)
  {
    fprintf(fp, "%e ", pm[i] );
  }
  fprintf(fp, "\n");
  return;
}

/*!
 * This function read the particle from the file.
 */
void read_particle_line2d(FILE *fp, void *model)
{
  int j;
  double *psample = (double *)model;

  for(j=0; j < dnest_num_params; j++)
  {
    if(fscanf(fp, "%lf", psample+j) < 1)
    {
      printf("%f\n", *psample);
      fprintf(stderr, "#Error: Cannot read file %s.\n", options.sample_file);
      exit(0);
    }
  }
  return;
}

/*!
 * this function perturbs parameters.
 */
double perturb_line2d(void *model)
{
  double *pm = (double *)model;
  double logH = 0.0, limit1, limit2, width, rnd;
  int which, which_level;

  /*
   * fixed parameters need not to update
   * perturb important parameters more frequently
   */
  do
  {
    rnd = dnest_rand();
    if(rnd < fmax(0.2, 1.0*(num_params_blr+num_params_var)/num_params))
      which = dnest_rand_int(num_params_blr + num_params_var);
    else
      which = dnest_rand_int(parset.n_con_recon) + num_params_blr + num_params_var;
  }while(par_fix[which] == 1);

  which_parameter_update = which;

  /* level-dependent width */
  which_level_update = dnest_get_which_level_update();
  which_level = which_level_update > (size_levels-10)?(size_levels-10):which_level_update;

  if( which_level > 0)
  {
    limit1 = limits[(which_level-1) * num_params *2 + which *2];
    limit2 = limits[(which_level-1) * num_params *2 + which *2 + 1];
    width = limit2 - limit1;
  }
  else
  {
    width = ( par_range_model[which][1] - par_range_model[which][0] );
  }

  width /= (2.35);

  if(par_prior_model[which] == GAUSSIAN)
  {
    logH -= (-0.5*pow((pm[which] - par_prior_gaussian[which][0])/par_prior_gaussian[which][1], 2.0) );
    pm[which] += dnest_randh() * width;
    dnest_wrap(&pm[which], par_range_model[which][0], par_range_model[which][1]);
    logH += (-0.5*pow((pm[which] - par_prior_gaussian[which][0])/par_prior_gaussian[which][1], 2.0) );
  }
  else
  {
    pm[which] += dnest_randh() * width;
    dnest_wrap(&(pm[which]), par_range_model[which][0], par_range_model[which][1]);
  }

  return logH;
}

/*!
 * this function calculates likelihood.
 */
double log_likelihoods_cal_line2d_exam(const void *model)
{
  return 0.0;
}

void accept_action_2d()
{
  int param;
  double *ptemp;

  // the parameter previously updated
  param = which_parameter_update;
  // continuum parameter is updated
  if(param >= num_params_blr)
  {
    /*
     *note that (response) Fline is also changed as long as Fcon is changed.
     *num_params_blr-the parameter is the systematic error of continuum.
     *the change of this parameter also changes continuum reconstruction.
     */

    ptemp = Fcon_rm_particles[which_particle_update];
    Fcon_rm_particles[which_particle_update] = Fcon_rm_particles_perturb[which_particle_update];
    Fcon_rm_particles_perturb[which_particle_update] = ptemp;

    if(force_update != 1)
    {
      ptemp = Fline_at_data_particles[which_particle_update];
      Fline_at_data_particles[which_particle_update] = Fline_at_data_particles_perturb[which_particle_update];
      Fline_at_data_particles_perturb[which_particle_update] = ptemp;
    }
  }
  else if( param < num_params_blr -1 && force_update != 1)
  {
    /* BLR parameter is updated
     * Note a) that the (num_par_blr-1)-th parameter is systematic error of line.
     * when this parameter is updated, Trans2D and Fline are unchanged.
     *      b) Fline is always changed, except for param = num_params_blr-1.
     */

    {
      ptemp = TransTau_particles[which_particle_update];
      TransTau_particles[which_particle_update] = TransTau_particles_perturb[which_particle_update];
      TransTau_particles_perturb[which_particle_update] = ptemp;

      ptemp = Trans2D_at_veldata_particles[which_particle_update];
      Trans2D_at_veldata_particles[which_particle_update] = Trans2D_at_veldata_particles_perturb[which_particle_update];
      Trans2D_at_veldata_particles_perturb[which_particle_update] = ptemp;

      ptemp = Fline_at_data_particles[which_particle_update];
      Fline_at_data_particles[which_particle_update] = Fline_at_data_particles_perturb[which_particle_update];
      Fline_at_data_particles_perturb[which_particle_update] = ptemp;

    }
  }

  return;
}

/*
 * action when particle i is killed in cdnest sampling.
 * particle i_copy's properties is copyed to particle i.
 */
void kill_action_2d(int i, int i_copy)
{
  memcpy(Fcon_rm_particles[i], Fcon_rm_particles[i_copy], parset.n_con_recon * sizeof(double));
  memcpy(Fline_at_data_particles[i], Fline_at_data_particles[i_copy], n_line_data * n_vel_data_ext  * sizeof(double));
  memcpy(TransTau_particles[i], TransTau_particles[i_copy], parset.n_tau*sizeof(double));
  memcpy(Trans2D_at_veldata_particles[i], Trans2D_at_veldata_particles[i_copy], parset.n_tau * n_vel_data_ext * sizeof(double));
  return;
}


/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
#ifdef PADDLE_WITH_MKLML
#include "paddle/fluid/platform/dynload/mklml.h"
#endif

#ifdef PADDLE_USE_OPENBLAS
#include <cblas.h>
// remove typedef in openblas
#undef FLOAT
#undef INT
#undef SIZE
#endif

#include <cmath>
#include <vector>

#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/framework/tensor_util.h"
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/enforce.h"

namespace paddle {
namespace operators {
namespace math {
template <typename DeviceContext, typename T, int Rank>
struct Transpose {
  void operator()(const DeviceContext& context, const framework::Tensor& in,
                  framework::Tensor* out, const std::vector<int>& axis);
};

template <typename DeviceContext, typename T>
struct SetConstant {
  void operator()(const DeviceContext& context, framework::Tensor* tensor,
                  T num);
};

template <typename Place>
void set_constant_with_place(const platform::DeviceContext& context,
                             framework::Tensor* tensor, float value);

void set_constant(const platform::DeviceContext& context,
                  framework::Tensor* tensor, float value);

template <typename DeviceContext, typename T>
struct RowwiseAdd {
  void operator()(const DeviceContext& context, const framework::Tensor& input,
                  const framework::Tensor& vec, framework::Tensor* output);
};

template <typename DeviceContext, typename T>
struct ColwiseSum {
  void operator()(const DeviceContext& context, const framework::Tensor& input,
                  framework::Tensor* vec);
};

template <typename DeviceContext, typename T>
struct RowwiseSum {
  void operator()(const DeviceContext& context, const framework::Tensor& input,
                  framework::Tensor* vec);
};

template <typename DeviceContext, typename T>
struct RowwiseMean {
  void operator()(const DeviceContext& context, const framework::Tensor& input,
                  framework::Tensor* vec);
};

}  // namespace math
}  // namespace operators
}  // namespace paddle


#pragma once

#include "SimpleEvaluator.h"
#include "BooleanDAG.h"
#include "Interface.h"
#include "SuggestionGenerator.h"
#include "OptimizationWrapper.h"
#include "MaxOptimizationWrapper.h"
#include "BoolNodeSimplifier.h"
#include "NumDebugger.h"
#include <map>
#include <set>

#ifndef _NOGSL
#include <gsl/gsl_vector.h>
#else
#include "FakeGSL.h"
#endif

#include "BoolBasedSampler.h"

class SuggestionGeneratorUsingMax: public SuggestionGenerator {
    Interface* interf;
    BooleanDAG* dag;
    map<string, int>& ctrls;
    FloatManager& fm;

    ActualEvaluators* actualEval;
    SmoothEvaluators* smoothEval;

    OptimizationWrapper* opt;
    MaxOptimizationWrapper* maxOpt;

    int minimizeNode;
    set<int> assertConstraints;

    vector<gsl_vector*> maxes;
    gsl_vector* tmp;

    int NUM_MAXES = 5;

    NumDebugger* debugger;
    const vector<vector<int>>& dependentInputs;

public:
    vector<gsl_vector*> dirGrads;
    SuggestionGeneratorUsingMax(BooleanDAG* _dag, FloatManager& _fm, Interface* _interface, map<string, int>& _ctrls, OptimizationWrapper* _opt, MaxOptimizationWrapper* _maxOpt, ActualEvaluators* _actualEval, SmoothEvaluators* _smoothEval, int ncontrols, const vector<vector<int>>& _dependentInputs, NumDebugger* _debugger): dag(_dag), fm(_fm), interf(_interface), ctrls(_ctrls), opt(_opt), maxOpt(_maxOpt), actualEval(_actualEval), smoothEval(_smoothEval), dependentInputs(_dependentInputs), debugger(_debugger) {

        minimizeNode = -1;

	    for (int i = 0; i < dag->size(); i++) { // TODO: this should also be set by caller class
	    	bool_node* n = (*dag)[i];
	    	if (n->type == bool_node::ASSERT && ((ASSERT_node*)n)->isHard()) {
	    		minimizeNode = i;
	    	} else if (n->type == bool_node::ASSERT || Util::isSqrt(n)) {
	    		assertConstraints.insert(i);
	    	} else if (n->type == bool_node::CTRL && n->getOtype() == OutType::BOOL) {
	    		assertConstraints.insert(i);
	    	}
	    }

        for (int i = 0; i < 20; i++) {
            dirGrads.push_back(gsl_vector_alloc(ncontrols));
        }

        for (int i = 0; i < NUM_MAXES; i++) {
            maxes.push_back(gsl_vector_alloc(ncontrols));
        }
        tmp = gsl_vector_alloc(ncontrols);
    }
    ~SuggestionGeneratorUsingMax() {
        for (int i = 0; i < dirGrads.size(); i++) {
            gsl_vector_free(dirGrads[i]);
        }

        for (int i = 0; i < maxes.size(); i++) {
            gsl_vector_free(maxes[i]);
        }
        gsl_vector_free(tmp);
    }

    pair<int, int> getSuggestion(const gsl_vector* state) {
        actualEval->run(state);

        vector<tuple<double, int, int>> s;

        for (auto it = interf->varsMapping.begin(); it != interf->varsMapping.end(); it++) {
            int nodeid = it->first;
            if (interf->hasValue(nodeid)) continue;

            bool_node* n = (*dag)[nodeid];
            bool hasArraccChild = Util::hasArraccChild(n);

            double dist = actualEval->dist(n->id);
            double cost = abs(dist);
            if (hasArraccChild && cost < 0.1) {
                return make_pair(nodeid, dist >= 0.0);
            }

            if (hasArraccChild) {
                cost = cost / 1000.0;
            }

            s.push_back(make_tuple(cost, nodeid, dist > 0));

        }
        cout << endl;

        sort(s.begin(), s.end());
        if (s.size() > 0) {
            return make_pair(get<1>(s[0]), get<2>(s[1]));
        } else {
            Assert(false, "weqhqp");
            return make_pair(-1, 0);
        }

    }

    pair<int, int> getUnsatSuggestion(const gsl_vector* state) {
        actualEval->run(state);
        double dassert;
        double d;
        vector<tuple<int, int>> conflict;
        for (auto it = assertConstraints.begin(); it != assertConstraints.end(); it++) {
            //dassert = actualEval->getErrorOnConstraint(*it);
            if (true) {
                // TODO: this is currently hacky
                const vector<int>& inputs = dependentInputs[*it];
                int minid = -1;
                double minerror  = 1e30;
                for (auto it1 = inputs.begin(); it1 != inputs.end(); it1++) {
                    if (interf->hasValue(*it1)) continue;
                    d = actualEval->dist(*it1);
                    //if (dassert < 0.0) {
                        cout << *it << " " << *it1 << " " << d << " " << Util::print(actualEval->grad(*it1)) << endl;
                    //}
                    if (d < minerror) {
                        minerror = d;
                        minid = *it1;
                    }
                }
                //Assert(minerror > 0.0, "something is wrong");
                if (minid != -1) {
                    cout << "Min " << *it << " " << minid << " " << minerror << endl;
                    conflict.push_back(make_tuple(*it, minid));
                }

                if ((*dag)[*it]->mother()->type == bool_node::NOT &&  minerror >= 0.2) {
                    break;
                }
            }

        }
        sort(conflict.begin(), conflict.end());
        reverse(conflict.begin(), conflict.end());
        Assert(conflict.size() > 0, "No failed asserts???");
        return make_pair(get<1>(conflict[0]), 1);

    }

    SClause* getUnsatClause(const gsl_vector* state, const gsl_vector* initState) {
        actualEval->run(state);
        double dassert;
        double d;
        vector<int> conflictingAsserts;
        for (auto it = assertConstraints.begin(); it != assertConstraints.end(); it++) {
            dassert = actualEval->getErrorOnConstraint(*it);
            if ( (*dag)[*it]->mother()->type != bool_node::NOT) {
                continue;
            }
            if (true) {
                if (dassert < 0.01) {
                    conflictingAsserts.push_back(*it);
                }
                if (dassert < -0.2) {
                    break;
                }
            }

        }

        SClause* clause = new SClause();

        actualEval->run(initState);
        for (auto it = conflictingAsserts.begin(); it != conflictingAsserts.end(); it++) {
                // TODO: this is currently hacky
                const vector<int>& inputs = dependentInputs[*it];
                int minid = -1;
                double minerror  = 1e30;
                for (auto it1 = inputs.begin(); it1 != inputs.end(); it1++) {
                    d = actualEval->dist(*it1);
                    cout << *it << " " << *it1 << " " << d << " " << Util::print(actualEval->grad(*it1)) << endl;
                    if (d < minerror) {
                        minerror = d;
                        minid = *it1;
                    }
                }
                //Assert(minerror > 0.0, "something is wrong");
                if (minid != -1 ) {
                    clause->add(*it, minid);
                }
        }

        return clause;

    }

    SClause* getUnsatClause(const gsl_vector* state) {
        actualEval->run(state);
        double dassert;
        double d;
        SClause* clause = new SClause();
        for (auto it = assertConstraints.begin(); it != assertConstraints.end(); it++) {
            dassert = actualEval->getErrorOnConstraint(*it);
            if (true) {
                // TODO: this is currently hacky
                const vector<int>& inputs = dependentInputs[*it];
                int minid = -1;
                double minerror  = 1e30;
                for (auto it1 = inputs.begin(); it1 != inputs.end(); it1++) {
                    d = actualEval->dist(*it1);
                    if (dassert < 0.0) {
                        cout << *it << " " << *it1 << " " << d << " " << Util::print(actualEval->grad(*it1)) << endl;
                    }
                    if (d < minerror) {
                        minerror = d;
                        minid = *it1;
                    }
                }
                //Assert(minerror > 0.0, "something is wrong");
                if (minid != -1 && (*dag)[*it]->mother()->type == bool_node::NOT &&  minerror >= -0.01) {
                    clause->add(*it, minid);
                }

                if ((*dag)[*it]->mother()->type == bool_node::NOT &&  minerror >= 0.2) {
                    break;
                }
            }

        }

        return clause;

    }

    BooleanDAG* getNewDag(BooleanDAG* origDag, int origNid, const gsl_vector* s) {
        BooleanDAG* newDag = origDag->clone();
        SimpleGradEvaluator* seval = actualEval->getEvaluator(origDag);
        seval->run(s);
        BoolNodeSimplifier simplifier(*newDag, fm);
        simplifier.process(*newDag, origNid, seval);
        DagOptim cse(*newDag, fm);
        cse.process(*newDag);
        newDag->cleanup();
        //newDag->lprint(cout);
        return newDag;
    }

    Predicate* make_pure(Predicate* p, const gsl_vector* s) {
        cout << "Making pure: " << p->print() << " at " << Util::print(s) << endl;
        if (p->isBasic()) {
            BasicPredicate* bp = (BasicPredicate*) p;
            BooleanDAG* newDag = getNewDag(bp->dag, bp->nid, s);
            // register newDag
            actualEval->addEvaluator(newDag);
            smoothEval->addEvaluator(newDag);
            int cid = -1;
            for (int i = newDag->size() - 1; i >= 0; i--) {
                bool_node* n = (*newDag)[i];
                if (n->type == bool_node::ASSERT) {
                    cid = n->mother()->id;
                    break;
                }
            }
            Assert(cid != -1, "No assert in new dag");
            return new BasicPredicate(newDag, cid, bp);
        } else {
            DiffPredicate* dp = (DiffPredicate*)p;
            Predicate* new_p1;
            Predicate* new_p2;
            if (dp->p1->isPure()) {
                new_p1 = dp->p1;
            } else {
                new_p1 = make_pure(dp->p1, s);
            }
            if (dp->p2->isPure()) {
                new_p2 = dp->p2;
            } else {
                new_p2 = make_pure(dp->p2, s);
            }
            int ncontrols = ctrls.size();
            if (ncontrols == 0) {
                ncontrols = 1;
            }
            return new DiffPredicate(new_p1, new_p2, ncontrols);
        }
    }

    int getDirectionsLevel0(const gsl_vector* state, double beta) {
        GradUtil::BETA = beta;
        GradUtil::ALPHA = -beta;
        smoothEval->run(state);

        vector<tuple<double, Predicate*, int>> s;

        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;

            double dist = smoothEval->dist(p);
            double cost = abs(dist);
            if (cost >= 0.1) {
                s.push_back(make_tuple(cost, p, dist > 0));
            }
        }

        sort(s.begin(), s.end());
        cout << "Distances: " << endl;
        int end = 5;
        if (end > s.size()) {
            end = s.size();
        }
        for (int i = 0; i < end; i++) {
            cout << get<1>(s[i])->print() << " " << get<0>(s[i]) << " " << get<2>(s[i]) << endl;
        }

        int numDirs = 0;
        for (int i = 0; i < end; i++) {
            if (numDirs == 5) break;
            double dist = get<0>(s[i]);
            Predicate* p = get<1>(s[i]);
            int val = !get<2>(s[i]);
            if (dist < 0.1) continue;
            smoothEval->getErrorOnPredicate(p, val, dirGrads[numDirs]);
            cout << "G: " << Util::print(dirGrads[numDirs]) << endl;
            bool newDir = true;
            for (int j = 0; j < numDirs; j++) {
                if (Util::sameDir(dirGrads[j], dirGrads[numDirs])) {
                    newDir = false;
                    cout << "Not a new dir" << endl;
                    break;
                }
            }
            if (newDir) {
                numDirs++;
            }
        }
        Assert(numDirs > 0, "No dirs?");
        return numDirs;
    }

    void updateState(gsl_vector* state, int idx, double dx) {
        gsl_vector_set(state, idx, gsl_vector_get(state, idx) + dx);
    }

    void randomize(gsl_vector* tmp) {
        for (int i = 0; i < tmp->size; i++) {
            double r = -1.0 + (rand() % 20)/10.0;
            gsl_vector_set(tmp, i, r);
        }
    }

    int getMaxes(gsl_vector* state, double beta) {
        int numMaxes = 0;
        for (int i = 0; i < NUM_MAXES; i++) {
            randomize(tmp);
            gsl_vector_scale(tmp, 1.0/gsl_blas_dnrm2(tmp));
            gsl_vector_add(state, tmp);
            cout << Util::print(state) << endl;
            maxOpt->maximize(interf, state, tmp, assertConstraints, minimizeNode, beta, -1, i);
            gsl_vector_memcpy(maxes[i], maxOpt->getMinState());
            gsl_vector_sub(state, tmp);
            cout << Util::print(state) << endl;
            numMaxes++;
        }
        return numMaxes;
    }

    vector<tuple<double, Predicate*, int>> getDistancesLevel0(const gsl_vector* state) {
        actualEval->run(state);

        vector<tuple<double, Predicate*, int>> s;

        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;

            double dist = actualEval->dist(p);
            double cost = abs(dist);
            s.push_back(make_tuple(cost, p, dist > 0));
        }

        sort(s.begin(), s.end());
        cout << "Distances: " << endl;
        int end = 5;
        if (end > s.size()) {
            end = s.size();
        }
        for (int i = 0; i < end; i++) {
            cout << get<1>(s[i])->print() << " " << get<0>(s[i]) << " " << get<2>(s[i]) << endl;
        }

        return s;
    }

    int chooseBeta(LocalState* ls) {
       double error1 = ls->errors[0];
       double error2 = ls->errors[1];
       double error3 = ls->errors[2];
       cout << "Errors: " << error1 << " " << error2 << " " << error3 << endl;

       if (error3 <= 0.0) {
            Assert(false, "Zero error?");
       }
       if (error3 < 1.0 || error2 <= 0.01) {
            return 2;
       }
       if (error2 < 5.0 || error1 <= 0.01) {
            return 1;
       } else {
            return 0;
       }
    }

    /*virtual IClause* getConflictClause(int level, LocalState* ls) {
        Assert(level == -1, "Other levels is not yet supported");
        int betaIdx = chooseBeta(ls);
        double beta;
        gsl_vector* state;
        if (betaIdx == 0) {
            beta = -1;
            state = ls->localSols[0];
        } else if (betaIdx == 1) {
            beta = -10;
            state = ls->localSols[1];
        } else if (betaIdx == 2) {
            beta = -50;
            state = ls->localSols[2];
        }
        cout << "Choosing beta = " << beta << endl;

        vector<tuple<double, Predicate*, int>> preds;
        actualEval->run(state);
        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;
            cout << p->print() << " ";
            double dist = actualEval->dist(p);
            cout << dist << " ";
            cout << Util::print(actualEval->grad(p)) << endl;
            double cost = abs(dist);
            double norm = actualEval->grad_norm(p);
            if (norm < 0.01) continue;
            if (cost/norm < 0.01) continue;
            preds.push_back(make_tuple(cost/norm, p, dist >= 0));
        }
        sort(preds.begin(), preds.end());

        int numPreds = min(20, preds.size());
        IClause* c = new IClause(GradUtil::counter);
        for (int i = 0; i < numPreds; i++) {
            Predicate* p = get<1>(preds[i]);
            cout << p->print() << " ";
            if (!p->isPure()) {
                p = make_pure(p, state);
            }
            c->add(p, !get<2>(preds[i]));
        }
        cout << endl;

        if (c->size() == 0) {
            Assert(false, "Empty clause?");
        }
        return c;
    }*/

    virtual IClause* getConflictClause(int level, LocalState* ls) {
        int betaIdx = chooseBeta(ls);
        double beta;
        gsl_vector* state;
        if (betaIdx == 0) {
            beta = -1;
            state = ls->localSols[0];
        } else if (betaIdx == 1) {
            beta = -10;
            state = ls->localSols[1];
        } else if (betaIdx == 2) {
            beta = -50;
            state = ls->localSols[2];
        }
        cout << "Choosing beta = " << beta << endl;
        int numMaxes = getMaxes(state, beta);
        cout << "Num maxes: " << numMaxes << endl;

        map<int, vector<tuple<double, Predicate*, int>>> maxesToPredicates;
        map<Predicate*, vector<tuple<double, int, int>>> predicatesToMaxes;

        map<Predicate*, tuple<double, int>> lsPreds;

        actualEval->run(state);
        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;
            double dist = actualEval->dist(p);
            double cost = abs(dist);
            double norm = actualEval->grad_norm(p);
            if (norm < 0.01) {
                cost = 1000.0;
            } else {
                cost = cost/norm;
            }
            lsPreds[p] = make_tuple(cost, dist >= 0);
        }

        for (int i = 0; i < numMaxes; i++) {
            actualEval->run(maxes[i]);
            vector<tuple<double, Predicate*, int>> preds;
            for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
                Predicate* p = *it;
                double dist = actualEval->dist(p);
                double cost = abs(dist);
                double norm = actualEval->grad_norm(p);
                if (norm < 0.01) {
                    cost = 1000.0;
                } else {
                    cost = cost/norm;
                }
                auto& lsTup = lsPreds[p];
                bool same = false;
                if (get<1>(lsTup) == (dist > 0)) {
                    same = true;
                }
                if (get<0>(lsTup) <= 0.1) {
                    same = true;
                }
                if (cost <= 0.1) {
                    same = true;
                }
                if (!same) {
                    preds.push_back(make_tuple(cost, p, dist >= 0));
                    if (predicatesToMaxes.find(p) == predicatesToMaxes.end()) {
                        vector<tuple<double, int, int>> maxstates;
                        predicatesToMaxes[p] = maxstates;
                    }
                    predicatesToMaxes[p].push_back(make_tuple(cost, i, dist>=0));
                }

            }
            sort(preds.begin(), preds.end());
            maxesToPredicates[i] = (preds);
        }

        set<int> selected;

        IClause* c = new IClause(GradUtil::counter);

        for (int i = 0; i < numMaxes; i++) {
            if (selected.find(i) != selected.end()) continue;
            auto& l = maxesToPredicates[i];
            if (l.size() == 0) continue;
            //Assert(l.size() != 0, "No preds to eliminate max");
            auto& tup = l[0];
            Predicate* p = get<1>(tup);
            auto& m = predicatesToMaxes[p];
            selected.insert(i);
            for (int k = 0; k < m.size(); k++) {
                selected.insert(get<1>(m[k]));
            }
            if (!p->isPure()) {
                p = make_pure(p, state);
            }
            c->add(p, get<2>(tup));
        }

        if (c->size() == 0) {
            Assert(false, "Empty clause?");
        }

        // add predicates that satisfy all maxes and the current local state
        vector<tuple<double, Predicate*>> preds;
        actualEval->run(state);
        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;
            auto& m = predicatesToMaxes[p];
            if (m.size() != 0) {
                continue;
            }
            double dist = actualEval->dist(p);
            double cost = abs(dist);
            double norm = actualEval->grad_norm(p);
            if (norm < 0.01) continue;
            preds.push_back(make_tuple(cost/norm, p));
        }
        sort(preds.begin(), preds.end());
        int numDirs = 0;
        vector<double> dists;
        dists.resize(dirGrads.size());
        for (int i = 0; i < preds.size(); i++) {
            if (numDirs >= dirGrads.size()) break;
            Predicate* p = get<1>(preds[i]);
            dists[numDirs] = actualEval->dist(p);
            gsl_vector_memcpy(dirGrads[numDirs], actualEval->grad(p));
            bool newDir = true;
            for (int j = 0; j < numDirs; j++) {
                if (Util::sameDir(dirGrads[j], dirGrads[numDirs])) {
                    if (dists[j] * dists[numDirs] >= 0.0) {
                        newDir = false;
                        //cout << "Not a new dir" << endl;
                        break;
                    }
                }
            }
            if (newDir) {
                cout << "New dir from ls: " << p->print() << " " << actualEval->dist(p) << " " << Util::print(actualEval->grad(p)) << endl;
                if (!p->isPure()) {
                    p = make_pure(p, state);
                }

                c->add(p, dists[numDirs] < 0);

                numDirs++;
            }
        }
        return c;

    }

    /*virtual IClause* getConflictClause(int level, LocalState* ls) {
        gsl_vector* state = ls->localSols[2];

        vector<tuple<double, Predicate*, int>> preds;
        actualEval->run(state);
        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;
            double dist = actualEval->dist(p);
            double cost = abs(dist);
            double norm = actualEval->grad_norm(p);
            if (norm < 0.01) continue;
            if (cost/norm < 0.1) continue;
            preds.push_back(make_tuple(cost/norm, p, dist >= 0));
        }

        int numPreds = min(5, preds.size());
        IClause* c = new IClause(GradUtil::counter);
        for (int i = 0; i < numPreds; i++) {
            int ridx = rand() % preds.size();
            Predicate* p = get<1>(preds[ridx]);
            cout << p->print() << " " << actualEval->dist(p) << " " << Util::print(actualEval->grad(p)) << endl;
            if (!p->isPure()) {
                p = make_pure(p, state);
            }
            c->add(p, !get<2>(preds[ridx]));
        }

        if (c->size() == 0) {
            Assert(false, "Empty clause?");
        }
        return c;

    }*/

    /*virtual IClause* getConflictClause(int level, LocalState * ls) {
        Assert(level == -1, "Other levels is not yet supported");
        int betaIdx = chooseBeta(ls);
        double beta;
        gsl_vector* state;
        if (betaIdx == 0) {
            beta = -1;
            state = ls->localSols[0];
        } else if (betaIdx == 1) {
            beta = -10;
            state = ls->localSols[1];
        } else if (betaIdx == 2) {
            beta = -50;
            state = ls->localSols[2];
        }
        cout << "Choosing beta = " << beta << endl;
        int numMaxes = getMaxes(state, beta);
        cout << "Num maxes: " << numMaxes << endl;

        // collect close preds
        set<Predicate*> close_preds;

        vector<tuple<double, Predicate*>> preds;
        actualEval->run(state);
        for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
            Predicate* p = *it;
            double dist = actualEval->dist(p);
            double cost = abs(dist);
            double norm = actualEval->grad_norm(p);
            if (norm < 0.01) continue;
            preds.push_back(make_tuple(cost/norm, p));
        }
        sort(preds.begin(), preds.end());
        int numDirs = 0;
        vector<double> dists;
        dists.resize(dirGrads.size());
        for (int i = 0; i < preds.size(); i++) {
            if (numDirs >= dirGrads.size()) break;
            Predicate* p = get<1>(preds[i]);
            dists[numDirs] = actualEval->dist(p);
            gsl_vector_memcpy(dirGrads[numDirs], actualEval->grad(p));
            bool newDir = true;
            for (int j = 0; j < numDirs; j++) {
                if (Util::sameDir(dirGrads[j], dirGrads[numDirs])) {
                    if (dists[j] * dists[numDirs] >= 0.0) {
                        newDir = false;
                        //cout << "Not a new dir" << endl;
                        break;
                    }
                }
            }
            if (newDir) {
                cout << "New dir from ls: " << p->print() << " " << actualEval->dist(p) << " " << Util::print(actualEval->grad(p)) << endl;
                numDirs++;
                if (p->isPure()) {
                    close_preds.insert(p);
                } else {
                    p = make_pure(p, state);
                    close_preds.insert(p);
                }
            }
        }

        for (int k = 0; k < numMaxes; k++) {
            preds.clear();
            actualEval->run(maxes[k]);
            for (auto it = interf->levelPredicates[0].begin(); it != interf->levelPredicates[0].end(); it++) {
                Predicate* p = *it;

                double dist = actualEval->dist(p);
                double cost = abs(dist);
                double norm = actualEval->grad_norm(p);
                if (norm < 0.01) continue;
                preds.push_back(make_tuple(cost/norm, p));
            }
            sort(preds.begin(), preds.end());
            int numDirs = 0;
            for (int i = 0; i < preds.size(); i++) {
                if (numDirs >= 20) break;
                Predicate* p = get<1>(preds[i]);
                dists[numDirs] = actualEval->dist(p);
                gsl_vector_memcpy(dirGrads[numDirs], actualEval->grad(p));
                bool newDir = true;
                for (int j = 0; j < numDirs; j++) {
                    if (Util::sameDir(dirGrads[j], dirGrads[numDirs])) {
                        if (dists[j] * dists[numDirs] >= 0.0) {
                            newDir = false;
                            //cout << "Not a new dir" << endl;
                            break;
                        }
                    }
                }
                if (newDir) {
                    cout << "New dir from max: " << p->print() << " " << actualEval->dist(p) << " " << Util::print(actualEval->grad(p)) << endl;
                    numDirs++;
                    if (p->isPure()) {
                        close_preds.insert(p);
                    } else {
                        p = make_pure(p, maxes[k]);
                        close_preds.insert(p);
                    }
                }
            }
        }
        // collect common conditions satisfied by state, maxes
        map<Predicate*, vector<tuple<double, int>>> common_preds;

        actualEval->run(state);
        for (auto it = close_preds.begin(); it != close_preds.end(); it++) {
            Predicate* p = *it;
            cout << "P: " << p->print() << " ";
            double dist = actualEval->dist(p);
            cout << dist << " ";
            cout << Util::print(actualEval->grad(p)) << endl;
            double cost = abs(dist);
            int val = dist >= 0;
            vector<tuple<double, int>> costValPairs;
            costValPairs.push_back(make_tuple(cost, val));
            common_preds[p] = costValPairs;
        }

        for (int i = 0; i < numMaxes; i++) {
            actualEval->run(maxes[i]);
            for (auto it = close_preds.begin(); it != close_preds.end(); it++) {
                Predicate* p = *it;

                double dist = actualEval->dist(p);
                double cost = abs(dist);
                int val = dist >= 0;
                common_preds[p].push_back(make_tuple(cost, val));
            }
        }

        vector<tuple<double, Predicate*, int>> ss;

        for (auto it = common_preds.begin(); it != common_preds.end(); it++) {
            Predicate* p = it->first;

            vector<tuple<double, int>>& pairs = common_preds[p];
            int val = -1;
            double cost = 1e10;

            cout << "Predicate " << p->print() << " " ;

            for (int i = 0; i < pairs.size(); i++) {
                auto& tup = pairs[i];
                double tcost = get<0>(tup);
                double tval = get<1>(tup);
                cout << "(" << tval << "," << tcost << ")" << " ";
                if (tcost < cost) {
                    cost = tcost;
                }
                if (tcost > 0.01) {
                    if (val == -1) {
                        val = tval;
                    } else if (val != tval) {
                        val = -2;
                        break;
                    }
                }
            }
            cout << "Final: " << val << "," << cost << endl;
            if (val == 0 || val == 1) {
                ss.push_back(make_tuple(cost, p, val));
            }
        }
        sort(ss.begin(), ss.end());

        cout << "Num common preds: " << ss.size() << endl;

        int numPreds = min(20, ss.size());
        IClause* c = new IClause(GradUtil::counter);
        for (int i = 0; i < numPreds; i++) {
            c->add(get<1>(ss[i]), !get<2>(ss[i]));
        }

        if (c->size() == 0) {
            Assert(false, "Empty clause?");
        }
        return c;
    }*/

};


#ifndef CTETRA_INPUT_H
#define CTETRA_INPUT_H

#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>

typedef void (*InputFn)(double k[3], gsl_vector *values);

typedef void (*UEInputFn)(double k[3], gsl_vector *evalues, gsl_matrix_complex *evecs);

#endif // CTETRA_INPUT_H