/*
Copyright (C) 2019-2020 JingWeiZhangHuai <jingweizhanghuai@163.com>
Licensed under the Apache License, Version 2.0; 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.
*/

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

#include <cblas.h>
#include "morn_tensor.h"

struct TensorBatchNormPara
{
    MLayer *prev;

    int res_valid;

    float rate;
    float decay;
    float momentum;
};

void *mTensorBatchNormPara(MList *ini,char *name)
{
    struct TensorBatchNormPara *para = (struct TensorBatchNormPara *)mMalloc(sizeof(struct TensorBatchNormPara));

    char *value = mINIRead(ini,name,"prev");
    para->prev = mNetworkLayer(ini,value);
    mException((para->prev == NULL),EXIT,"invalid prev");

    para->res_valid = (strcmp("Input",mLayerType(para->prev))!=0);

    value = mINIRead(ini,name,"rate");
    if(value != NULL) para->rate = atof(value);
    else
    {
        value = mINIRead(ini,"para","rate");
        if(value != NULL) para->rate = atof(value);
        else              para->rate = 0.001;
    }

    value = mINIRead(ini,name,"decay");
    if(value != NULL) para->decay = atof(value);
    else
    {
        value = mINIRead(ini,"para","decay");
        if(value != NULL) para->decay = atof(value);
        else              para->decay = 0.01;
    }
    mException((para->decay<0.0f)||(para->decay>=1.0f),EXIT,"invalid para decay");

    value = mINIRead(ini,name,"momentum");
    if(value != NULL) para->momentum = atof(value);
    else
    {
        value = mINIRead(ini,"para","momentum");
        if(value != NULL) para->momentum = atof(value);
        else              para->momentum = 0.9;
    }

    return para;
}

struct HandleTensorBatchNorm
{
    float *k;
    float *b;

    float *k_update;
    float *b_update;

    double *mean;
    double *var;
    float *roll_mean;
    float *roll_var;
};
void endTensorBatchNorm(void *info)
{
    struct HandleTensorBatchNorm *handle = (struct HandleTensorBatchNorm *)info;
    if(handle->k        != NULL) mFree(handle->k);
    if(handle->b        != NULL) mFree(handle->b);
    if(handle->mean     != NULL) mFree(handle->mean);
    if(handle->var      != NULL) mFree(handle->var );
    if(handle->roll_mean!= NULL) mFree(handle->roll_mean);
    if(handle->roll_var != NULL) mFree(handle->roll_var );
}
#define HASH_TensorBatchNorm 0xde8952f6

