#include "gadget.h"    //All the required standard header files are in here
#include "optinfo.h"
#include "mathfunc.h"
#include "doublevector.h"
#include "intvector.h"
#include "errorhandler.h"
#include "ecosystem.h"
#include "global.h"
#include "pso.h"

#ifdef _OPENMP
#include "omp.h"
#endif

extern Ecosystem* EcoSystem;
#ifdef _OPENMP
extern Ecosystem** EcoSystems;
#endif


//==============================================================
//// calulate swarm size based on dimensionality
int OptInfoPso::pso_calc_swarm_size(int dim) {
    int size = 10. + 2. * sqrt(dim);
        return (size > PSO_MAX_SIZE ? PSO_MAX_SIZE : size);
     }

void OptInfoPso::OptimiseLikelihood() {

  handle.logMessage(LOGINFO, "\nStarting PSO optimisation algorithm\n");
  double oldf, newf, steplength, tmp;
  int    i, offset,rchange,rcheck,rnumber;
  int nvars = EcoSystem->numOptVariables();
  DoubleVector x(nvars);
  DoubleVector trialx(nvars);
  DoubleVector bestx(nvars);
  struct Best {double bestf; int index;};
  struct Best best;
  double bestf;
  DoubleVector lowerb(nvars);
  DoubleVector upperb(nvars);
  DoubleVector init(nvars);
  //DoubleVector initialstep(nvars, rho);
  IntVector param(nvars, 0);
  IntVector lbound(nvars, 0);
  IntVector rbounds(nvars, 0);
  IntVector trapped(nvars, 0);

/**
 * Internal variables for PSO
 */
  DoubleMatrix pos(size,nvars,0.0); // position matrix
  DoubleMatrix vel(size,nvars,0.0); // velocity matrix
  DoubleMatrix pos_b(size,nvars,0.0); // best position matrix
  DoubleVector fit(size); // particle fitness vector
  DoubleVector fit_b(size); // best fitness vector
  DoubleMatrix pos_nb(size,nvars,0.0); // what is the best informed position for each particle
  IntMatrix comm(size,size,0); // communications:who informs who rows : those who inform cols : those who are informed 
  int improved; // whether solution->error was improved during the last iteration
  int d, step;
  double a, b;
  double rho1, rho2; //random numbers (coefficients)
  double w; //current omega
  Inform_fun inform_fun;
  Calc_inertia_fun calc_inertia_fun; // inertia weight update function
  int psosteps=0;

  if (this->seed!=0)
     srand(this->seed);

  handle.logMessage(LOGINFO,"Starting PSO with particles ",size,"\n");
  switch (nhood_strategy)
	{
	case PSO_NHOOD_GLOBAL :
	    // comm matrix not used
            handle.logMessage(LOGINFO,"PSO will use global communication\n");
	    inform_fun = &OptInfoPso::inform_global;
	    break;
	case PSO_NHOOD_RING :
            handle.logMessage(LOGINFO,"PSO will use ring communication\n");
	    init_comm_ring(comm);
	    inform_fun = &OptInfoPso::inform_ring;
	    break;
	case PSO_NHOOD_RANDOM :
            handle.logMessage(LOGINFO,"PSO will use random communication\n");
	    init_comm_random(comm);
	    inform_fun = &OptInfoPso::inform_random;
	    break;
	}
  // SELECT APPROPRIATE INERTIA WEIGHT UPDATE FUNCTION
  switch (w_strategy)
  {
	case PSO_W_CONST : 
             handle.logMessage(LOGINFO,"PSO will use constant inertia\n");
	     calc_inertia_fun = &OptInfoPso::calc_inertia_const; 
	     break; 
	case PSO_W_LIN_DEC :
             handle.logMessage(LOGINFO,"PSO will use linear decreasing inertia\n");
	    calc_inertia_fun = &OptInfoPso::calc_inertia_lin_dec;
	    break;
  }
  //
  //initialize omega using standard value
  //
  w = PSO_INERTIA;
  handle.logMessage(LOGINFO,"PSO initial inertia ",w,"\n");
  handle.logMessage(LOGINFO,"PSO c1",c1,"\n");
  handle.logMessage(LOGINFO,"PSO c2",c2,"\n");
  handle.logMessage(LOGINFO,"PSO goal",goal,"\n");

  EcoSystem->scaleVariables();
#ifdef _OPENMP
  int numThr = omp_get_max_threads ( );
#pragma omp parallel for
  for (i=0;i<numThr;i++)
      EcoSystems[i]->scaleVariables();
#endif
  EcoSystem->getOptScaledValues(x);
  EcoSystem->getOptLowerBounds(lowerb);
  EcoSystem->getOptUpperBounds(upperb);
  EcoSystem->getOptInitialValues(init);

  for (i = 0; i < nvars; i++) {
    // Scaling the bounds, because the parameters are scaled
    lowerb[i] = lowerb[i] / init[i];
    upperb[i] = upperb[i] / init[i];
    if (lowerb[i] > upperb[i]) {
      tmp = lowerb[i];
      lowerb[i] = upperb[i];
      upperb[i] = tmp;
    }
    bestx[i] = x[i];
    trialx[i] = x[i];
    param[i] = i;
  }

  bestf = EcoSystem->SimulateAndUpdate(trialx);
  if (bestf != bestf) { //check for NaN
    handle.logMessage(LOGINFO, "Error starting PSO optimisation with f(x) = infinity");
    converge = -1;
    iters = 1;
    return;
  }

  offset = EcoSystem->getFuncEval();  //number of function evaluations done before loop
  
  iters = 0;

  //
  // Initialize the particles
  //
  i=0;
  handle.logMessage(LOGINFO,"Initialising PSO ",size,"particles\n");
  //
  // Initialize the fist particle with the best known position
  //
  for (d=0; d<nvars; d++) {
	    // generate two numbers within the specified range and around trialx
	    a = trialx[d] ;
	    // initialize position
	    pos[i][d] = a;
	    pos_b[i][d] = a;
	    pos_nb[i][d] = a;
	    a = lowerb[d] + (upperb[d] - lowerb[d]) * (rand()*1.0/RAND_MAX);
	    b = lowerb[d] + (upperb[d] - lowerb[d]) * (rand()*1.0/RAND_MAX);
	    vel[i][d] = (a-b) / 2.;
  }
  fit[i] = EcoSystem->SimulateAndUpdate(pos[i]);
  fit_b[i] = fit[i]; // this is also the personal best
  if (fit[i] < bestf) {
	    bestf = fit[i];
            best.bestf=bestf;
	    bestx=pos[i];
            best.index=i;
  }
//
// Now other particles with a random position in the solution space
//
best.bestf=bestf;
best.index=0;
handle.logMessage(LOGINFO,"before Best.bestf=",best.bestf);
#pragma omp declare reduction(bestomp: struct Best  : omp_out = omp_in.bestf < omp_out.bestf ? omp_in : omp_out)
#pragma omp parallel for default(shared) private(i,d,a,b) reduction(bestomp:best) firstprivate(bestf)
  for (i=1; i<size; i++) {
	// for each dimension
      ostringstream spos,svel;
      spos << "Initial position of particle " << i<<" [" ;
      svel << "Initial velocity of particle " << i<<" [" ;
	
	for (d=0; d<nvars; d++) {
	    // generate two numbers within the specified range and around trialx
	    a = lowerb[d] + (upperb[d] - lowerb[d]) * (rand()*1.0/RAND_MAX);
	    b = lowerb[d] + (upperb[d] - lowerb[d]) * (rand()*1.0/RAND_MAX);
	    // initialize position
	    pos[i][d] = a;
	    // best position is trialx 
	    pos_b[i][d] = a;
	    pos_nb[i][d] = a;
	    // initialize velocity
	    vel[i][d] = (a-b) / 2.;
            spos << pos[i][d] << " ";
            svel << vel[i][d] << " ";
	}
#ifdef _OPENMP
	fit[i] = EcoSystems[omp_get_thread_num ()]->SimulateAndUpdate(pos[i]);
	best.bestf = fit[i];
#else
	fit[i] = EcoSystem->SimulateAndUpdate(pos[i]);
#endif
	fit_b[i] = fit[i]; // this is also the personal best
	// update gbest??
//#pragma omp critical (best1)
	if (fit[i] < bestf) {
	    // update best fitness
 		handle.logMessage(LOGINFO,"iInside Best.bestf=",best.bestf);
	    // copy particle pos to gbest vector
	    //memmove((void *)&bestx, (void *)&pos[i],
	//	    sizeof(double) * nvars);
	    best.index=i;
            bestf=best.bestf;
	    //bestx=pos[i];
	}
	spos<<" ]";
	svel<<" ]";
        handle.logMessage(LOGDEBUG, spos.str().c_str());
        handle.logMessage(LOGDEBUG, svel.str().c_str());

	
  }
//
// Compare the reduction with the best value
//

  if (best.bestf < bestf){
  	bestx=pos[best.index];
	bestf=best.bestf;
  }
   

 // 
 // RUN ALGORITHM
 //
  psosteps=0;
  while (1) {
    
    handle.logMessage(LOGINFO, "PSO optimisation after", psosteps*size,"\n");
    if (isZero(bestf)) {
      iters = EcoSystem->getFuncEval() - offset;
      handle.logMessage(LOGINFO, "Error in PSO optimisation after", iters, "function evaluations, f(x) = 0");
      converge = -1;
      return;
    }
    /* if too many function evaluations occur, terminate the algorithm */

    iters = EcoSystem->getFuncEval() - offset;
    if (psosteps > psoiter) {
      handle.logMessage(LOGINFO, "\nStopping PSO optimisation algorithm\n");
      handle.logMessage(LOGINFO, "The optimisation stopped after", iters, "function evaluations");
      handle.logMessage(LOGINFO, "The optimisation stopped because the maximum number of PSO steps was reached");
      handle.logMessage(LOGINFO, "was reached and NOT because an optimum was found for this run");

      bestf = EcoSystem->SimulateAndUpdate(bestx);
      for (i = 0; i < nvars; i++)
        bestx[i] = bestx[i] * init[i];
      handle.logMessage(LOGINFO, "\nPSO finished with a likelihood score of", bestf);
      EcoSystem->storeVariables(bestf, bestx);
      return;
    }

    if (w_strategy)
      w = (this->*calc_inertia_fun)(psosteps);
	// check optimization goal
    if (bestf <= goal) {
      handle.logMessage(LOGINFO, "\nStopping PSO optimisation algorithm\n");
      handle.logMessage(LOGINFO, "Goal achieved!!! @ step ", psosteps);
      bestf = EcoSystem->SimulateAndUpdate(bestx);
      handle.logMessage(LOGINFO, "\nPSO finished with a likelihood score of", bestf);
      break;
      }

      // update pos_nb matrix (find best of neighborhood for all particles)
    (this->*inform_fun)(comm, pos_nb, pos_b, fit_b, bestx, improved);

	// the value of improved was just used; reset it
    improved = 0;

	// update all particles
    best.bestf=bestf;
#pragma omp parallel for default(shared) private(i,d,rho1,rho2)  reduction(bestomp:best) firstprivate(bestf)
    for (i=0; i<size; i++) {
	  // for each dimension
      handle.logMessage(LOGDEBUG, "In parallel",omp_get_thread_num(),"\n");
      ostringstream spos,svel;
      spos << "New position of particle " << i<<" [" ;
      svel << "New velocity of particle " << i<<" [" ;
      for (d=0; d<nvars; d++) {
	// calculate stochastic coefficients
	rho1 = c1 * (rand()*1.0)/RAND_MAX;
	rho2 = c2 * (rand()*1.0)/RAND_MAX;
	// update velocity
	vel[i][d] = w * vel[i][d] +	\
	    rho1 * (pos_b[i][d] - pos[i][d]) +	\
	    rho2 * (pos_nb[i][d] - pos[i][d]);
		// update position
	pos[i][d] += vel[i][d];
        spos << pos[i][d] <<" ";
        svel << vel[i][d] <<" ";
		// clamp position within bounds?
	if (clamp_pos) {
	    if (pos[i][d] < lowerb[d]) {
		pos[i][d] = lowerb[d];
		vel[i][d] = 0;
	    } else if (pos[i][d] > upperb[d]) {
		pos[i][d] = upperb[d];
		vel[i][d] = 0;
	    }
	} else {
		    // enforce periodic boundary conditions
	    if (pos[i][d] < lowerb[d]) {

		pos[i][d] = upperb[d] - fmod(lowerb[d] - pos[i][d],
					  upperb[d] - lowerb[d]);
		vel[i][d] = 0;

	    } else if (pos[i][d] > upperb[d]) {

		pos[i][d] = lowerb[d] + fmod(pos[i][d] - upperb[d],
					  upperb[d] - lowerb[d]);
		vel[i][d] = 0;

			
	    }
	  }
		    
         } //end for d
       spos<< "]" << endl;
       svel<< "]" << endl;
       handle.logMessage(LOGDEBUG,spos.str().c_str());
       handle.logMessage(LOGDEBUG,svel.str().c_str());
	    
	    // update particle fitness
       fit[i] = EcoSystems[omp_get_thread_num()]->SimulateAndUpdate(pos[i]);
	    // update personal best position?
       best.bestf = fit[i];
       if (fit[i] < fit_b[i]) {
	   fit_b[i] = fit[i];
		// copy contents of pos[i] to pos_b[i]
		//memmove((void *)&pos_b[i], (void *)&pos[i], 
	//		sizeof(double) * nvars);
	   pos_b[i]=pos[i];
       }
	    // update gbest??
        handle.logMessage(LOGDEBUG, "Particle ",i);
        handle.logMessage(LOGDEBUG, "The likelihood score is", bestf, "at the point");
        handle.logMessage(LOGDEBUG, "Its likelihood score is", fit[i], "at the point");
//#pragma omp critical (best2)
        if (fit[i] < bestf) {
	   improved = 1;
		// update best fitness
           best.index=i;
           bestf=fit[i];
		// copy particle pos to gbest vector
		//memmove((void *)&bestx, (void *)&pos[i],
	//		sizeof(double) * nvars);
	//   for (d = 0; d < nvars; d++)
        //     bestx[d] = pos[i][d] * init[d];
        //   iters = EcoSystem->getFuncEval() - offset;
      	//   handle.logMessage(LOGINFO, "\nNew optimum found after", iters, "function evaluations");
      	//   handle.logMessage(LOGINFO, "The likelihood score is", best.bestf, "at the point");
      	//   EcoSystem->storeVariables(best.bestf, bestx);
      	//   EcoSystem->writeBestValues();
      	//   bestx=pos[i];

        }
      } // end i
      if (best.bestf < bestf){
	   bestf=best.bestf;
	   bestx=pos[best.index];
	   for (d = 0; d < nvars; d++)
             bestx[d] = bestx[d] * init[d];
           iters = EcoSystem->getFuncEval() - offset; //To change. It is incorrect when is parallel.
      	   handle.logMessage(LOGINFO, "\nNew optimum found after", iters, "function evaluations");
      	   handle.logMessage(LOGINFO, "The likelihood score is", bestf, "at the point");
      	   EcoSystem->storeVariables(bestf, bestx);
      	   EcoSystem->writeBestValues();
      	   bestx=pos[best.index];
	}
	
      psosteps++;

    } // while (newf < bestf)

    iters = EcoSystem->getFuncEval() - offset;
    handle.logMessage(LOGINFO, "Existing PSO after ", iters, "function evaluations ...");
   for (d = 0; d < nvars; d++)
         bestx[d] = bestx[d] * init[d];
    EcoSystem->storeVariables(bestf, bestx);
    EcoSystem->writeBestValues();


}