void TensorBatchNormSet(MLayer *layer)
{
    if(layer->state != DFLT) return;
    struct TensorBatchNormPara *para = (struct TensorBatchNormPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out=layer->tns;

    MHandle *hdl=mHandle(out,TensorBatchNorm);
    struct HandleTensorBatchNorm *handle = (struct HandleTensorBatchNorm *)(hdl->handle);

    if(handle->k != NULL) {mFree(handle->k);} handle->k =(float *)mMalloc(in->channel*sizeof(float));
    if(handle->b != NULL) {mFree(handle->b);} handle->b =(float *)mMalloc(in->channel*sizeof(float));

    if(handle->mean!= NULL) {mFree(handle->mean);} handle->mean=(double *)mMalloc(in->channel*sizeof(double));
    if(handle->var != NULL) {mFree(handle->var );} handle->var =(double *)mMalloc(in->channel*sizeof(double));

    if(handle->roll_mean!= NULL) {mFree(handle->roll_mean);} handle->roll_mean=(float *)mMalloc(in->channel*sizeof(float));
    if(handle->roll_var != NULL) {mFree(handle->roll_var );} handle->roll_var =(float *)mMalloc(in->channel*sizeof(float));

    mTensorRedefine(out,in->batch,in->channel,in->height,in->width,NULL);
    if(morn_network_flag == MORN_TRAIN)
    {
        if(INVALID_TENSOR(res)) mTensorRedefine(res,in->batch,in->channel,in->height,in->width,out->data);
        else                    mTensorRedefine(res,in->batch,in->channel,in->height,in->width,NULL);

        handle->k_update = handle->roll_mean;
        handle->b_update = handle->roll_var ;
        memset(handle->k_update,0,in->channel*sizeof(float));
        memset(handle->b_update,0,in->channel*sizeof(float));
    }

    if(morn_network_parafile==NULL)
    {
        for(int i=0;i<in->channel;i++){handle->k[i] = 1.0f;handle->b[i] = 1.0f;}
        // printf("handle->k[0] is %f\n",handle->k[0]);
        // printf("handle->b[0] is %f\n",handle->b[0]);
    }
    else
    {
        mNetworkParaRead(layer,"scale",handle->k,in->channel*sizeof(float));
        mNetworkParaRead(layer,"bias" ,handle->b,in->channel*sizeof(float));
        // printf("handle->k[0] is %f\n",handle->k[0]);
        // printf("handle->b[0] is %f\n",handle->b[0]);
    }
    hdl->valid = 1;
}

void mTensorBatchNormForward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("BatchNorm",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorBatchNormPara *para = (struct TensorBatchNormPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *out=layer->tns;

    TensorBatchNormSet(layer);

    MHandle *hdl=mHandle(out,TensorBatchNorm);
    struct HandleTensorBatchNorm *handle = (struct HandleTensorBatchNorm *)(hdl->handle);

    int in_size = in->height*in->width;

    memset(handle->mean,0,in->channel*sizeof(double));
    memset(handle->var ,0,in->channel*sizeof(double));
    for(int b=0;b<in->batch;b++)for(int c=0;c<in->channel;c++)for(int i=0;i<in_size;i++)
    {
        double data = (double)(in->data[b][c*in_size+i]);
        handle->mean[c] += data;
        handle->var [c] += data*data;
    }
    for(int c=0;c<in->channel;c++)
    {
        handle->mean[c] = handle->mean[c]/(double)(in->batch*in_size);
        handle->var [c] = handle->var [c]/(double)(in->batch*in_size);
        handle->var [c]-= (handle->mean[c]*handle->mean[c]);
        handle->var [c] = (handle->var[c]<=0.0f)?0.00001f:sqrt(handle->var[c]);
        // printf("handle->mean[c] is %f,handle->var[c] is %f\n",handle->mean[c],handle->var[c]);
    }

    if(morn_network_flag == MORN_PREDICT)
    {
        for(int c=0;c<in->channel;c++)
        {
            if(layer->state == DFLT)
            {
                handle->roll_mean[c] = handle->mean[c];
                handle->roll_var [c] = handle->var [c];
            }

            handle->mean[c] = handle->roll_mean[c]*0.99f+handle->mean[c]*0.01;
            handle->var [c] = handle->roll_var [c]*0.99f+handle->var [c]*0.01;
            handle->roll_mean[c] = handle->mean[c];
            handle->roll_var [c] = handle->var [c];
        }
    }

    for(int b=0;b<in->batch;b++)for(int c=0;c<in->channel;c++)for(int i=0;i<in_size;i++)
    {
        float data = in->data[b][c*in_size+i];
        out->data[b][c*in_size+i] = ((data-handle->mean[c])/handle->var[c])*handle->k[c]+handle->b[c];
    }

    layer->state = MORN_FORWARD;
}

void mTensorBatchNormBackward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("BatchNorm",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorBatchNormPara *para = (struct TensorBatchNormPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out= layer->res;

    MHandle *hdl=mHandle(layer->tns,TensorBatchNorm);
    struct HandleTensorBatchNorm *handle = (struct HandleTensorBatchNorm *)(hdl->handle);
    mException((hdl->valid == 0),EXIT,"no forward operate");

    mNetworkParaWrite(layer,"scale",handle->k,in->channel*sizeof(float));
    mNetworkParaWrite(layer,"bias" ,handle->b,in->channel*sizeof(float));

    int in_size = in->height*in->width;

    for(int c=0;c<in->channel;c++)
    {
        handle->k_update[c] *= para->momentum;
        handle->b_update[c] *= para->momentum;
    }
    for(int b=0;b<in->batch;b++)for(int c=0;c<in->channel;c++)for(int i=0;i<in_size;i++)
    {
        int idx = c*in_size+i;
        handle->k_update[c] += out->data[b][idx]*((in->data[b][idx]-handle->mean[c])/handle->var[c]);
        handle->b_update[c] += out->data[b][idx];
    }
    for(int c=0;c<in->channel;c++)
    {
        handle->k[c] = (handle->k[c]*(1.0f-(para->decay*para->rate)))-(handle->k_update[c]*para->rate/(float)(in->batch*in_size));
        handle->b[c] = (handle->b[c]*(1.0f-(para->decay*para->rate)))-(handle->b_update[c]*para->rate/(float)(in->batch*in_size));
    }

    if(para->res_valid==0) return;

    for(int b=0;b<in->batch;b++)for(int c=0;c<in->channel;c++)for(int i=0;i<in_size;i++)
    {
        int idx = c*in_size+i;
        float data;
        data = (in->data[b][idx] - handle->mean[c])/handle->var[c];
        data = (out->data[b][idx]*handle->k[c]/handle->var[c])*(1.0+(1.0f-data*data)/((float)(in->batch*in_size)));

        res->data[b][idx]=(para->prev->state == MORN_FORWARD)?data:(res->data[b][idx]+data);
    }

    para->prev->state = MORN_BACKWARD;
}


#ifndef CWANNIER_SPINORBIT_H
#define CWANNIER_SPINORBIT_H

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_matrix.h>
#include "HTightBinding.h"

typedef struct {
    gsl_complex val;
    int row;
    int col;
} socElem;

HTightBinding* HamiltonianWithSOC(double soc_strength, double theta, double phi, HTightBinding *Hrs_up, HTightBinding *Hrs_dn);

void addSOCElems(gsl_matrix_complex *Hr_onsite, double soc_strength, double theta, double phi);

gsl_matrix_complex* onSiteSOC_SpinZ();

void setWithHermitianConjugate(gsl_matrix_complex *M, int row, int col, gsl_complex val);

gsl_matrix_complex* onSiteSOC(double theta, double phi);

gsl_matrix_complex *RotationMatrix(double theta, double phi);
#endif // CWANNIER_SPINORBIT_H


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <pthread.h>

#include <gsl/gsl_sf.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h>

#include "SIShift.h"

static int  verbose  = FALSE;

static int  nLines = 10;

static char *usage[] = {"SIShift - \n",
                        "Required parameters:\n",
			"   -a    abundance data\n",
			"   -in   filename      parameter file  \n",
                        "Optional:\n",
			"   -b    integer       length of sample file to ignore\n",
			"   -s    integer       sampling frequency\n",
			"   -seed long          seed random number generator\n",
			"   -v                  verbose\n"};

void updateExpectations(double* adExpect, int nMax, double dAlpha, double dBeta, double dGamma, int nS, t_Data *ptData)
{
  int i = 0;

  for(i = 1; i <= nMax; i++){
    int nA = i;
    double dLog = 0.0, dP = 0.0;

    dLog = logLikelihood(nA, dAlpha, dBeta, dGamma);

    dP = exp(dLog);

    adExpect[i - 1] += dP*nS;
  }
}

int main(int argc, char* argv[]){
  t_Params    tParams;
  t_SIParams* atSIParams;
  int         i = 0, nSamples = 0, nMax = 0;
  t_Data      tData;
  double*     adExpect = NULL;
  double      dShift = 0.0;

  gsl_set_error_handler_off();

  getCommandLineParams(&tParams, argc, argv);

  /*allocate memory for samples*/
  atSIParams = (t_SIParams *) malloc(MAX_SAMPLES*sizeof(t_SIParams));
  if(!atSIParams)
    goto memoryError;

  /*read in Monte-Carlo samples*/
  readSamples(&tParams, atSIParams, &nSamples);

  readAbundanceData(tParams.szAbundFile, &tData);

  nMax = 1000;//tData.aanAbund[tData.nNA - 1][0];
  nMax = floor(pow(2.0,ceil(log((double) nMax)/log(2.0)) + 2.0) + 1.0e-7);

  adExpect = (double *) malloc(sizeof(double)*nMax);

  for(i = 0; i < nMax; i++){
    adExpect[i] = 0.0;
  }

  dShift = 5.0e5/(double) tData.nJ;
  for(i = 0; i < nSamples; i++){
    updateExpectations(adExpect, nMax,
		       atSIParams[i].dAlpha*sqrt(dShift),
		       atSIParams[i].dBeta*dShift,
		       atSIParams[i].dGamma,
		       atSIParams[i].nS,
		       &tData);
  }

  for(i = 1; i <= nMax; i++){
    printf("%d %f\n",i,adExpect[i - 1]/((double) nSamples));
  }

  free(adExpect);
  exit(EXIT_SUCCESS);

 memoryError:
  fprintf(stderr, "Failed to allocate memory in main aborting ...\n");
  fflush(stderr);
  exit(EXIT_FAILURE);
}

int intCompare(const void *pvA, const void *pvB)
{
  int* pnA = (int *) pvA, *pnB = (int *) pvB;

  if(*pnA < *pnB)
    return +1;
  else if(*pnA == *pnB){
    return 0;
  }
  else{
    return -1;
  }
}

int doubleCompare(const void *pvA, const void *pvB)
{
  double* pnA = (double *) pvA, *pnB = (double *) pvB;

  if(*pnA < *pnB)
    return +1;
  else if(*pnA == *pnB){
    return 0;
  }
  else{
    return -1;
  }
}

void writeUsage(FILE* ofp)
{
  int i = 0;
  char *line;

  for(i = 0; i < nLines; i++){
    line = usage[i];
    fputs(line,ofp);
  }
}

char *extractParameter(int argc, char **argv, char *param,int when)
{
  int i = 0;

  while((i < argc) && (strcmp(param,argv[i]))){
    i++;
  }

  if(i < argc - 1){
    return(argv[i + 1]);
  }

  if((i == argc - 1) && (when == OPTION)){
    return "";
  }

  if(when == ALWAYS){
    fprintf(stdout,"Can't find asked option %s\n",param);
  }

  return (char *) NULL;
}

void getCommandLineParams(t_Params *ptParams,int argc,char *argv[])
{
  char *szTemp = NULL;
  char *cError = NULL;

  /*get parameter file name*/
  ptParams->szInputFile  = extractParameter(argc,argv, INPUT_FILE,ALWAYS);
  if(ptParams->szInputFile == NULL)
    goto error;

 /*get abundance file name*/
  ptParams->szAbundFile  = extractParameter(argc,argv, ABUND_FILE,ALWAYS);
  if(ptParams->szAbundFile == NULL)
    goto error;

  /*get long seed*/
  szTemp  = extractParameter(argc,argv,SEED,OPTION);
  if(szTemp != NULL){
    ptParams->lSeed = strtol(szTemp,&cError,10);
    if(*cError != '\0'){
      goto error;
    }
  }
  else{
    ptParams->lSeed = 0;
  }

  /*get burn*/
  szTemp  = extractParameter(argc, argv, BURN, OPTION);
  if(szTemp != NULL){
    ptParams->nBurn = strtol(szTemp,&cError,10);
    if(*cError != '\0'){
      goto error;
    }
  }
  else{
    ptParams->nBurn = DEF_BURN;
  }

 /*get long seed*/
  szTemp  = extractParameter(argc, argv, SAMPLE, OPTION);
  if(szTemp != NULL){
    ptParams->nSample = strtol(szTemp,&cError,10);
    if(*cError != '\0'){
      goto error;
    }
  }
  else{
    ptParams->nSample = DEF_SAMPLE;
  }

  /*verbosity*/
  szTemp = extractParameter(argc, argv, VERBOSE, OPTION);
  if(szTemp != NULL){
    verbose = TRUE;
  }

  return;

 error:
  writeUsage(stdout);
  exit(EXIT_FAILURE);
}

void readSamples(t_Params *ptParams, t_SIParams *atSIParams, int *pnSamples)
{
  int  nSamples = 0;
  char *szInputFile = ptParams->szInputFile;
  char szLine[MAX_LINE_LENGTH];
  FILE *ifp = NULL;

  ifp = fopen(szInputFile, "r");

  if(ifp){
    while(fgets(szLine, MAX_LINE_LENGTH, ifp)){
      char *szTok = NULL, *szBrk = NULL, *pcError = NULL;
      int  nTime  = 0;

      /*remove trailing new line*/
      szBrk = strpbrk(szLine, "\n"); (*szBrk) = '\0';

      szTok = strtok(szLine, DELIM);

      nTime = strtol(szTok, &pcError, 10);
      if(*pcError != '\0') goto fileFormatError;

      if(nTime > ptParams->nBurn && nTime % ptParams->nSample == 0){

	szTok = strtok(NULL,DELIM);
	atSIParams[nSamples].dAlpha = strtod(szTok, &pcError);
	if(*pcError != '\0') goto fileFormatError;

	szTok = strtok(NULL,DELIM);
	atSIParams[nSamples].dBeta = strtod(szTok, &pcError);
	if(*pcError != '\0') goto fileFormatError;

	szTok = strtok(NULL,DELIM);
	atSIParams[nSamples].dGamma = strtod(szTok, &pcError);
	if(*pcError != '\0') goto fileFormatError;

	szTok = strtok(NULL,DELIM);
	atSIParams[nSamples].nS = strtol(szTok, &pcError, 10);
	if(*pcError != '\0') goto fileFormatError;

	nSamples++;
      }
    }

    fclose(ifp);
  }
  else{
    fprintf(stderr, "Failed to open file %s aborting\n", szInputFile);
    fflush(stderr);
    exit(EXIT_FAILURE);
  }

  (*pnSamples) = nSamples;
  return;

 fileFormatError:
  fprintf(stderr, "Incorrectly formatted input file aborting\n");
  fflush(stderr);
  exit(EXIT_FAILURE);
}

void readAbundanceData(const char *szFile, t_Data *ptData)
{
  int **aanAbund = NULL;
  int  i = 0, nNA = 0, nA = 0, nC = 0;
  int  nL = 0, nJ = 0;
  char szLine[MAX_LINE_LENGTH];
  FILE* ifp = NULL;

  ifp = fopen(szFile, "r");

  if(ifp){
    char* szTok   = NULL;
    char* pcError = NULL;

    fgets(szLine, MAX_LINE_LENGTH, ifp);

    szTok = strtok(szLine, DELIM2);

    nNA = strtol(szTok,&pcError,10);
    if(*pcError != '\0'){
      goto formatError;
    }

    aanAbund = (int **) malloc(nNA*sizeof(int*));

    for(i = 0; i < nNA; i++){
      aanAbund[i] = (int *) malloc(sizeof(int)*2);

      fgets(szLine, MAX_LINE_LENGTH, ifp);

      szTok = strtok(szLine, DELIM2);

      nA = strtol(szTok,&pcError,10);
      if(*pcError != '\0'){
	goto formatError;
      }

      szTok = strtok(NULL, DELIM2);

      nC = strtol(szTok,&pcError,10);
      if(*pcError != '\0'){
	goto formatError;
      }

      nL += nC;
      nJ += nC*nA;

      aanAbund[i][0]  = nA;
      aanAbund[i][1]  = nC;
    }
  }
  else{
    fprintf(stderr, "Failed to open abundance data file %s aborting\n", szFile);
    fflush(stderr);
    exit(EXIT_FAILURE);
  }

  ptData->nJ          = nJ;
  ptData->nL          = nL;
  ptData->aanAbund    = aanAbund;
  ptData->nNA         = nNA;
  return;

 formatError:
  fprintf(stderr, "Incorrectly formatted abundance data file\n");
  fflush(stderr);
  exit(EXIT_FAILURE);
}

double fX(double x, double dA, double dB, double dNDash)
{
  double dTemp1 = (dA*(x - dB)*(x - dB))/x;

  return log(x) - (1.0/dNDash)*(x + dTemp1);
}

double f2X(double x, double dA, double dB, double dNDash)
{
  double dRet = 0.0, dTemp = 2.0*dA*dB*dB;

  dRet = (1.0/(x*x))*(1.0 + (1.0/dNDash)*(dTemp/x));

  return -dRet;
}

double sd(int n, double dAlpha, double dBeta, double dGamma)
{
 double dA = 0.5*(-1.0 + sqrt(1.0 + (dAlpha*dAlpha)/(dBeta*dBeta)));
  double dN = (double) n, dNDash = dN + dGamma - 1.0, dRN = 1.0/dN;
  double dTemp1 = (0.5*dN)/(1.0 + dA), dTemp2 = 4.0*dRN*dRN*(1.0 + dA)*dA*dBeta*dBeta;
  double dXStar = dTemp1*(1.0 + sqrt(1.0 + dTemp2));
  double dFX = fX(dXStar, dA, dBeta, dNDash);
  double d2FX = -dNDash*f2X(dXStar, dA, dBeta, dNDash);
  double dLogK = 0.0, dGamma1 = dGamma;

  if(dGamma1 < 0.0){
    dGamma1 *= -1.0;
  }

  dLogK = gsl_sf_bessel_lnKnu(dGamma1,2.0*dA*dBeta);

  return -2.0*dA*dBeta -log(2.0) -dLogK -dGamma*log(dBeta) + dNDash*dFX + 0.5*log(2.0*M_PI) - 0.5*log(d2FX);
}

int bessel(double* pdResult, int n, double dAlpha, double dBeta, double dGamma)
{
  double dResult = 0.0;
  double dOmega = 0.0, dGamma2 = 0.0;
  double dLogK1 = 0.0, dLogK2 = 0.0;
  double dN = (double) n, dNu = dGamma + dN;
  double dTemp1 = 0.0;

  if(dNu < 0.0){
    dNu = -dNu;
  }

  if(dGamma < 0.0){
    dGamma2 = -dGamma;
  }
  else{
    dGamma2 = dGamma;
  }

  dOmega = sqrt(dBeta*dBeta + dAlpha*dAlpha) - dBeta;

  dLogK2 = gsl_sf_bessel_lnKnu(dNu, dAlpha);

  if(!gsl_finite(dLogK2)){
    if(dAlpha < 0.1*sqrt(dNu + 1.0)){
      //printf("l ");
      dLogK2 = gsl_sf_lngamma(dNu) + (dNu - 1.0)*log(2.0) - dNu*log(dAlpha);
    }
    else{
      //printf("sd ");
      (*pdResult) = dResult;
      return FALSE;
    }
  }

  dLogK1 = dGamma*log(dOmega/dAlpha) -gsl_sf_bessel_lnKnu(dGamma2,dOmega);

  dTemp1 = log((dBeta*dOmega)/dAlpha);

  dResult = dN*dTemp1 + dLogK2 + dLogK1;
  (*pdResult) = dResult;
  return TRUE;
}

double logLikelihood(int n, double dAlpha, double dBeta, double dGamma)
{
  double dLogFacN = 0.0;
  int status      = 0;
  double dRet     = 0.0;

  if(n < 50){
    dLogFacN = gsl_sf_fact(n);
    dLogFacN = log(dLogFacN);
  }
  else{
    dLogFacN = gsl_sf_lngamma(((double) n) + 1.0);
  }

  status = bessel(&dRet,n, dAlpha,dBeta,dGamma);
  if(status == FALSE){
    dRet = sd(n, dAlpha,dBeta,dGamma);
  }

  return dRet - dLogFacN;
}


#pragma once
#include "IntervalGrad.h"

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

#include "BooleanNodes.h"
#include "BooleanDAG.h"
#include "NodeVisitor.h"
#include "VarStore.h"
#include <map>
#include "FloatSupport.h"

#include <iostream>

using namespace std;

#define FINE_GRAIN_RANGES 1

class RangeDiff: NodeVisitor
{
	FloatManager& floats;
	BooleanDAG& bdag;
  map<string, int> floatCtrls; // Maps float ctrl names to indices with grad vectors
  int nctrls; // number of float ctrls
	gsl_vector* ctrls; // ctrl values
  vector<IntervalGrad*> ranges; // Keeps track of ranges for each node
	vector<DistanceGrad*> distances; // Keeps track of distance metric for boolean nodes
	map<int, int> inputValues; // Maps node id to values set by the SAT solver
	double error = 0.0;
	gsl_vector* errorGrad;

	int DEFAULT_INP = -1;
	int assertCtr;
	bool foundFailure;

	static constexpr float ASSERT_PENALTY = 1.0;

public:
	int failedAssert;

  RangeDiff(BooleanDAG& bdag_p, FloatManager& _floats, const map<string, int>& floatCtrls_p);
  ~RangeDiff(void);

  virtual void visit( SRC_node& node );
  virtual void visit( DST_node& node );
  virtual void visit( CTRL_node& node );
  virtual void visit( PLUS_node& node );
  virtual void visit( TIMES_node& node );
  virtual void visit( ARRACC_node& node );
  virtual void visit( DIV_node& node );
  virtual void visit( MOD_node& node );
  virtual void visit( NEG_node& node );
  virtual void visit( CONST_node& node );
  virtual void visit( LT_node& node );
  virtual void visit( EQ_node& node );
	virtual void visit( AND_node& node );
	virtual void visit( OR_node& node );
	virtual void visit( NOT_node& node );
  virtual void visit( ARRASS_node& node );
  virtual void visit( UFUN_node& node );
  virtual void visit( TUPLE_R_node& node );
	virtual void visit( ASSERT_node& node );

  double run(const gsl_vector* ctrls_p, map<int, int>& inputValues_p, gsl_vector* errorGrad_p);

	void setrange(bool_node& bn, IntervalGrad* r) {
		ranges[bn.id] = r;
	}

	IntervalGrad* r(bool_node& bn) {
		IntervalGrad* interval = ranges[bn.id];
		if (interval == NULL) {
			gsl_vector* l = gsl_vector_alloc(nctrls);
			gsl_vector* h = gsl_vector_alloc(nctrls);
			interval = new IntervalGrad(0, 0, l, h);
			setrange(bn, interval);
		}
		return interval;
	}

	IntervalGrad* r(bool_node* bn) {
		return r(*bn);
	}

	void setdistance(bool_node& bn, DistanceGrad* d) {
		distances[bn.id] = d;
	}

	DistanceGrad* d(bool_node& bn) {
		DistanceGrad* dist = distances[bn.id];
		if (dist == NULL) {
			gsl_vector* g = gsl_vector_alloc(nctrls);
			dist = new DistanceGrad(0, g);
			setdistance(bn, dist);
		}
		return dist;
	}

	DistanceGrad* d(bool_node* bn) {
		return d(*bn);
	}

	void print() {
		for (int i = 0; i < bdag.size(); i++) {
			cout << bdag[i]->lprint() << " ";
			if (bdag[i]->getOtype() == OutType::FLOAT) {
				cout << r(bdag[i])->print() << endl;
			} else {
				cout << d(bdag[i])->print() << endl;
			}
		}
	}
	void printFull() {
    for (int i = 0; i < bdag.size(); i++) {
			cout << bdag[i]->lprint() << endl;
			if (bdag[i]->getOtype() == OutType::FLOAT) {
				cout << r(bdag[i])->printFull() << endl;
			} else {
				cout << d(bdag[i])->printFull() << endl;
			}
		}
  }

	bool isFloat(bool_node& bn) {
		return (bn.getOtype() == OutType::FLOAT);
	}

	bool isFloat(bool_node* bn) {
		return (bn->getOtype() == OutType::FLOAT);
	}

	int getInputValue(bool_node& bn) {
		if (inputValues.find(bn.id) != inputValues.end()) {
			int val = inputValues[bn.id];
			Assert(val == 0 || val == 1, "NYI: Integer values");
			return val;
		} else {
			return DEFAULT_INP;
		}
	}

	int getInputValue(bool_node* bn) {
		return getInputValue(*bn);
	}
	void computeError(double dist, int expected, gsl_vector* dg, bool_node& node, bool relax = false);

};


#include <gsl/gsl_math.h>
#include <gsl/gsl_cblas.h>
#include "cblas.h"

void
cblas_ssyr2k (const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo,
              const enum CBLAS_TRANSPOSE Trans, const int N, const int K,
              const float alpha, const float *A, const int lda,
              const float *B, const int ldb, const float beta, float *C,
              const int ldc)
{
#define BASE float
#include "source_syr2k_r.h"
#undef BASE
}


//
//  global.h
//  Agile_RRT
//
//  Created by Timothy M. Caldwell.
//  Copyright (c) 2015 Timothy M. Caldwell. Refer to license.txt
//

#ifndef __Agile_RRT__global__
#define __Agile_RRT__global__

#include <fstream>
//#include <string>
//#include <sstream>

#include <iostream>
#include <iomanip>

#include <stdlib.h>
#include <math.h>
#include <random>

//boost libraries
#include <boost/numeric/odeint.hpp>

//gsl libraries
#include <gsl/gsl_spline.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h>

//Eigen libraries
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Core>
#include <eigen3/Eigen/Geometry>

//namespaces
using namespace std;
namespace odeint = boost::numeric::odeint;

extern const int NI;                              // Constant; Number of inputs to the dynamic system
extern const int NS;                              // Constant; Number of states to the dynamic system

extern const double kLARGENUM;                    // Constant; A sufficiently large number

// Constants; For integration using odeint
extern const int kMAX_NUM_INTEGRATION_STEPS;
extern const double kINTEGRATION_DT_START;

//integrate state types
extern const int kFREEDYNAMICS_TYPE;
extern const int kFEEDFORWARD_TYPE;
extern const int kTRACKPOINT_TYPE;
extern const int kTRACKTRAJECTORY_TYPE;
extern const int kLINEARSTEERINGPROJECTION_TYPE;
extern const int kLINEARSTEERINGPROJECTIONWCOST_TYPE;

//integrate linear state and ricatti types
extern const int kLTI_TYPE;
extern const int kLTV_TYPE;

//STMSlot2 types
extern const int kLTI_OPENLOOP_TYPE;
extern const int kLTI_CLOSEDLOOP_TYPE;
extern const int kLTV_OPENLOOP_TYPE;
extern const int kLTV_CLOSEDLOOP_TYPE;

//control type for linear state equation
extern const int kFREE_CONTROLTYPE;
extern const int kFEEDFORWARD_CONTROLTYPE;
extern const int kMINENERGY_OPENLOOP_CONTROLTYPE;
extern const int kMINENERGY_CLOSEDLOOP_CONTROLTYPE;

//Gramian types
extern const int kWW_GRAMIANTYPE;
extern const int kWWK_GRAMIANTYPE;
extern const int kSSK_GRAMIANTYPE;

typedef vector<double> ode_state_type;            // odeint integrator state_type

///////////////////////////////
//Currently implemented as global. Needs to be updated for NS and NI
///////////////////////////////
// Eigen Matrix types of size NS or NI by NS or NI
typedef Eigen::Matrix< double , 8 , 8 > NSxNS_type;
typedef Eigen::Matrix< double , 8 , 1 > NSxNI_type;
typedef Eigen::Matrix< double , 1 , 1 > NIxNI_type;
typedef Eigen::Matrix< double , 1 , 8 > NIxNS_type;
typedef Eigen::Matrix< double , 8 , 1 > NSx1_type;
typedef Eigen::Matrix< double , 1 , 1 > NIx1_type;

struct constraints_struct {
  Eigen::Matrix<double, 8, 2> xbounds;              // list of upper and lower bounds for each state
  Eigen::Matrix<double, 1, 2> ubounds;              // list of upper and lower bounds for each input

  // list of circle obstacles of radius obs_radii[.] at point (obs_Xs[.], obs_Ys[.])
  vector< double > obs_radii;
  vector< double > obs_Xs;
  vector< double > obs_Ys;
};
// list of upper and lower bounds for each state and input. Should be the same as those set in constraints_struct
struct sampling_bounds_struct {
  Eigen::Matrix<double, 8, 2> xbounds;
  Eigen::Matrix<double, 1, 2> ubounds;
};

// stepper type used by odeint for integration
//typedef odeint::controlled_runge_kutta< odeint::runge_kutta_dopri5<ode_state_type> > stepper_type; // Alternative stepper choice
typedef odeint::controlled_runge_kutta< odeint::runge_kutta_cash_karp54<ode_state_type> > stepper_type;

// Printing parameters
extern const int kPRINT_PRECISION;                // Printing precision number
extern const double kPRINT_CHOP_BOUND;            // Lower bound for chopping (i.e. if abs(x)<kPRINT_CHOP_BOUND, then 0 is returned)

// prints a vec to stream in a form easily readable by Mathematica.
extern void PrintVec(const vector<double> & vec, ostream * stream = &(std::cout));
extern void PrintVec(const vector<double> & vec, const string & name, ostream * stream);

#endif // __Agile_RRT__global__


#include <stdio.h>
#include <math.h>
#include <gsl/gsl_sf_bessel.h>
#include "global.h"
#include "utility.h"
#include "constants.h"
#include <gsl/gsl_cblas.h>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <omp.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include "steppmethods.h"

#define NUMBER_SAMPLES 1
#define NUMBER_VOLUMES 1
#define SHIFT_LATE 200
#define BINS 100
#define PDF_NR 12

// regular main but with multiple sampling runs over multiple volume fractions phi.

int
main (void){
	FILE *source, *target;
  	source = fopen("global.h", "r"); // get source global.h at program start so as to not copy edits during run

  	if( source == NULL )
  	{
  		printf("cannot open global.h\n");
  		exit(EXIT_FAILURE);
  	}
  	SETNUMTHREADS
  	int i,j,l,k,m ,n;
	Constants();									// draw coupling constants and other derived constants
	double *t = calloc(LENGTH_T, sizeof(double));
	time_t tme = time(NULL);
	char ordnerLABEL[300];
	char ueberordnerLABEL[300];
	char ordnerLABEL_original[300];
	strftime(ordnerLABEL, sizeof ordnerLABEL, "%A %c", localtime(&tme));
  	double squares_all_calcs[LENGTH_T][NUMBER_SAMPLES];			// for printing all squares in shared table
  	i = 0;
  	printf("\n%s\n", ordnerLABEL);
  	deleteSpaces(ordnerLABEL,ordnerLABEL);
  	strcpy (ordnerLABEL_original, ordnerLABEL);    	// keep original, change ordnerLABEL per run
  	double alpha_global = ALPHA;
  	double temp_global = TEMP;						// macros from global make trouble when used as print arguments, hence copy to memory
  	sprintf(ueberordnerLABEL,"N=%d_SIM=%d_T=%1.1E_ALPH=%1.1f", OSSZI, SIMULATIONEN, temp_global, alpha_global);
  	char tempLABEL[300];
  	strcpy (tempLABEL, ordnerLABEL);
  	strcat(tempLABEL, LABELPREFIX);
  	strcat(tempLABEL, ueberordnerLABEL);
  	strcpy (ueberordnerLABEL, tempLABEL);
  	struct stat st2 = {0};

  	double prob_density[BINS][PDF_NR];
  	for (i = 0; i < BINS ; i++)
  	{
  		for (j = 0; j < PDF_NR; j++)
  		{
  			prob_density[i][j] = 0.0;
  		}
  	}
  	int check_times[PDF_NR];
  	for (j = 0; j < PDF_NR; j++)
  	{
  		check_times[j] = (int) ( 10.0 * (j+1)) ;
  	}
  	check_times[PDF_NR-1] = LENGTH_T - 1;
  	for (j = 0; j < PDF_NR; j++)
  	{
  		printf("check times nr %d is %d \n", j, check_times[j]);
  	}
  	printf("\nprint results into ");
  	printf(ueberordnerLABEL);

  	if (stat(ueberordnerLABEL, &st2) == -1)
	{												//Teste ob Ordner existiert, erstelle Ordner
		mkdir(ueberordnerLABEL, 0700);
	}
	if (chdir(ueberordnerLABEL))						// change into directory of simulation
	{
		printf("Error changing directory");
		return 1;
	}

  	// reserviere speicher fuer eine Vollständige trajektorie ------------------
	m = LENGTH_T; n = ORDER;
	double **z = (double **) malloc(m * sizeof(double *));
	z[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) z[i] = z[0] + n *i;

		double *squares = calloc(LENGTH_T, sizeof(double));
	double *short_correlation_x = calloc(LENGTH_T, sizeof(double));	//short term (1 timestep) correlations
	double *short_correlation_y = calloc(LENGTH_T, sizeof(double));
	double *total_momentum_x = calloc(LENGTH_T, sizeof(double));
	double *total_momentum_y = calloc(LENGTH_T, sizeof(double));
	double *long_correlation_x = calloc(LENGTH_T, sizeof(double)); //longterm correlation from first timestep
	double *long_correlation_y = calloc(LENGTH_T, sizeof(double));
	double *squares_increase = calloc(LENGTH_T, sizeof(double));	// Increase a from assumed form x^2 0 = a t
	double real_volumes[10];

	double *bathp = calloc(LENGTH_T, sizeof(double));
	double *bathq = calloc(LENGTH_T, sizeof(double));

	double *px_correlation = calloc(LENGTH_T, sizeof(double));	// save correlation of impulse_x to px(t=0)
	double *py_correlation = calloc(LENGTH_T, sizeof(double));
	double *px_correlation_late =  calloc(LENGTH_T, sizeof(double));	// corelation after time SHIFT_LATE
	double *px_correlation_late2 = calloc(LENGTH_T, sizeof(double));	// corelation after time SHIFT_LATE

	int n_zplots = SIMULATIONEN;
	if (SIMULATIONEN > MAX_NR_PLOTS ) n_zplots = MAX_NR_PLOTS;    // speichere maximal 30*DIM trajectorien
  	// reserviere speicher fuer samplepfade massives Teilchen zum plotten ------------------
	m = LENGTH_T; n = DIM*n_zplots;
	double **zplots = (double **) malloc(m * sizeof(double *));
	zplots[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) zplots[i] = zplots[0] + n *i;

	// reserviere speicher fuer samplepfade 1. bad  Teilchen zum plotten ------------------
		m = LENGTH_T; n = DIM*n_zplots;
	double **qplots = (double **) malloc(m * sizeof(double *));
	qplots[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) qplots[i] = qplots[0] + n *i;

	// reserviere speicher fuer samplepfade letztes bad  Teilchen zum plotten ------------------
		m = LENGTH_T; n = DIM*n_zplots;
	double **qlplots = (double **) malloc(m * sizeof(double *));
	qlplots[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) qlplots[i] = qlplots[0] + n *i;

	// reserviere speicher fuer samplepfade zusätzlicher Kac_Zwanzig Kraft Term ------------------
		m = LENGTH_T; n = DIM*n_zplots;
	double **Zwanzig_Force = (double **) malloc(m * sizeof(double *));
	Zwanzig_Force[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) Zwanzig_Force[i] = Zwanzig_Force[0] + n *i;

	// reserviere speicher fuer probability Density ------------------
		m = 80; n = 80;
	double **P_density = (double **) malloc(m * sizeof(double *));
	P_density[0] = (double *) malloc(m* n * sizeof(double));
	for (i=1; i<m; i++) P_density[i] = P_density[0] + n *i;
  	// reserviere speicher fuer Gitterkoordinaten ------------------

		double DENSITY_POINTS = 0.0;

	for(m=0; m<80; m++)
	{
		for(n=0; n<80; n++)
		{
			P_density[m][n] = 0.0;
		}
	}

	double *EKIN = calloc(LENGTH_T, sizeof(double));
	double *EBAD = calloc(LENGTH_T, sizeof(double));
	double *ETOT = calloc(LENGTH_T, sizeof(double));
	double *PTOT = calloc(LENGTH_T, sizeof(double));
	double *PTOTY = calloc(LENGTH_T, sizeof(double));
	double *LTOT = calloc(LENGTH_T, sizeof(double));

	// GSL random number Setup for Taus Generator
	const gsl_rng_type * T;
	gsl_rng * r;
	gsl_rng_env_setup();
	T = gsl_rng_ranlxd2;
	r = gsl_rng_alloc (T);
	gsl_rng_set(r, time(NULL)); // Seed with time
	// RNG setup End //
	// GSL random number Setup for Taus Generator for global rand
	const gsl_rng_type * T2;
	T2 = gsl_rng_ranlxd2;
	RAND_GLOBAL = gsl_rng_alloc (T2);
	gsl_rng_set(RAND_GLOBAL, time(NULL)); // Seed with time
	// RNG setup End //

	clock_t start, end;

	printf("evenly spread t \n");	//set up time vec
  	for (i = 0; i < 10; i++) // first ten values in small steps
  	{
  		t[i] =  i * TIME_STEPS/10.0;
  	}
  	for (i = 10; i < LENGTH_T; i++)
  	{
  		t[i] = (i-9) * TIME_STEPS;
  	}
  	double squarespace = 0.0;

  	DIFF_COEFF = KBOLTZ * TEMP/GAMMA * sin(M_PI * ALPHA ) / (M_PI * ALPHA);
  	if(ALPHA > 0.95)
  	{
  		DIFF_COEFF = KBOLTZ * TEMP/GAMMA;
  	}
  	LATTICE_SPACING =  sqrt(2*DIM*DIFF_COEFF * pow( TIME_ARRIVAL, ALPHA) );
  	printf("LATTICE_SPACING = %3.3E\n", LATTICE_SPACING);
  	double latlength = LATTICE_SPACING;

		//------------------------------EIGENTLICHE SIMULATION ------------------------------------------------------------------
  	int sim_count = 0;
  	start = clock();
  	time_first_contact = 0.0;
  	int error_count = 0;
		//
  	for (i = 0; i < SIMULATIONEN; i++)
  	{

		vec_zero_i(lattice_position,DIM);		// set stating cell to zero !
		TARGET_CONTACT = 0;
		FLAG_DUMP_VALUE = 0;					// set to 1 if error occures
		// -------------------------BAD
		Bath_Setup(y, LABEL, r);				// draw intial conditions
		//---------------------Bad aufgesetzt
		VVerlet_parallel(ORDER, y, z, t, Poti_Handle);
		end = clock();
		//------------------------------auswertung Integrationsergebnisse
		if ( !(FLAG_DUMP_VALUE))			// only count result without err
		{
		  	double *kahan_c = calloc(LENGTH_T, sizeof(double));                //
		  	double *kahan_t = calloc(LENGTH_T, sizeof(double));
		  	double *kahan_y = calloc(LENGTH_T, sizeof(double));
		  	#pragma omp parallel for
		  	for(l = 0; l < LENGTH_T; l++)
		  	{
		  		if (l>10)
		  		{
		  			long_correlation_x[l] += z[1][0] *z[l][0]/((double) SIMULATIONEN);
		  			long_correlation_y[l] += z[1][1] *z[l][1]/((double) SIMULATIONEN);
		  		}
		  		kahan_y[l] = z[10][DIM] * z[l][DIM] / ((double) SIMULATIONEN) / (mass * KBOLTZ * TEMP)\
		  					- kahan_c[l];
		  		kahan_t[l] = px_correlation[l] + kahan_y[l];
		  		kahan_c[l] = (kahan_y[l] - px_correlation[l]) - kahan_y[l];
		  		px_correlation[l] = kahan_t[l];
		  		py_correlation[l] += z[10][DIM + 1] * z[l][DIM + 1] / ((double) SIMULATIONEN) / (mass * KBOLTZ * TEMP);
		  		px_correlation_late[l] += z[SHIFT_LATE][DIM] * z[l][DIM] / ((double) SIMULATIONEN) / (mass * KBOLTZ * TEMP);
		  		px_correlation_late2[l] += z[SHIFT_LATE * 2][DIM] * z[l][DIM] / ((double) SIMULATIONEN) / (mass * KBOLTZ * TEMP);
		  		for(j=0; j< DIM; j++)
		  		{
		  			if (i < n_zplots)
		  			{
		  				zplots[l][DIM * i + j] = z[l][j];
		  				qplots[l][DIM * i + j] = z[l][2*DIM + j];
		  				double Force_TEMP = 0.0;
		  				for(int os = 0; os < OSSZI; os++)
		  					{
		  						Force_TEMP += (pow(coupling[os],2.0)/pow(ommega[os],2.0) - coupling[os]) * z[l][(2 + 2 * os) * DIM + j];
		  					}
		  					Zwanzig_Force[l][DIM * i + j]=Force_TEMP;
		  					int i_last = OSSZI -1;
		  					qlplots[l][DIM * i + j] = z[l][2*OSSZI*DIM + j];
		  			}
		  			squares[l] = squares[l] + pow(z[l][j],2.0)/((double) SIMULATIONEN);
		  			EKIN[l] = EKIN[l] + pow(z[l][DIM+j],2.0)/(mass * SIMULATIONEN * KBOLTZ *TEMP * DIM);
		  			for(k=0;k<OSSZI;k++)
		  			{
		  				bathq[l] = bathq[l] + pow(z[l][(2 + 2*k) *DIM +j],2.0)/( (double) (OSSZI*SIMULATIONEN)\
		  						*KBOLTZ*TEMP) * pow(ommega[k],2.0)/2.0;// mittlere potentielle E pro Badteilchen
		  																//reskaliert durch kT
		  				bathp[l] = bathp[l] + pow(z[l][(3 + 2*k) *DIM +j],2.0)/( (double) (OSSZI*SIMULATIONEN)\
		  						*KBOLTZ*TEMP) /2.0/massq[k];						// mittlere kinetische E pro Badteilchen
		  																//reskaliert durch kT
		  				EBAD[l] = EBAD[l] + pow(z[l][(3 + 2*k) *DIM +j], 2.0)/2.0\
		  						+ 0.5  * pow(ommega[k],2.0) \
		  						* pow(z[l][(2 + 2*k) *DIM +j] - coupling[k]/(pow(ommega[k],2.0) ) * z[l][j] , 2.0) /((double) SIMULATIONEN);
		  			}
		  		}
			  	for(k=0;k<OSSZI;k++)	// add p and angular momentum L for bath
			  	{
			  		PTOT[l] 	+= z[l][(3 + 2*k) *DIM + 0];
			  		PTOTY[l]	+= z[l][(3 + 2*k) *DIM + 1];
			  		LTOT[l] += 	(z[l][(2 + 2*k) *DIM + 0] * z[l][(3 + 2*k) *DIM + 1]\
			  					- z[l][(2 + 2*k) *DIM + 1] * z[l][(3 + 2*k) *DIM + 0]) /((double) SIMULATIONEN);
			  	}
			  	PTOT[l] += z[l][2];
			  	PTOT[l] += z[l][3];
			  	LTOT[l] += (z[l][0] * z[l][3] - z[l][1] * z[l][2]) / ((double) SIMULATIONEN);
			}
			if(DIM == 1)
			{
			  	// setup pdf checks
			  	double bins = BINS * 1.0;
			  	double dx = LATTICE_SPACING/bins;
			  	double sims = SIMULATIONEN;
			  	for (j = 0; j < PDF_NR; j++)
			  	{
			  		int t_check = check_times[j];
			  		for (int i_bin = 0; i_bin < BINS; i_bin++)
			  		{
			  			double lower = -LATTICE_SPACING / 2.0 + i_bin * dx;
			  			double upper =  -LATTICE_SPACING / 2.0 + (i_bin + 1) * dx;
			  			if 	( 	( z[t_check][0] <= upper) &&
			  					( z[t_check][0] > lower)
			  				)
			  			{
			  				prob_density[i_bin][j] += 1.0/sims;
			  			}
			  		}
			  	}
			}
			sim_count += 1;
			//printf("%d  und t %4.2f \n", i ,((double) (end - start)));
			printf("\r%d von %d mit t/count = %4.2f s und average t_rest =%4.2f h ", sim_count , SIMULATIONEN, (double) (end - start) /  sim_count/ THREADS / CLOCKS_PER_SEC ,\
			  		((double) (end - start) / sim_count/ THREADS/ CLOCKS_PER_SEC * (SIMULATIONEN - sim_count)/3600));
			printf(" %d hits on Lat  and avrgtcntct = %4.2f",TARGET_CONTACT, time_first_contact/(i+1) );
			fflush(stdout);
			for(l = 0; l < LENGTH_T; l++)
			{
			  	for (int oss_i=0; oss_i < OSSZI; oss_i++)
			  	{
			  		total_momentum_x[l] += z[l][(3 + 2*oss_i) *DIM + 0];
			  		total_momentum_y[l] += z[l][(3 + 2*oss_i) *DIM + 1];
			  	}
			  	total_momentum_x[l] += z[l][(0 + 2*0) *DIM + 0];
			  	total_momentum_y[l] += z[l][(0 + 2*0) *DIM + 1];
			}
			// -------------- end evaluation if----------------------
			free(kahan_y); free(kahan_t); free(kahan_c);
		}else
		{
			error_count++;
			printf("\n err occured at calc %d, calculation dumped, %d totatl errors\n",  i,error_count );
			i -= 1;		// do one more calculation
		}

	}
	tme = time(NULL);
	char end_label[90];
	strftime( end_label, sizeof end_label, "%A %c", localtime(&tme));
	printf("\n%s\n", end_label);
	for(l = 0; l < LENGTH_T; l++)
	{
		ETOT[l] = EKIN[l] + EBAD[l];
			  	PTOT[l] = sqrt(pow(PTOT[l]/((double) SIMULATIONEN),2.0) + pow(PTOTY[l]/((double) SIMULATIONEN),2.0));
	}
	for(l = 1; l < LENGTH_T; l++)
	{
		squares_increase[l] = (squares[l]- squares[l-1])/(t[l] - t[l-1]);
	}
			//------------------------------ENDE SIMULATION, speichere daten ------------------------------------------------------------------
	for(m=0; m<80; m++)
	{
		for(n=0; n<80; n++)
		{
			P_density[m][n] *= 1.0/DENSITY_POINTS;
		}
	}
	FILE *fp;

	struct stat st = {0};

	if (stat(ordnerLABEL, &st) == -1){		//Teste ob Ordner existiert, erstelle Ordner
		mkdir(ordnerLABEL, 0700);
	}

	fp = fopen ("shellscript.sh", "w");		//create Shellscript to start Gnuplot from c main()
	fprintf(fp, "cd %s\n", ordnerLABEL);
	fprintf(fp, "gnuplot gnuplot.txt\n");
	fprintf(fp, "cd ..");
	fclose (fp);

	 // copy global.h 	/

   	// copy to same name into directory, clould be anything different

   	if (chdir(ordnerLABEL))			// change into directory of simulation
   	{
   		printf("Error changing directory");
   		return 1;
   	}

   	mkdir("plots", 0700);
   	mkdir("trajec", 0700);

   	target = fopen("global.h", "w");

   	if( target == NULL )
   	{
   		fclose(source);
   		printf("Press any key to exit...\n");
   		exit(EXIT_FAILURE);
   	}
   	char ch;
   	while( ( ch = fgetc(source) ) != EOF )
   	{
   		fputc(ch, target);
   	}

   	printf("File copied successfully.\n");

   	fclose(target);

   	fp = fopen ("latlength.dat", "w");
   	fprintf(fp, "%lf  \n", latlength);
   	fclose (fp);

   	fp = fopen ("squares_rohdaten.dat", "w");
   	for(l = 0; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l], squares[l]);
   	}
   	fclose (fp);

   	fp = fopen ("ekin.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, EKIN[l]);
   	}
   	fclose (fp);

   	fp = fopen ("ekinbath.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, bathp[l]);
   	}
   	fclose (fp);

   	fp = fopen ("ommega.dat", "w");
   	for(l = 0; l < OSSZI; l++){
   		fprintf(fp, "%lf\n", ommega[l]);
   	}
   	fclose (fp);

   	fp = fopen ("PTOT.dat", "w");
   	for(l = 0; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, PTOT[l]);
   	}
   	fclose (fp);

   	fp = fopen ("P_X.dat", "w");
   	for(l = 0; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %1.3e\n", t[l]/TIME_ARRIVAL, total_momentum_x[l]);
   	}
   	fclose (fp);

   	fp = fopen ("P_Y.dat", "w");
   	for(l = 0; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %1.3e\n", t[l]/TIME_ARRIVAL, total_momentum_y[l]);
   	}
   	fclose (fp);

   	fp = fopen ("xshortcorellation.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, short_correlation_x[l]/latlength/latlength);
   	}
   	fclose (fp);

   	fp = fopen ("yshortcorellation.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, short_correlation_y[l]/latlength/latlength);
   	}
   	fclose (fp);

   	fp = fopen ("xlongcorellation.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, long_correlation_x[l]/latlength/latlength);
   	}
   	fclose (fp);

   	fp = fopen ("ylongcorellation.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l]/TIME_ARRIVAL, long_correlation_y[l]/latlength/latlength);
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, px_correlation[l]);
   	}
   	fclose (fp);

   	fp = fopen ("PY_corr.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, py_correlation[l]);
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr_late.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, px_correlation_late[l]);
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr_late2.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, px_correlation_late2[l]);
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr_abs.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, fabs(px_correlation[l]));
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr_late_abs.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, fabs(px_correlation_late[l]));
   	}
   	fclose (fp);

   	fp = fopen ("PX_corr_late2_abs.dat", "w");
   	for(l = 1; l < LENGTH_T; l++){
   		fprintf(fp, "%1.3E  %1.3E\n", t[l]/TIME_ARRIVAL, fabs(px_correlation_late2[l]));
   	}
   	fclose (fp);

   	if (DIM == 1)
   	{
   		fp = fopen ("PDF_1D.dat", "w");
   		for (int pos = 0; pos < BINS; pos++)
   		{
   			fprintf(fp, "%1.3E ", -LATTICE_SPACING/2.0 + pos*LATTICE_SPACING/((double) BINS));
   			for(j = 0; j < PDF_NR; j++)
   			{
   				fprintf(fp, "%1.3E ", prob_density[pos][j]);
   			}
   			fprintf(fp,"\n");
   		}
   		fclose (fp);
   	}

	// finde größtes s^2/t^alpha zum endzeitpunkt
   	int t_start = 10;
   	double temp_ende = 0.0;
   	for(int i_alpha = 1; i_alpha < 12; i_alpha ++){
   		double alpha_temp = 0.45 + 0.05 * i_alpha;
   		l = LENGTH_T-1;
   		if (temp_ende < (squares[l] / (pow(t[l], alpha_temp))) )
   		{
   			temp_ende = squares[l] / (pow(t[l], alpha_temp));
   		}
   	}
   	if (DIM>1)
   	{

   		char zplotsLABEL[30];
   		for (i = 0; i < n_zplots; i++)
   		{
   			sprintf(zplotsLABEL, "trajec/zplots%d.dat",i);
   			fp = fopen (zplotsLABEL, "w");
   			for(l = 0; l < LENGTH_T; l++){
   				for (j = 0; j < DIM ; j++){
   					fprintf(fp, "%lf  ", ((zplots[l][j +i*DIM])/latlength) );
   				}
   				fprintf(fp, "\n");
   			}
   			fclose (fp);
   		}
   		for (i = 0; i < n_zplots; i++)
   		{
   			sprintf(zplotsLABEL, "trajec/qplots%d.dat",i);
   			fp = fopen (zplotsLABEL, "w");
   			for(l = 0; l < LENGTH_T; l++){
   				for (j = 0; j < DIM ; j++){
   					fprintf(fp, "%lf  ", ((qplots[l][j +i*DIM])) );
   				}
   				fprintf(fp, "\n");
   			}
   			fclose (fp);
   		}
   		for (i = 0; i < n_zplots; i++)
   		{
   			sprintf(zplotsLABEL, "trajec/Zwanzig_Force%d.dat",i);
   			fp = fopen (zplotsLABEL, "w");
   			for(l = 0; l < LENGTH_T; l++){
   				for (j = 0; j < DIM ; j++){
   					fprintf(fp, "%lf  ", ((Zwanzig_Force[l][j +i*DIM])) );
   				}
   				fprintf(fp, "\n");
   			}
   			fclose (fp);
   		}
   		for (i = 0; i < n_zplots; i++)
   		{
   			sprintf(zplotsLABEL, "trajec/qlplots%d.dat",i);
   			fp = fopen (zplotsLABEL, "w");
   			for(l = 0; l < LENGTH_T; l++){
   				for (j = 0; j < DIM ; j++){
   					fprintf(fp, "%lf  ", ((qlplots[l][j +i*DIM])) );
   				}
   				fprintf(fp, "\n");
   			}
   			fclose (fp);
   		}
   	}
   	if (DIM==1)
   	{

   		char zplotsLABEL[30];
   		for (i = 0; i < n_zplots; i++)
   		{
   			sprintf(zplotsLABEL, "trajec/zplots%d.dat",i);
   			fp = fopen (zplotsLABEL, "w");
   			for(l = 0; l < LENGTH_T; l++){
   				fprintf(fp, " %lf  %lf \n" ,t[l]/TIME_ARRIVAL, zplots[l][i*DIM]/latlength);
   			}
   			fclose (fp);
   		}
   	}

   	fp = fopen ("ETOT.dat", "w");
   	for(l = 0; l < LENGTH_T; l++){
   		fprintf(fp, "%lf  %lf\n", t[l], ETOT[l]);
   	}
   	fclose (fp);

   	fp = fopen ("DENSITY.dat", "w");
   	for(m=0; m<80; m++)
   	{
   		for(n=0; n<80; n++)
   		{
   			fprintf(fp, "%1.2e ", P_density[m][n]);
   		}
   		fprintf(fp, "\n");
   	}
   	fclose (fp);

	gsl_rng_free (r); gsl_rng_free (RAND_GLOBAL);
	free(zplots[0]); free(zplots);
	free(qplots[0]); free(qplots);
	free(Zwanzig_Force[0]); free(Zwanzig_Force);
	free(qlplots[0]); free(qlplots);
	free (P_density[0]); free (P_density);
	free(t);
	free(bathp); free(bathq);
	free(squares_increase);
	free(short_correlation_x);
	free(total_momentum_x);
	free(total_momentum_y);
	free(short_correlation_y);
	free(long_correlation_x);
	free(px_correlation); free(py_correlation);
	free(px_correlation_late);
	free(px_correlation_late2);
	free(long_correlation_y);
	free(ETOT); free(EKIN); free(EBAD);
	free(PTOT); free(PTOTY);free(LTOT);
	free(squares);
	free(z[0]); free(z);
}