void OptInfoPso::OptimiseLikelihoodREP() {
      OptimiseLikelihood();

}

void OptInfoPso::OptimiseLikelihoodOMP() {
      OptimiseLikelihood();
  }

/**
 * \brief calculate constant  inertia weight equal to w_max
 */
double OptInfoPso::calc_inertia_const(int step) {
        return w_max;
}
/**
 * \brief calculate linearly decreasing inertia weight
 */
double OptInfoPso::calc_inertia_lin_dec(int step) {

    int dec_stage = 3 * psoiter / 4;
    if (step <= dec_stage)
	return w_min + (w_max - w_min) *	\
	    (dec_stage - step) / dec_stage;
    else
	return w_min;
}



/**
 * \brief NEIGHBORHOOD (COMM) MATRIX STRATEGIES. global neighborhood
 */
void OptInfoPso::inform_global(IntMatrix& comm, DoubleMatrix& pos_nb,
		   DoubleMatrix& pos_b, DoubleVector& fit_b,
		   DoubleVector& gbest,int improved)
{

    int i;
    // all particles have the same attractor (gbest)
    // copy the contents of gbest to pos_nb
    for (i=0; i<size; i++)
	//memmove((void *)&pos_nb[i][0], &gbest[0],
//		sizeof(double) * gbest.Size());
	pos_nb[i]=gbest;
}


/**
 * \brief  general inform function :: according to the connectivity matrix COMM, it copies the best position (from pos_b) of the informers of each particle to the pos_nb matrix
 */
void OptInfoPso::inform(IntMatrix& comm, DoubleMatrix& pos_nb, DoubleMatrix& pos_b, DoubleVector& fit_b, int improved) 
{
    int i, j;
    int b_n; // best neighbor in terms of fitness

    // for each particle
    for (j=0; j<size; j++) {
	b_n = j; // self is best
	// who is the best informer??
	for (i=0; i<size; i++)
	    // the i^th particle informs the j^th particle
	    if (comm[i][j] && fit_b[i] < fit_b[b_n])
		// found a better informer for j^th particle
		b_n = i;
	// copy pos_b of b_n^th particle to pos_nb[j]
	//memmove((void *)&pos_nb[j][0],
	//	(void *)&pos_b[b_n][0],
	//	sizeof(double) * pos_nb.Ncol(j));
	pos_nb[j]=pos_b[b_n];    
}
}    

/**
 * \brief topology initialization :: this is a static (i.e. fixed) topology
 */
void OptInfoPso::init_comm_ring(IntMatrix& comm) {
    int i;
    // reset array
    comm.setToZero(); 

    // choose informers
    for (i=0; i<size; i++) {
	// set diagonal to 1
	comm[i][i] = 1;
	if (i==0) {
	    // look right
	    comm[i][1] = 1;
	    // look left
	    comm[i][size-1] = 1;
	} else if (i==size-1) {
	    // look right
	    comm[i][0] = 1;
	    // look left
	    comm[i][i-1] = 1;
	} else {
	    // look right
	    comm[i][i+1] = 1;
	    // look left
	    comm[i][i-1] = 1;
	}

    }
    // Print Matrix
    handle.logMessage(LOGDEBUG, "\nPSO ring communication\n");
    for (i=0;i<size;i++){
       string message="[";
       for (int j=0;j<size;j++)
       {
          std::ostringstream ss;
	  ss<< comm[i][j];
	  message=message+" "+ss.str();
	}
       message=message+"]\n";
       handle.logMessage(LOGDEBUG, message.c_str());
    }

}