/* Copyright (c) 2011-2012, Jérémy Fix. All rights reserved. */

/* Redistribution and use in source and binary forms, with or without */
/* modification, are permitted provided that the following conditions are met: */

/* * Redistributions of source code must retain the above copyright notice, */
/* this list of conditions and the following disclaimer. */
/* * Redistributions in binary form must reproduce the above copyright notice, */
/* this list of conditions and the following disclaimer in the documentation */
/* and/or other materials provided with the distribution. */
/* * None of the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. */

/* THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND */
/* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED */
/* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE */
/* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE */
/* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL */
/* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR */
/* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER */
/* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, */
/* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE */
/* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */

#ifndef EKF_TYPES_H
#define EKF_TYPES_H

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

/**
 * @short Extended Kalman Filter in the case of additive noise, the notations follow Van Der Merwe, phD, p. 36
 */
namespace ekf
{

  class EvolutionNoise;

  /**
   * @short Structure holding the parameters of the statistical linearization
   *
   */
  typedef struct
  {

    /**
     * @short Evolution noise type
     */
    EvolutionNoise *evolution_noise;

    /**
     * @short Covariance of the observation noise
     */
    double observation_noise;

    /**
     * @short Prior estimate of the covariance matrix
     */
    double prior_pk;

    /**
     * @short Number of parameters to estimate
     */
    int n;

    /**
     * @short Dimension of the output
     */
    int no;

    /**
     * @short Is the observation gradient diagonal ? In that case, simplifications can be introduced
     */
    bool observation_gradient_is_diagonal;

  } ekf_param;

  typedef struct
  {
    /**
     * @short Current estimate of the state
     */
    gsl_vector * xk;

    /**
     * @short Predicted state
     */
    gsl_vector * xkm;

    /**
     * @short Variance-Covariance of the state
     */
    gsl_matrix * Pxk;

    /**
     * @short Jacobian of the evolution function
     */
    gsl_matrix *Fxk;

    /**
     * @short Variance covariance of the evolution noise
     */
    gsl_matrix * Rv;

    /**
     * @short Current observation
     */
    gsl_vector * yk;

    /**
     * @short Current innovations
     */
    gsl_vector * ino_yk;

    /**
     * @short Jacobian of the observation function
     */
    gsl_matrix *Hyk;

    /**
     * @short Variance covariance of the observation noise
     */
    gsl_matrix * Rn;

    /**
     * @short Kalman gain
     */
    gsl_matrix *Kk;

    /**
     * @short Temporary matrices
     */
    gsl_matrix * temp_n_n;
    gsl_matrix * temp_n_1;
    gsl_matrix * temp_no_no;
    gsl_matrix * temp_n_no;
    gsl_matrix * temp_2_n_n;
    gsl_vector * temp_no;

    /**
     * @short Optional parameters (this must be allocated and initialized from the user side!
     */
    gsl_vector * params;

  } ekf_state;