void OptInfoPso::inform_ring(IntMatrix& comm, DoubleMatrix& pos_nb,
		 DoubleMatrix& pos_b, DoubleVector& fit_b,
		 DoubleVector& gbest, int improved)
{

    // update pos_nb matrix
    inform(comm, pos_nb, pos_b, fit_b, improved);
    
}

/**
 * \brief random neighborhood topology
 */
void OptInfoPso::init_comm_random(IntMatrix& comm) {

    int i, j, k;
    // reset array
    comm.setToZero();

    // choose informers
    for (i=0; i<size; i++) {
	// each particle informs itself
	comm[i][i] = 1;
	// choose kappa (on average) informers for each particle
	for (k=0; k<nhood_size; k++) {
	    // generate a random index
	    //j = gsl_rng_uniform_int(settings->rng, settings->size);
	    j=rand()*size/RAND_MAX;
	    // particle i informs particle j
	    comm[i] [j] = 1;
	}
    }
}



void OptInfoPso::inform_random(IntMatrix& comm, DoubleMatrix& pos_nb,
		   DoubleMatrix& pos_b, DoubleVector& fit_b,
		   DoubleVector& gbest, int improved)
{


    // regenerate connectivity??
    if (!improved)
	init_comm_random(comm);
    inform(comm, pos_nb, pos_b, fit_b, improved);

}