  /**
   * @short Mother class from which the evolution noises inherit
   */
  class EvolutionNoise
  {
  protected:
    double _initial_value;
  public:
  EvolutionNoise(double initial_value) : _initial_value(initial_value) {};
    void init(ekf_param &p, ekf_state &s)
    {
      gsl_matrix_set_identity(s.Rv);
      gsl_matrix_scale(s.Rv, _initial_value);
    }

    virtual void updateEvolutionNoise(ekf_param &p, ekf_state &s) = 0;
  };

  /**
   * @short Annealing type evolution noise
   */
  class EvolutionAnneal : public EvolutionNoise
  {
    double _decay, _lower_bound;
  public:
  EvolutionAnneal(double initial_value, double decay, double lower_bound) : EvolutionNoise(initial_value),
      _decay(decay),
      _lower_bound(lower_bound)
      { };

    void updateEvolutionNoise(ekf_param &p, ekf_state &s)
    {
      for(int i = 0 ; i < p.n ; ++i)
	{
	  for(int j = 0 ; j < p.n ; ++j)
	    {
	      if(i == j)
		gsl_matrix_set(s.Rv, i, i, ukf::math::max(_decay * gsl_matrix_get(s.Rv,i,i),_lower_bound));
	      else
		gsl_matrix_set(s.Rv, i, j, 0.0);
	    }
	}
    }
  };

  /**
   * @short Forgetting type evolution noise
   */
  class EvolutionRLS : public EvolutionNoise
  {
    double _decay;
  public:
  EvolutionRLS(double initial_value, double decay) : EvolutionNoise(initial_value), _decay(decay)
    {
      if(ukf::math::cmp_equal(_decay, 0.0))
	printf("Forgetting factor should not be null !!\n");
    };

    void updateEvolutionNoise(ekf_param &p, ekf_state &s)
    {
      gsl_matrix_memcpy(s.Rv, s.Pxk);
      gsl_matrix_scale(s.Rv, 1.0 / _decay - 1.0);
    }
  };

  /**
   * @short Robbins-Monro evolution noise
   */
  class EvolutionRobbinsMonro : public EvolutionNoise
  {
    double _alpha;
  public:
  EvolutionRobbinsMonro(double initial_value, double alpha) : EvolutionNoise(initial_value),
      _alpha(alpha)
      { };

    void updateEvolutionNoise(ekf_param &p, ekf_state &s)
    {
      // Compute Kk * ino_yk
      gsl_matrix_view mat_view = gsl_matrix_view_array(s.ino_yk->data, p.no, 1);
      gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, s.Kk, &mat_view.matrix, 0.0, s.temp_n_1);
      // Compute : Rv = (1 - alpha) Rv + alpha . (Kk * ino_yk) * (Kk * ino_yk)^T
      gsl_blas_dgemm(CblasNoTrans, CblasTrans, _alpha, s.temp_n_1, s.temp_n_1, (1.0 - _alpha),s.Rv);
    }
  };

}

#endif // EKF_TYPES_H


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_statistics_double.h>
#include <math.h>
#include <utilities.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_stack_code.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"
#include "tz_neurotrace.h"
#include "tz_stack_threshold.h"
#include "tz_iarray.h"
#include "tz_darray.h"

/*
 * trace_neuron - trace neuron from given seeds
 *
 * trace_neuron [!wtr] seed_file -Dsave_dir
 *   -r: write intermediate results
 *
 */
int main(int argc, char* argv[])
{
  static char *Spec[] = {
    "[!wtr] [-canvas <string>] [-mask <string>] [-res <string>] [-minr <int>]",
    "-minlen <double>",
    " <image:string> -S<string> -D<string>",
    NULL};

  Process_Arguments(argc, argv, Spec, 1);

  char *dir = Get_String_Arg("-D");

  char file_path[100];
  sprintf(file_path, "%s/%s", dir, Get_String_Arg("-S"));
  printf("%s\n", file_path);

  Geo3d_Scalar_Field *seed = Read_Geo3d_Scalar_Field(file_path);

  int idx;

  sprintf(file_path, "%s/%s.bn", dir, "max_r");
  double max_r;
  int tmp;
  if (fexist(file_path)) {
    darray_read2(file_path, &max_r, &tmp);
  } else {
    max_r = darray_max(seed->values, seed->size, &idx);
  }

  printf("%g\n", max_r);

  max_r *= 1.5;

  /*
  sprintf(file_path, "%s/%s", dir, "soma0.bn");
  if (!fexist(file_path)) {
    max_r *= 2.0;
  }
  */

  Set_Neuroseg_Max_Radius(max_r);

  Stack *signal = Read_Stack(Get_String_Arg("image"));

  dim_type dim[3];
  dim[0] = signal->width;
  dim[1] = signal->height;
  dim[2] = signal->depth;
  /*
  IMatrix *chord = Make_IMatrix(dim, 3);

  Stack *code = Make_Stack(GREY16,
			   signal->width, signal->height, signal->depth);
  */
  Rgb_Color color;
  Set_Color(&color, 255, 0, 0);

  Stack *canvas = NULL;

  char trace_file_path[100];
  sprintf(trace_file_path, "%s/%s", dir, Get_String_Arg("-canvas"));

  if (fexist(trace_file_path) == 1) {
    canvas = Read_Stack((char *) trace_file_path);
  } else {
    canvas = Copy_Stack(signal);
    Stretch_Stack_Value_Q(canvas, 0.999);
    Translate_Stack(canvas, COLOR, 1);
  }

  Stack *traced = NULL;

  char trace_mask_path[100];
  sprintf(trace_mask_path, "%s/%s", dir, Get_String_Arg("-mask"));

  if (fexist(trace_mask_path) == 1) {
    traced = Read_Stack((char *) trace_mask_path);
  } else {
    traced = Make_Stack(GREY, signal->width, signal->height, signal->depth);
    One_Stack(traced);
  }

  //Object_3d *obj = NULL;
  int seed_offset = -1;

  Neurochain *chain = NULL;

  double z_scale = 1.0;

  if (Is_Arg_Matched("-res")) {
    sprintf(file_path, "%s", Get_String_Arg("-res"));

    if (fexist(file_path)) {
      double res[3];
      int length;
      darray_read2(file_path, res, &length);
      if (res[0] != res[1]) {
	perror("Different X-Y resolutions.");
	TZ_ERROR(ERROR_DATA_VALUE);
      }
      z_scale = res[0] / res[2];
    }
  }

  //sprintf(file_path, "%s/%s", dir, Get_String_Arg("-M"));
  //Stack *stack = Read_Stack(file_path);

  tic();

  FILE *fp = NULL;
  char chain_file_path[100];
  char vrml_file_path[100];

  double min_chain_length = 25.0;

  if (Is_Arg_Matched("-minlen")) {
    min_chain_length = Get_Double_Arg("-minlen");
  }

  int *indices = iarray_malloc(seed->size);
  double *values = darray_malloc(seed->size);
  int i;

  Local_Neuroseg *locseg = (Local_Neuroseg *)
    malloc(seed->size * sizeof(Local_Neuroseg));

  int index = 0;
  for (i = 0; i < seed->size; i++) {
    printf("-----------------------------> seed: %d / %d\n", i, seed->size);
    indices[i] = i;
    index = i;
    int x = (int) seed->points[index][0];
    int y = (int) seed->points[index][1];
    int z = (int) seed->points[index][2];

    double width = seed->values[index];

    chain = New_Neurochain();

    seed_offset = Stack_Util_Offset(x, y, z, signal->width, signal->height,
				    signal->depth);

    if (width < 3.0) {
      width += 0.5;
    }
    Set_Neuroseg(&(locseg[i].seg), width, width, 12.0,
		 0.0, 0.0, 0.0);

    double cpos[3];
    cpos[0] = x;
    cpos[1] = y;
    cpos[2] = z;
    cpos[2] *= z_scale;

    Set_Neuroseg_Position(&(locseg[i]), cpos, NEUROSEG_CENTER);
    Stack_Fit_Score fs;
    fs.n = 1;
    fs.options[0] = 1;
    values[i] = Local_Neuroseg_Orientation_Search_C(&(locseg[i]), signal, z_scale, &fs);
  }

  darray_qsort(values, indices, seed->size);

  /*
  for (i = 0; i < seed->size; i++) {
    indices[i] = i;
  }
  darraycpy(values, seed->values, 0, seed->size);
  darray_qsort(values, indices, seed->size);
  */

  int counter = 0;

  //  for (i = seed->size - 1; i >= seed->size - 231; i--) {
  for (i = seed->size - 1; i >= 0; i--) {
    index = indices[i];

    printf("-----------------------------> seed: %d / %d\n", i, seed->size);

    sprintf(chain_file_path, "%s/chain%d.bn", dir, index);
    sprintf(vrml_file_path, "%s/chain%d.wrl", dir, index);

    if (fexist(chain_file_path) == 1) {
      chain = Read_Neurochain(chain_file_path);
      if (Neurochain_Geolen(chain) >= min_chain_length) {
	Write_Neurochain_Vrml(vrml_file_path, chain);
	Neurochain_Label(canvas, chain, z_scale);
	Neurochain_Erase_E(traced, chain, z_scale, 0,
			   Neurochain_Length(chain, FORWARD),
			   1.5, 0.0);
      }

      Free_Neurochain(chain);
      printf("chain exists\n");
      continue;
    }

    int x = (int) seed->points[index][0];
    int y = (int) seed->points[index][1];
    int z = (int) seed->points[index][2];

    if (*STACK_PIXEL_8(traced, x, y, z, 0) == 0) {
      printf("traced \n");
      continue;
    }

    double width = seed->values[index];

    if (width > max_r) {
      printf("too thick\n");
      continue;
    }

    if (Is_Arg_Matched("-minr")) {
      int max_level = (int) (width + 0.5);
      if (max_level <= Get_Int_Arg("-minr")) {
	printf("too thin\n");
	continue;
      }
    }
    /*
    seed_offset = Stack_Util_Offset(x, y, z, signal->width, signal->height,
				    signal->depth);
    */

    chain = New_Neurochain();
    /*
    Stack_Level_Code_Constraint(stack, code, chord->array, &seed_offset, 1,
				max_level + 1);

    Voxel_t v;
    v[0] = x;
    v[1] = y;
    v[2] = z;

    Stack *tmp_stack = Copy_Stack(stack);
    obj = Stack_Grow_Object_Constraint(tmp_stack, 1, v, chord, code,
				       max_level);
    Free_Stack(tmp_stack);

    Print_Object_3d_Info(obj);

    double vec[3];
    Object_3d_Orientation_Zscale(obj, vec, MAJOR_AXIS, z_scale);

    double theta, psi;
    Geo3d_Vector obj_vec;
    Set_Geo3d_Vector(&obj_vec, vec[0], vec[1], vec[2]);

    Geo3d_Vector_Orientation(&obj_vec, &theta, &psi);
    */

    /*
    if (width < 3.0) {
      width += 0.5;
    }
    Set_Neuroseg(&(chain->locseg.seg), width, width, 12.0,
		 0.0, 0.0, 0.0);

    double cpos[3];
    cpos[0] = x;
    cpos[1] = y;
    cpos[2] = z;
    cpos[2] *= z_scale;

    //Set_Neuroseg_Position(&(chain->locseg), cpos, NEUROSEG_BOTTOM);
    Set_Neuroseg_Position(&(chain->locseg), cpos, NEUROSEG_CENTER);
    Stack_Fit_Score fs;
    fs.n = 1;
    fs.options[0] = 1;
    Local_Neuroseg_Orientation_Search_C(&(chain->locseg), signal, z_scale,
					&fs);
    //fs.options[0] = 1;
    */

    Copy_Local_Neuroseg(&(chain->locseg), &(locseg[index]));
    Neurochain *chain_head = chain;

    if (Initialize_Tracing(signal, chain, NULL, z_scale) >= MIN_SCORE) {
      if ((Neuroseg_Hit_Traced(&(chain->locseg), traced, z_scale) == FALSE) &&
	  (chain->locseg.seg.r1 < max_r) &&
	  (chain->locseg.seg.r2 < max_r)) {
	//Initialize_Tracing(signal, chain, NULL, z_scale);
	chain = Trace_Neuron2(signal, chain, BOTH, traced, z_scale, 500);

	//Neurochain *chain_head = Neurochain_Head(chain);
	chain_head = Neurochain_Remove_Overlap_Segs(chain);
	chain_head = Neurochain_Remove_Turn_Ends(chain_head, 0.5);
	/*
	if (i == seed->size - 231) {
	  Print_Neurochain(chain_head);
	}
	*/

	fp = fopen(chain_file_path, "w");
	Neurochain_Fwrite(chain_head, fp);
	fclose(fp);
	if (Neurochain_Geolen(chain_head) >= min_chain_length) {
	  Write_Neurochain_Vrml(vrml_file_path, chain_head);

	  Neurochain_Erase_E(traced, chain_head, z_scale, 0,
			     Neurochain_Length(chain_head, FORWARD),
			     1.5, 0.0);
	  Neurochain_Label(canvas, chain_head, z_scale);

	  counter += Neurochain_Length(chain_head, FORWARD);
	  if (counter > 500) {
	    if (Is_Arg_Matched("-r")) {
	      Write_Stack((char *) trace_mask_path, traced);
	    }

	    if (Is_Arg_Matched("-r")) {
	    Write_Stack((char *) trace_file_path, canvas);
	    }

	    counter = 0;
	  }
	}
      }
    }

    Free_Neurochain(chain_head);

    //Kill_Object_3d(obj);
  }

  Write_Stack((char *) trace_file_path, canvas);
  if (Is_Arg_Matched("-r")) {
    Write_Stack((char *) trace_mask_path, traced);
  }

  Kill_Geo3d_Scalar_Field(seed);

  printf("Time passed: %lld\n", toc());

  return 0;
}


/*
  C code for the adiabatic approximation
*/
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_roots.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_integration.h>
#ifdef _OPENMP
#include <omp.h>
#endif
#define CHUNKSIZE 10
//Potentials
#include <galpy_potentials.h>
#include <actionAngle.h>
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
/*
  Structure Declarations
*/
struct JRAdiabaticArg{
  double ER;
  double Lz22;
  int nargs;
  struct potentialArg * actionAngleArgs;
};
struct JzAdiabaticArg{
  double Ez;
  double R;
  int nargs;
  struct potentialArg * actionAngleArgs;
};
/*
  Function Declarations
*/
void actionAngleAdiabatic_actions(int,double *,double *,double *,double *,
				 double *,int,int *,double *,double,
				 double *,double *,int *);
void calcJRAdiabatic(int,double *,double *,double *,double *,double *,
		     int,struct potentialArg *,int);
void calcJzAdiabatic(int,double *,double *,double *,double *,int,
		     struct potentialArg *,int);
void calcRapRperi(int,double *,double *,double *,double *,double *,
		  int,struct potentialArg *);
void calcZmax(int,double *,double *,double *,double *,int,
	      struct potentialArg *);
double JRAdiabaticIntegrandSquared(double,void *);
double JRAdiabaticIntegrand(double,void *);
double JzAdiabaticIntegrandSquared(double,void *);
double JzAdiabaticIntegrand(double,void *);
double evaluateVerticalPotentials(double, double,int, struct potentialArg *);
/*
  Actual functions, inlines first
*/
inline void calcEREzL(int ndata,
		      double *R,
		      double *vR,
		      double *vT,
		      double *z,
		      double *vz,
		      double *ER,
		      double *Ez,
		      double *Lz,
		      int nargs,
		      struct potentialArg * actionAngleArgs){
  int ii;
  UNUSED int chunk= CHUNKSIZE;
#pragma omp parallel for schedule(static,chunk) private(ii)
  for (ii=0; ii < ndata; ii++){
    *(ER+ii)= evaluatePotentials(*(R+ii),0.,
				 nargs,actionAngleArgs)
      + 0.5 * *(vR+ii) * *(vR+ii)
      + 0.5 * *(vT+ii) * *(vT+ii);
    *(Ez+ii)= evaluateVerticalPotentials(*(R+ii),*(z+ii),
					 nargs,actionAngleArgs)
      + 0.5 * *(vz+ii) * *(vz+ii);
    *(Lz+ii)= *(R+ii) * *(vT+ii);
  }
}
/*
  MAIN FUNCTIONS
 */
void actionAngleAdiabatic_actions(int ndata,
				  double *R,
				  double *vR,
				  double *vT,
				  double *z,
				  double *vz,
				  int npot,
				  int * pot_type,
				  double * pot_args,
				  double gamma,
				  double *jr,
				  double *jz,
				  int * err){
  int ii;
  //Set up the potentials
  struct potentialArg * actionAngleArgs= (struct potentialArg *) malloc ( npot * sizeof (struct potentialArg) );
  parse_actionAngleArgs(npot,actionAngleArgs,pot_type,pot_args);
  //ER, Ez, Lz
  double *ER= (double *) malloc ( ndata * sizeof(double) );
  double *Ez= (double *) malloc ( ndata * sizeof(double) );
  double *Lz= (double *) malloc ( ndata * sizeof(double) );
  calcEREzL(ndata,R,vR,vT,z,vz,ER,Ez,Lz,npot,actionAngleArgs);
  //Calculate peri and apocenters
  double *rperi= (double *) malloc ( ndata * sizeof(double) );
  double *rap= (double *) malloc ( ndata * sizeof(double) );
  double *zmax= (double *) malloc ( ndata * sizeof(double) );
  calcZmax(ndata,zmax,z,R,Ez,npot,actionAngleArgs);
  calcJzAdiabatic(ndata,jz,zmax,R,Ez,npot,actionAngleArgs,10);
  //Adjust planar effective potential for gamma
  UNUSED int chunk= CHUNKSIZE;
#pragma omp parallel for schedule(static,chunk) private(ii)
  for (ii=0; ii < ndata; ii++){
    *(Lz+ii)= fabs( *(Lz+ii) ) + gamma * *(jz+ii);
    *(ER+ii)+= 0.5 * *(Lz+ii) * *(Lz+ii) / *(R+ii) / *(R+ii)
      - 0.5 * *(vT+ii) * *(vT+ii);
  }
  calcRapRperi(ndata,rperi,rap,R,ER,Lz,npot,actionAngleArgs);
  calcJRAdiabatic(ndata,jr,rperi,rap,ER,Lz,npot,actionAngleArgs,10);
  for (ii=0; ii < npot; ii++) {
    if ( (actionAngleArgs+ii)->i2d )
      interp_2d_free((actionAngleArgs+ii)->i2d) ;
    if ((actionAngleArgs+ii)->accx )
      gsl_interp_accel_free ((actionAngleArgs+ii)->accx);
    if ((actionAngleArgs+ii)->accy )
      gsl_interp_accel_free ((actionAngleArgs+ii)->accy);
    free((actionAngleArgs+ii)->args);
  }
  free(actionAngleArgs);
  free(ER);
  free(Ez);
  free(Lz);
  free(rperi);
  free(rap);
  free(zmax);
}
void calcJRAdiabatic(int ndata,
		     double * jr,
		     double * rperi,
		     double * rap,
		     double * ER,
		     double * Lz,
		     int nargs,
		     struct potentialArg * actionAngleArgs,
		     int order){
  int ii, tid, nthreads;
#ifdef _OPENMP
  nthreads = omp_get_max_threads();
#else
  nthreads = 1;
#endif
  gsl_function * JRInt= (gsl_function *) malloc ( nthreads * sizeof(gsl_function) );
  struct JRAdiabaticArg * params= (struct JRAdiabaticArg *) malloc ( nthreads * sizeof (struct JRAdiabaticArg) );
  for (tid=0; tid < nthreads; tid++){
    (params+tid)->nargs= nargs;
    (params+tid)->actionAngleArgs= actionAngleArgs;
  }
  //Setup integrator
  gsl_integration_glfixed_table * T= gsl_integration_glfixed_table_alloc (order);
  UNUSED int chunk= CHUNKSIZE;
#pragma omp parallel for schedule(static,chunk)				\
  private(tid,ii)							\
  shared(jr,rperi,rap,JRInt,params,T,ER,Lz)
  for (ii=0; ii < ndata; ii++){
#ifdef _OPENMP
    tid= omp_get_thread_num();
#else
    tid = 0;
#endif
    if ( *(rperi+ii) == -9999.99 || *(rap+ii) == -9999.99 ){
      *(jr+ii)= 9999.99;
      continue;
    }
    if ( (*(rap+ii) - *(rperi+ii)) / *(rap+ii) < 0.000001 ){//circular
      *(jr+ii) = 0.;
      continue;
    }
    //Setup function
    (params+tid)->ER= *(ER+ii);
    (params+tid)->Lz22= 0.5 * *(Lz+ii) * *(Lz+ii);
    (JRInt+tid)->function = &JRAdiabaticIntegrand;
    (JRInt+tid)->params = params+tid;
    //Integrate
    *(jr+ii)= gsl_integration_glfixed (JRInt+tid,*(rperi+ii),*(rap+ii),T)
      * sqrt(2.) / M_PI;
  }
  free(JRInt);
  free(params);
  gsl_integration_glfixed_table_free ( T );
}
void calcJzAdiabatic(int ndata,
		     double * jz,
		     double * zmax,
		     double * R,
		     double * Ez,
		     int nargs,
		     struct potentialArg * actionAngleArgs,
		     int order){
  int ii, tid, nthreads;
#ifdef _OPENMP
  nthreads = omp_get_max_threads();
#else
  nthreads = 1;
#endif
  gsl_function * JzInt= (gsl_function *) malloc ( nthreads * sizeof(gsl_function) );
  struct JzAdiabaticArg * params= (struct JzAdiabaticArg *) malloc ( nthreads * sizeof (struct JzAdiabaticArg) );
  for (tid=0; tid < nthreads; tid++){
    (params+tid)->nargs= nargs;
    (params+tid)->actionAngleArgs= actionAngleArgs;
  }
  //Setup integrator
  gsl_integration_glfixed_table * T= gsl_integration_glfixed_table_alloc (order);
  UNUSED int chunk= CHUNKSIZE;
#pragma omp parallel for schedule(static,chunk)				\
  private(tid,ii)							\
  shared(jz,zmax,JzInt,params,T,Ez,R)
  for (ii=0; ii < ndata; ii++){
#ifdef _OPENMP
    tid= omp_get_thread_num();
#else
    tid = 0;
#endif
    if ( *(zmax+ii) == -9999.99 ){
      *(jz+ii)= 9999.99;
      continue;
    }
    if ( *(zmax+ii) < 0.000001 ){//circular
      *(jz+ii) = 0.;
      continue;
    }
    //Setup function
    (params+tid)->Ez= *(Ez+ii);
    (params+tid)->R= *(R+ii);
    (JzInt+tid)->function = &JzAdiabaticIntegrand;
    (JzInt+tid)->params = params+tid;
    //Integrate
    *(jz+ii)= gsl_integration_glfixed (JzInt+tid,0.,*(zmax+ii),T)
      * 2 * sqrt(2.) / M_PI;
  }
  free(JzInt);
  free(params);
  gsl_integration_glfixed_table_free ( T );
}
void calcRapRperi(int ndata,
		  double * rperi,
		  double * rap,
		  double * R,
		  double * ER,
		  double * Lz,
		  int nargs,
		  struct potentialArg * actionAngleArgs){
  int ii, tid, nthreads;
#ifdef _OPENMP
  nthreads = omp_get_max_threads();
#else
  nthreads = 1;
#endif
  double peps, meps;
  gsl_function * JRRoot= (gsl_function *) malloc ( nthreads * sizeof(gsl_function) );
  struct JRAdiabaticArg * params= (struct JRAdiabaticArg *) malloc ( nthreads * sizeof (struct JRAdiabaticArg) );
  //Setup solver
  int status;
  int iter, max_iter = 100;
  const gsl_root_fsolver_type *T;
  double R_lo, R_hi;
  struct pragmasolver *s= (struct pragmasolver *) malloc ( nthreads * sizeof (struct pragmasolver) );
  T = gsl_root_fsolver_brent;
  for (tid=0; tid < nthreads; tid++){
    (params+tid)->nargs= nargs;
    (params+tid)->actionAngleArgs= actionAngleArgs;
    (s+tid)->s= gsl_root_fsolver_alloc (T);
  }
  UNUSED int chunk= CHUNKSIZE;
  gsl_set_error_handler_off();
#pragma omp parallel for schedule(static,chunk)				\
  private(tid,ii,iter,status,R_lo,R_hi,meps,peps)			\
  shared(rperi,rap,JRRoot,params,s,R,ER,Lz,max_iter)
  for (ii=0; ii < ndata; ii++){
#ifdef _OPENMP
    tid= omp_get_thread_num();
#else
    tid = 0;
#endif
    //Setup function
    (params+tid)->ER= *(ER+ii);
    (params+tid)->Lz22= 0.5 * *(Lz+ii) * *(Lz+ii);
    (JRRoot+tid)->params = params+tid;
    (JRRoot+tid)->function = &JRAdiabaticIntegrandSquared;
    //Find starting points for minimum
    if ( fabs(GSL_FN_EVAL(JRRoot+tid,*(R+ii))) < 0.0000001){ //we are at rap or rperi
      peps= GSL_FN_EVAL(JRRoot+tid,*(R+ii)+0.0000001);
      meps= GSL_FN_EVAL(JRRoot+tid,*(R+ii)-0.0000001);
      if ( fabs(peps) < 0.00000001 && fabs(meps) < 0.00000001 && peps*meps >= 0.) {//circular
	*(rperi+ii) = *(R+ii);
	*(rap+ii) = *(R+ii);
      }
      else if ( peps < 0. && meps > 0. ) {//umax
	*(rap+ii)= *(R+ii);
	R_lo= 0.9 * (*(R+ii) - 0.0000001);
	R_hi= *(R+ii) - 0.00000001;
	while ( GSL_FN_EVAL(JRRoot+tid,R_lo) >= 0. && R_lo > 0.000000001){
	  R_hi= R_lo; //this makes sure that brent evaluates using previous
	  R_lo*= 0.9;
	}
	//Find root
	status = gsl_root_fsolver_set ((s+tid)->s, JRRoot+tid, R_lo, R_hi);
	if (status == GSL_EINVAL) {
	  *(rperi+ii) = 0.;//Assume zero if below 0.000000001
	  continue;
	}
	iter= 0;
	do
	  {
	    iter++;
	    status = gsl_root_fsolver_iterate ((s+tid)->s);
	    R_lo = gsl_root_fsolver_x_lower ((s+tid)->s);
	    R_hi = gsl_root_fsolver_x_upper ((s+tid)->s);
	    status = gsl_root_test_interval (R_lo, R_hi,
					     9.9999999999999998e-13,
					     4.4408920985006262e-16);
	  }
	while (status == GSL_CONTINUE && iter < max_iter);
	// LCOV_EXCL_START
	if (status == GSL_EINVAL) {//Shouldn't ever get here
	  *(rperi+ii) = -9999.99;
	  *(rap+ii) = -9999.99;
	  continue;
	}
	// LCOV_EXCL_STOP
	*(rperi+ii) = gsl_root_fsolver_root ((s+tid)->s);
      }
      else if ( peps > 0. && meps < 0. ){//umin
	*(rperi+ii)= *(R+ii);
	R_lo= *(R+ii) + 0.0000001;
	R_hi= 1.1 * (*(R+ii) + 0.0000001);
	while ( GSL_FN_EVAL(JRRoot+tid,R_hi) >= 0. && R_hi < 37.5) {
	  R_lo= R_hi; //this makes sure that brent evaluates using previous
	  R_hi*= 1.1;
	}
	//Find root
	status = gsl_root_fsolver_set ((s+tid)->s, JRRoot+tid, R_lo, R_hi);
	if (status == GSL_EINVAL) {
	  *(rperi+ii) = -9999.99;
	  *(rap+ii) = -9999.99;
	  continue;
	}
	iter= 0;
	do
	  {
	    iter++;
	    status = gsl_root_fsolver_iterate ((s+tid)->s);
	    R_lo = gsl_root_fsolver_x_lower ((s+tid)->s);
	    R_hi = gsl_root_fsolver_x_upper ((s+tid)->s);
	    status = gsl_root_test_interval (R_lo, R_hi,
					     9.9999999999999998e-13,
					     4.4408920985006262e-16);
	  }
	while (status == GSL_CONTINUE && iter < max_iter);
	// LCOV_EXCL_START
	if (status == GSL_EINVAL) {//Shouldn't ever get here
	  *(rperi+ii) = -9999.99;
	  *(rap+ii) = -9999.99;
	  continue;
	}
	// LCOV_EXCL_STOP
	*(rap+ii) = gsl_root_fsolver_root ((s+tid)->s);
      }
    }
    else {
      R_lo= 0.9 * *(R+ii);
      R_hi= *(R+ii);
      while ( GSL_FN_EVAL(JRRoot+tid,R_lo) >= 0. && R_lo > 0.000000001){
	R_hi= R_lo; //this makes sure that brent evaluates using previous
	R_lo*= 0.9;
      }
      R_hi= (R_lo < 0.9 * *(R+ii)) ? R_lo / 0.9 / 0.9: *(R+ii);
      //Find root
      status = gsl_root_fsolver_set ((s+tid)->s, JRRoot+tid, R_lo, R_hi);
      if (status == GSL_EINVAL) {
	*(rperi+ii) = 0.;//Assume zero if below 0.000000001
      } else {
	iter= 0;
	do
	  {
	    iter++;
	    status = gsl_root_fsolver_iterate ((s+tid)->s);
	    R_lo = gsl_root_fsolver_x_lower ((s+tid)->s);
	    R_hi = gsl_root_fsolver_x_upper ((s+tid)->s);
	    status = gsl_root_test_interval (R_lo, R_hi,
					     9.9999999999999998e-13,
					     4.4408920985006262e-16);
	  }
	while (status == GSL_CONTINUE && iter < max_iter);
	// LCOV_EXCL_START
	if (status == GSL_EINVAL) {//Shouldn't ever get here
	  *(rperi+ii) = -9999.99;
	  *(rap+ii) = -9999.99;
	  continue;
	}
	// LCOV_EXCL_STOP
	*(rperi+ii) = gsl_root_fsolver_root ((s+tid)->s);
      }
      //Find starting points for maximum
      R_lo= *(R+ii);
      R_hi= 1.1 * *(R+ii);
      while ( GSL_FN_EVAL(JRRoot+tid,R_hi) > 0. && R_hi < 37.5) {
	R_lo= R_hi; //this makes sure that brent evaluates using previous
	R_hi*= 1.1;
      }
      R_lo= (R_hi > 1.1 * *(R+ii)) ? R_hi / 1.1 / 1.1: *(R+ii);
      //Find root
      status = gsl_root_fsolver_set ((s+tid)->s, JRRoot+tid, R_lo, R_hi);
      if (status == GSL_EINVAL) {
	*(rperi+ii) = -9999.99;
	*(rap+ii) = -9999.99;
	continue;
      }
      iter= 0;
      do
	{
	  iter++;
	  status = gsl_root_fsolver_iterate ((s+tid)->s);
	  R_lo = gsl_root_fsolver_x_lower ((s+tid)->s);
	  R_hi = gsl_root_fsolver_x_upper ((s+tid)->s);
	  status = gsl_root_test_interval (R_lo, R_hi,
					   9.9999999999999998e-13,
					   4.4408920985006262e-16);
	}
      while (status == GSL_CONTINUE && iter < max_iter);
      // LCOV_EXCL_START
      if (status == GSL_EINVAL) {//Shouldn't ever get here
	*(rperi+ii) = -9999.99;
	*(rap+ii) = -9999.99;
	continue;
      }
      // LCOV_EXCL_STOP
      *(rap+ii) = gsl_root_fsolver_root ((s+tid)->s);
    }
  }
  gsl_set_error_handler (NULL);
  for (tid=0; tid < nthreads; tid++)
    gsl_root_fsolver_free( (s+tid)->s);
  free(s);
  free(JRRoot);
  free(params);
}
void calcZmax(int ndata,
	      double * zmax,
	      double * z,
	      double * R,
	      double * Ez,
	      int nargs,
	      struct potentialArg * actionAngleArgs){
  int ii, tid, nthreads;
#ifdef _OPENMP
  nthreads = omp_get_max_threads();
#else
  nthreads = 1;
#endif
  gsl_function * JzRoot= (gsl_function *) malloc ( nthreads * sizeof(gsl_function) );
  struct JzAdiabaticArg * params= (struct JzAdiabaticArg *) malloc ( nthreads * sizeof (struct JzAdiabaticArg) );
  //Setup solver
  int status;
  int iter, max_iter = 100;
  const gsl_root_fsolver_type *T;
  double z_lo, z_hi;
  struct pragmasolver *s= (struct pragmasolver *) malloc ( nthreads * sizeof (struct pragmasolver) );
  T = gsl_root_fsolver_brent;
  for (tid=0; tid < nthreads; tid++){
    (params+tid)->nargs= nargs;
    (params+tid)->actionAngleArgs= actionAngleArgs;
    (s+tid)->s= gsl_root_fsolver_alloc (T);
  }
  UNUSED int chunk= CHUNKSIZE;
  gsl_set_error_handler_off();
#pragma omp parallel for schedule(static,chunk)				\
  private(tid,ii,iter,status,z_lo,z_hi)				\
  shared(zmax,JzRoot,params,s,z,Ez,R,max_iter)
  for (ii=0; ii < ndata; ii++){
#ifdef _OPENMP
    tid= omp_get_thread_num();
#else
    tid = 0;
#endif
    //Setup function
    (params+tid)->Ez= *(Ez+ii);
    (params+tid)->R= *(R+ii);
    (JzRoot+tid)->function = &JzAdiabaticIntegrandSquared;
    (JzRoot+tid)->params = params+tid;
    //Find starting points for minimum
    if ( fabs(GSL_FN_EVAL(JzRoot+tid,*(z+ii))) < 0.0000001){ //we are at zmax
      *(zmax+ii)= fabs( *(z+ii) );
    }
    else {
      z_lo= fabs ( *(z+ii) );
      z_hi= ( *(z+ii) == 0. ) ? 0.1: 1.1 * fabs( *(z+ii) );
      while ( GSL_FN_EVAL(JzRoot+tid,z_hi) >= 0. && z_hi < 37.5) {
	z_lo= z_hi; //this makes sure that brent evaluates using previous
	z_hi*= 1.1;
      }
      //Find root
      status = gsl_root_fsolver_set ((s+tid)->s, JzRoot+tid, z_lo, z_hi);
      if (status == GSL_EINVAL) {
	*(zmax+ii) = -9999.99;
	continue;
      }
      iter= 0;
      do
	{
	  iter++;
	  status = gsl_root_fsolver_iterate ((s+tid)->s);
	  z_lo = gsl_root_fsolver_x_lower ((s+tid)->s);
	  z_hi = gsl_root_fsolver_x_upper ((s+tid)->s);
	  status = gsl_root_test_interval (z_lo, z_hi,
					   9.9999999999999998e-13,
					   4.4408920985006262e-16);
	}
      while (status == GSL_CONTINUE && iter < max_iter);
      // LCOV_EXCL_START
      if (status == GSL_EINVAL) {//Shouldn't ever get here
	*(zmax+ii) = -9999.99;
	continue;
      }
      // LCOV_EXCL_STOP
      *(zmax+ii) = gsl_root_fsolver_root ((s+tid)->s);
    }
  }
  gsl_set_error_handler (NULL);
  for (tid=0; tid < nthreads; tid++)
    gsl_root_fsolver_free( (s+tid)->s);
  free(s);
  free(JzRoot);
  free(params);
}
double JRAdiabaticIntegrand(double R,
			   void * p){
  return sqrt(JRAdiabaticIntegrandSquared(R,p));
}
double JRAdiabaticIntegrandSquared(double R,
				  void * p){
  struct JRAdiabaticArg * params= (struct JRAdiabaticArg *) p;
  return params->ER - evaluatePotentials(R,0.,params->nargs,params->actionAngleArgs) - params->Lz22 / R / R;
}
double JzAdiabaticIntegrand(double z,
			    void * p){
  return sqrt(JzAdiabaticIntegrandSquared(z,p));
}
double JzAdiabaticIntegrandSquared(double z,
				   void * p){
  struct JzAdiabaticArg * params= (struct JzAdiabaticArg *) p;
  return params->Ez - evaluateVerticalPotentials(params->R,z,
						 params->nargs,
						 params->actionAngleArgs);
}
double evaluateVerticalPotentials(double R, double z,
				  int nargs,
				  struct potentialArg * actionAngleArgs){
  return evaluatePotentials(R,z,nargs,actionAngleArgs)
    -evaluatePotentials(R,0.,nargs,actionAngleArgs);
}


//#include "readArray.h"
//#include "tensor_train.h"

#include "matrix.h"
#include "tensor.h"
#include "sketch.h"
#include "tt.h"
#include "VTime.h"
#include "readArray_MPI.h"
#include "PSTT.h"
#include "my_io.h"
#include "SSTT.h"

#include <omp.h>

#include <stdio.h>
#include <cblas.h>
#include <lapacke.h>
#include <math.h>
#include <string.h>
#include <gperftools/heap-profiler.h>

#ifndef HEAD
#define HEAD ((int) 0)
#endif

#ifndef FILE_BUF
#define FILE_BUF ((int) 256)
#endif

VTime_gathered* get_TT(tensor_train* tt, MPI_tensor* ten, int mid, int solvetype, int *heap_profile, char* path, long* total_alloc_ptr);

void MyHeapProfilerStart(int heap_profile, char* path, char* prefix, int rank)
{
    if (heap_profile){
        char* profile_path = (char*) calloc(FILE_BUF, sizeof(char));
        snprintf(profile_path, FILE_BUF, "%s%s%d", path, prefix, rank);
        HeapProfilerStart(profile_path);
    }
}

void MyHeapProfilerStop(int* heap_profile, char* path, char* prefix, int rank, long* total_alloc_ptr)
{
    if (*heap_profile){
        HeapProfilerDump("");
        HeapProfilerStop();

        char file_name[256];
        sprintf(file_name, "%s%s%d.000%d.heap", path, prefix, rank, *heap_profile);
        FILE *f = fopen(file_name, "r");
        while(f){
            *heap_profile = *heap_profile + 1;
            sprintf(file_name, "%s%s%d.000%d.heap", path, prefix, rank, *heap_profile);
            f = fopen(file_name, "r");
        }
        sprintf(file_name, "%s%s%d.000%d.heap", path, prefix, rank, *heap_profile - 1);
        f = fopen(file_name, "r");

        char total_alloc_str[256];
        char c = fgetc(f);
        while( c != '['){
            c = fgetc(f);
        }
        while( c != ':'){
            c = fgetc(f);
        }
        c = fgetc(f);
        int ii = 0;
        while( c != ']'){
            if (c != ' '){
                total_alloc_str[ii] = c;
                ii = ii+1;
            }
            c = fgetc(f);
        }
        total_alloc_str[ii] = '\0';

        int info = sscanf(total_alloc_str, "%ld", total_alloc_ptr);
//        printf("total_alloc_str = %s, total_alloc = %ld\n", total_alloc_str, *total_alloc_ptr);

//        long MB = 1;
//        for (int ii = 0; ii < 20; ++ii){
//            MB = MB * 2;
//        }
//        printf("size = %f MB\n ", (double) *total_alloc_ptr / MB);
    }
}

long* gather_heap_data(int heap_profile, long* total_allocs, int size, int rank, MPI_Comm comm)
{

    int head = 0;
    long* total_allocs_matrix = NULL;

    if (heap_profile){
        if (rank == head){
            total_allocs_matrix = (long*) calloc((size + 1)*3, sizeof(long));
        }
        MPI_Gather(total_allocs, 3, MPI_LONG, total_allocs_matrix + 3, 3, MPI_LONG, head, comm);

        for (int ii = 0; ii < size; ++ii){
            for (int jj = 0; jj < 3; ++jj){
                total_allocs[jj + ii*3];
            }
        }

        if (rank == head){
            for (int jj = 1; jj < size + 1; ++jj){
                for (int ii = 0; ii < 3; ++ii){
                    total_allocs_matrix[ii] += total_allocs_matrix[ii + jj*3];
                }
            }
            long MB = 1;
            for (int ii = 0; ii < 20; ++ii){
                MB = MB*2;
            }
            printf("TOTAL ALGORITHM SIZE = %f\n", (double) total_allocs_matrix[2]/MB);
        }
    }

    return total_allocs_matrix;
}

//Input:
// argv[1] = solvetype
// argv[2] = tensortype
// -d      = dimension
// -n      = sizes
// -r      = ranks
// -p      = path
// -e      = nps
// -m      = mid (optional)
// -h      = heap_profile
// -M      = M for gaussian bumps (only for tensortype == 2)
// -g      = gamma for gaussian bumps (only for tensortype == 2)
int main (int argc, char *argv[])
{
    MPI_Init(&argc, &argv);

    int world_rank;
    int world_size;
    MPI_Comm comm_world = MPI_COMM_WORLD;
    MPI_Comm_rank(comm_world, &world_rank);
    MPI_Comm_size(comm_world, &world_size);
    if (world_rank == HEAD){ printf("Starting MPI_timing\n");}

    int solvetype = atoi(argv[1]);
    int tensortype = atoi(argv[2]);

    int* dvec = NULL;
    int d = 0;

    int* midvec = NULL;
    int mid = -1;

    int* heap_profile_vec = NULL;
    int heap_profile = 0;

    int* nps = NULL;

    int* n = NULL;
    int* r = NULL;

    void (*f_ten)(double* restrict, int*, int*, const void*);
    void* parameters = NULL;

    long* total_allocs = NULL;

    double acc = 1e-10;

    // Taken chars: 'd', 'm', 'n', 'r', 'p', 'e', 'g'

    dvec = get_int_star_input('d', 1, argc, argv);
    if (dvec == NULL){
        printf("No input d for Hilbert Tensor\n");
        MPI_Finalize();
        return 0;
    }
    else{
        d = dvec[0];
    }

    midvec = get_int_star_input('m', 1, argc, argv);
    if (midvec != NULL){
        mid = midvec[0];
    }

    // 'e' is for 'eeeeeeeeh-nps'
    nps = get_int_star_input('e', d, argc, argv);
    if (nps == NULL){
        printf("No input nps\n");
        MPI_Finalize();
        return 0;
    }

    n = get_int_star_input('n', d, argc, argv);
    if (n == NULL){
        printf("No input n\n");
        MPI_Finalize();
        return 0;
    }

    r = get_int_star_input('r', d+1, argc, argv);
    if (r == NULL){
        printf("No input r\n");
        MPI_Finalize();
        return 0;
    }

    char* path = get_string_input ('p', argc, argv);
    if (path == NULL){
        printf("Must input path with -p\n");
        MPI_Finalize();
        return 0;
    }

    heap_profile_vec = get_int_star_input('h', 1, argc, argv);
    if (heap_profile_vec != NULL){
        heap_profile = *heap_profile_vec;
        total_allocs = (long*) calloc(3, sizeof(long));
    }

    // Get the tensor pointer (I have _no_ idea how to make a function for this anymore)
    if (tensortype == 1){
        f_ten = &f_hilbert;
        parameters = p_hilbert_init(d, n);
    }
    else if (tensortype == 2){
        int* Mvec = NULL;
        int M;
        Mvec = get_int_star_input('M', 1, argc, argv);
        if (Mvec != NULL){
            M = Mvec[0];
        }
        else{
            printf("Must input M with -M when tensortype == 2\n");
            MPI_Finalize();
            return 0;
        }

        double* gammavec = NULL;
        double gamma;
        gammavec = get_double_star_input('g', 1, argc, argv);
        if (gammavec != NULL){
            gamma = gammavec[0];
        }
        else{
            printf("Must input gamma with -g when tensortype == 2\n");
            MPI_Finalize();
            return 0;
        }

        int seed = 168234591; // Arbitrary number
        parameters = unit_random_p_gaussian_bumps_init(d, n, M, gamma, seed);
        f_ten = &f_gaussian_bumps;

        free(gammavec); gammavec = NULL;
        free(Mvec);     Mvec = NULL;
    }

    MyHeapProfilerStart(heap_profile, path, ".tt", world_rank);
    tensor_train* tt = TT_init_rank(d, n, r);
    MyHeapProfilerStop(&heap_profile, path, ".tt", world_rank, total_allocs);

    MyHeapProfilerStart(heap_profile, path, ".ten", world_rank);
    MPI_tensor* ten = MPI_tensor_init(d, n, nps, MPI_COMM_WORLD, f_ten, parameters);
    MyHeapProfilerStop(&heap_profile, path, ".ten", world_rank, total_allocs + 1);

    VTime_gathered* tms = get_TT(tt, ten, mid, solvetype, &heap_profile, path, total_allocs + 2);
    if (world_rank == HEAD){ printf("Got the tensor train\n");}

    if (solvetype == 1){
        ten = MPI_tensor_init(d, n, nps, MPI_COMM_WORLD, f_ten, parameters);
    }
    double err = tt_error(tt, ten);
    if (world_rank == HEAD){
        printf("ERROR = %e\n", err);
        printf("TIME = %f\n", tms->t[0]);
    }
    MPI_tensor_free(ten);

    long* total_allocs_matrix = gather_heap_data(heap_profile, total_allocs, world_size, world_rank, comm_world);
    if (world_rank == HEAD){
        my_tt_write(tt, tms, err, nps, tensortype, solvetype, path, total_allocs_matrix);
    }

    if (tms != NULL){;
        VTime_gathered_free(tms); tms = NULL;
    }
    TT_free(tt); tt = NULL;

    if (tensortype == 1){
        p_hilbert_free(parameters); parameters = NULL;
    }
    else if (tensortype == 2){
        p_gaussian_bumps_free(parameters); parameters = NULL;
    }

    free(r);      r = NULL;
    free(n);      n = NULL;
    free(nps);    nps = NULL;
    free(midvec); midvec = NULL;
    free(dvec);   dvec = NULL;
    free(total_allocs); total_allocs = NULL;
    free(heap_profile_vec); heap_profile_vec = NULL;

    MPI_Finalize();

    return 0;
}

/* Key for argv[1]
 *  1 : streaming_SSTT
 *  2 : streaming_PSTT2
 *  3 : streaming_PSTT2_onepass
*/
VTime_gathered* get_TT(tensor_train* tt, MPI_tensor* ten, int mid, int solvetype, int *heap_profile, char* path, long* total_alloc_ptr)
{
    VTime* tm;
    int world_rank = ten->rank;
    MPI_Comm comm = ten->comm;

    if (world_rank==HEAD){ printf("Getting the tensor train ("); }
    if (solvetype == 1){
        if (world_rank==HEAD){ printf("SSTT)\n"); }
        MyHeapProfilerStart(*heap_profile, path, ".SSTT", world_rank);
        tm = SSTT(tt, ten);
        MyHeapProfilerStop(heap_profile, path, ".SSTT", world_rank, total_alloc_ptr);
    }
    if (solvetype == 2){
        if (world_rank==HEAD){ printf("PSTT2)\n"); }
        MyHeapProfilerStart(*heap_profile, path, ".PSTT", world_rank);
        tm = PSTT2(tt, ten, mid);
        MyHeapProfilerStop(heap_profile, path, ".PSTT", world_rank, total_alloc_ptr);
    }
    if (solvetype == 3){
        if (world_rank==HEAD){ printf("PSTT2_onepass)\n"); }
        MyHeapProfilerStart(*heap_profile, path, ".PSTT_onepass", world_rank);
        tm = PSTT2_onepass(tt, ten, mid);
        MyHeapProfilerStop(heap_profile, path, ".PSTT_onepass", world_rank, total_alloc_ptr);
    }
    if (world_rank==HEAD){ printf("Got tensor train\n");}

    VTime_gathered* tms = VTime_gather(tm, comm);

    if (tm != NULL){
        VTime_free(tm); tm = NULL;
    }

    return tms;
}


/* Copyright (c) 2011-2017, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#pragma once

namespace amgx
{
template <class Matrix>  typename Matrix::value_type estimate_largest_eigen_value(Matrix &A);
}

#include <norm.h>
#include <multiply.h>
#include <blas.h>

namespace amgx
{

template <class Matrix>
typename Matrix::value_type estimate_largest_eigen_value(Matrix &A)
{
    typedef typename Matrix::TConfig TConfig;
    typedef typename Matrix::value_type ValueTypeA;
    typedef typename TConfig::VecPrec ValueTypeB;
    typedef Vector<TConfig> VVector;
    VVector x(A.get_num_rows()), y(A.get_num_rows());
    fill(x, 1);

    for (int i = 0; i < 20; i++)
    {
        ValueTypeB Lmax = get_norm(A, x, LMAX);
        scal(x, ValueTypeB(1) / Lmax);
        multiply(A, x, y);
        x.swap(y);
    }

    ValueTypeB retval = get_norm(A, x, L2) / get_norm(A, y, L2);
    return retval;
}

} // namespace amgx


#include <nlopt.h>

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

/****************************************************************************/
/* test function from M. J. Box, "A new method of constrained optimization
   and a comparison with other methods," Computer J. 8 (1), 42-52 (1965) */

int testfuncs_verbose = 1;
int testfuncs_counter = 0;

static double testfuncs_status(int n, const double *x, double f)
{
    ++testfuncs_counter;
    if (testfuncs_verbose) {
        int i;
        printf("f_%d (%g", testfuncs_counter, x[0]);
        for (i = 1; i < n; ++i)
            printf(", %g", x[i]);
        printf(") = %g\n", f);
    }
    return f;
}

#define RETURN(f) return testfuncs_status(n, x, f);

static const double k1 = -145421.402, k2 = 2931.1506, k3 = -40.427932, k4 = 5106.192,
    k5 = 15711.36, k6 = -161622.577, k7 = 4176.15328, k8 = 2.8260078,
    k9 = 9200.476, k10 = 13160.295, k11 = -21686.9194, k12 = 123.56928,
    k13 = -21.1188894, k14 = 706.834, k15 = 2898.573, k16 = 28298.388,
    k17 = 60.81096, k18 = 31.242116, k19 = 329.574, k20 = -2882.082,
    k21 = 74095.3845, k22 = -306.262544, k23 = 16.243649, k24 = -3094.252,
    k25 = -5566.2628, k26 = -26237, k27 = 99, k28 = -0.42, k29 = 1300, k30 = 2100, k31 = 925548.252, k32 = -61968.8432, k33 = 23.3088196, k34 = -27096.648, k35 = -50843.766;

static double box(int n, const double *x, double *grad, void *data)
{
    double x1 = x[0], x2 = x[1], x3 = x[2], x4 = x[3], x5 = x[4];
    double b, x6, y1, y2, y3, y4, u;
    const double a0 = 9, a1 = 15, a2 = 50, a3 = 9.583, a4 = 20, a5 = 15, a6 = 6, a7 = 0.75;
    b = x2 + 0.01 * x3;
    x6 = (k1 + k2 * x2 + k3 * x3 + k4 * x4 + k5 * x5) * x1;
    y1 = k6 + k7 * x2 + k8 * x3 + k9 * x4 + k10 * x5;
    y2 = k11 + k12 * x2 + k13 * x3 + k14 * x4 + k15 * x5;
    y3 = k16 + k17 * x2 + k18 * x3 + k19 * x4 + k20 * x5;
    y4 = k21 + k22 * x2 + k23 * x3 + k24 * x4 + k25 * x5;
    u = a2 * y1 + a3 * y2 + a4 * y3 + a5 * y4 + 7840 * a6 - 100000 * a0 - 50800 * b * a7 + k31 + k32 * x2 + k33 * x3 + k34 * x4 + k35 * x5;
    if (grad) {
        int i;
        grad[0] = u + a1 * (k1 + k2 * x2 + k3 * x3 + k4 * x4 + k5 * x5);
        grad[1] = x1 * (a2 * k7 + a3 * k12 + a4 * k17 + a5 * k22 - 50800 * a7 + k32) + a1 * (k2 * x1);
        grad[2] = x1 * (a2 * k8 + a3 * k13 + a4 * k18 + a5 * k23 - 50800 * a7 * 0.01) + a1 * x1 * k3;
        grad[3] = x1 * (a2 * k9 + a3 * k14 + a4 * k19 + a5 * k24) + a1 * x1 * k4;
        grad[4] = x1 * (a2 * k10 + a3 * k15 + a4 * k20 + a5 * k25) + a1 * x1 * k5;
        for (i = 0; i < 5; ++i)
            grad[i] = -grad[i];
    }
    RETURN(-(u * x1 - 24345 + a1 * x6));
}

static double box_constraint(int n, const double *x, double *grad, void *data)
{
    int which_constraint = *((int *) data);
    double x1 = x[0], x2 = x[1], x3 = x[2], x4 = x[3], x5 = x[4];
    double x6, y1, y2, y3, x7, x8;
    int i;

    x6 = (k1 + k2 * x2 + k3 * x3 + k4 * x4 + k5 * x5) * x1;
    y1 = k6 + k7 * x2 + k8 * x3 + k9 * x4 + k10 * x5;
    y2 = k11 + k12 * x2 + k13 * x3 + k14 * x4 + k15 * x5;
    y3 = k16 + k17 * x2 + k18 * x3 + k19 * x4 + k20 * x5;
    x7 = (y1 + y2 + y3) * x1;
    x8 = (k26 + k27 * x2 + k28 * x3 + k29 * x4 + k30 * x5) * x1 + x6 + x7;

    if (grad) {
        grad[0] = grad[1] = grad[2] = grad[3] = grad[4] = 0;

        if (which_constraint != 2 && which_constraint != -2) {
            /* grad x6 */
            grad[0] += (k1 + k2 * x2 + k3 * x3 + k4 * x4 + k5 * x5);
            grad[1] += x1 * k2;
            grad[2] += x1 * k3;
            grad[3] += x1 * k4;
            grad[4] += x1 * k5;
        }
        if (which_constraint != 1 && which_constraint != -1) {
            /* grad x7 */
            grad[0] += (y1 + y2 + y3);
            grad[1] += x1 * (k7 + k12 + k17);
            grad[2] += x1 * (k8 + k13 + k18);
            grad[3] += x1 * (k9 + k14 + k19);
            grad[4] += x1 * (k10 + k15 + k20);
        }
        if (which_constraint == 3 || which_constraint == -3) {
            /* grad (x8 - x6 - x7) */
            grad[0] += k26 + k27 * x2 + k28 * x3 + k29 * x4 + k30 * x5;
            grad[1] += x1 * k27;
            grad[2] += x1 * k28;
            grad[3] += x1 * k29;
            grad[4] += x1 * k30;
        }
    }

    if (which_constraint == -1) {
        if (grad)
            for (i = 0; i < 5; ++i)
                grad[i] = -grad[i];
        return -x6;
    } else if (which_constraint == 1) {
        return x6 - 294000;
    } else if (which_constraint == -2) {
        if (grad)
            for (i = 0; i < 5; ++i)
                grad[i] = -grad[i];
        return -x7;
    } else if (which_constraint == 2) {
        return x7 - 294000;
    } else if (which_constraint == -3) {
        if (grad)
            for (i = 0; i < 5; ++i)
                grad[i] = -grad[i];
        return -x8;
    } else if (which_constraint == 3) {
        return x8 - 277200;
    }
    return 0;
}

int main(void)
{
    const double box_lb[5] = { 0, 1.2, 20, 9, 6.5 };
    const double box_ub[5] = { HUGE_VAL, 2.4, 60, 9.3, 7.0 };
    const double box_xmin[5] = { 4.53743, 2.4, 60, 9.3, 7.0 };

    int cdata[6] = { -1, 1, -2, 2, -3, 3 };
    double x[5] = { 2.52, 2, 37.5, 9.25, 6.8 };
    double minf;

    nlopt_minimize_constrained(NLOPT_LN_COBYLA, 5, box, NULL, 6, box_constraint, cdata, sizeof(int), box_lb, box_ub, x, &minf, -HUGE_VAL, 1e-10, 0, 0, NULL, 0, 0);
    printf("found f(%g,%g,%g,%g,%g) = %0.8f after %d iters\n", x[0], x[1], x[2], x[3], x[4], minf, testfuncs_counter);
    return 0;
}