#ifndef __PSO_H__
#define __PSO_H__

#include "image2d.h"
#include "mutualinformation.h"
#include <stdlib.h>
#include "transcache.h"
#include "transformation.h"
#include <algorithm>
#include "clustermanager.h"
#include "timer.h"
#include "psoparticle.h"
#include "psotracker.h"
#include "common.h"

//#define NUMTHREADS 8

/// pso thread data structure
struct psothreadd
{
  std::string alnname;                          ///< name attached to alignment for identification
  pthread_t tid;
  int tnum;                                     ///< thread number [0,totalthreads)
  int totalthreads;                             ///< total number of worker threads

  std::vector< image2d* > imstacka;             ///< image stack a
  std::vector< image2d* > imstackb;             ///< image stack b
  pthread_barrier_t *slicesync;                 ///< synchronize the particle calculation

  pthread_barrier_t *gsync;                     ///< global sync

  std::vector< psoparticle* > particles;        ///< list particles for each slice in the stack
  std::vector< psoparticle > *kcenters;         ///< kmeans cluster centers with best score
  std::vector< psoparticle > *besttrans;        ///< list of best transformations per slice

  std::vector< transcache* > tcache;            ///< transformation score cache for each slice

  std::vector< psoparticle > *tkcenters;        ///< array of temp thread kcenters

  image2d *mapa;
  image2d *mapb;

  int parnumlow;                                ///< responsible for these particles
  int parnumhigh;                               ///< responsible for these particles

  int nslices;

  int npar;                                     ///< particles per cluster
  int *clusflag;                                ///< cluster flag signal

  int *numclusters;

  double *clusterentropy;

  int *largestcluster;                          ///< pointer to largest cluster number

  mutualinformation *mi;                        ///< object for calculating mut. info.

  box searchspace;

  int *exit;                                    ///< exit flag

  psotracker *statdb;

  int clusterreset;

  std::vector< psoparticle* > parlist;          ///< vectors of pointers to every particle
  std::vector< transformation > *newtrans;      ///< array of vectors containing list of new transformations for each slice
  std::vector< psotrack_cachestat > *cachestat; ///< pointer to cache stat vector
};

double myrand();
void *psothread_start( void *d );
int pthread_psorun( psothreadd *d );
int pthread_psocluster( psothreadd *d );
int pthread_psoupdate( psoparticle *p, std::vector< psoparticle> *kc,
                       std::vector< psoparticle > *bs, psoparticle *lc,
                       double lgc, psothreadd *d );
int pthread_psobest( psothreadd *d, psoparticle *p, 
                     std::vector< psoparticle > *b,
                     int s, int np );
int pthread_psoupdate( psothreadd *d );
int pthread_psoparticle( psothreadd *d );
int pthread_psotcache( psothreadd *d, transcache *c, psoparticle *pl );
int pthread_psocluster( psothreadd *d, std::vector< psoparticle > *kc, int nclus );
int pthread_clusterclosest( psoparticle &p, std::vector< psoparticle > *kc, int nclus );
int pthread_psoclustermove( psothreadd *d, std::vector<psoparticle> &kc,
                            int numclus);
int pthread_savestats( psothreadd *d, int iteration );
int pthread_psoclusterbests( psothreadd *d, std::vector<psoparticle> &kc,
                             int numclusters );
int pthread_printpsostat( psothreadd *d, int lc, int iter, int v  );

///Particle Swarm Optimization for image alignment
class pso
{
 public:
  pso();
  ~pso();
  int del();

  /// sets the two images
  int init( std::vector< image2d* > ima, std::vector< image2d* > imb,
            image2d &ma, image2d &mb);

  /// produce the alignment
  /**
   * @param npar number of particles per slice
   * @param lowx search space low x limiter
   * @param lowy search space low y limiter
   * @param highx search space high x limiter
   * @param highx search space high y limiter
   * @param cm cluster manager depercated because of multithreaded function
   */
  int optimize( int npar, double lowx, double lowy,
                double highx, double highy, clustermanager *cm = NULL);

  /// returns the best transformation
  /**
   * @param tx will become the x transformation 
   * @param ty will become the y transformation
   * @param score the score of the transformation
   */
  int getresult( double &tx, double &ty, double &score );

  /// returns the results
  int getresult( std::vector< transformation > &trans );

  /// multithreaded alignment
  int mpoptimize( int npar, double lowx, double lowy,
                  double highx, double highy, std::string name = "",
                  psotracker *pt = NULL);

 private:

  /// returns random value between [0-1]
  double rand();

  int threadedoptmizer( );

  int initparticles( int npar );

  int calcparticle( psoparticle *p, int s );
  int updateparticle( psoparticle *p );

  int closestcluster( psoparticle &p );

  /// single threaded kmeans should look into using kmeans class
  int kmeans();

  int findbestcluster();

  int distributework();
  int distributeparticle( int p, int s );

  double bestx;    ///< best transformation x
  double besty;    ///< best transformation y
  double bscore;   ///< best score

  std::vector< psoparticle > kcenters;   ///< kmeans cluster centers with best score
  std::vector< psoparticle > besttrans;  ///< list of best transformations per slice

  std::vector< image2d* > imstacka;      ///< image stack a
  std::vector< image2d* > imstackb;      ///< image stack b

  image2d mapa;                          ///< map containing valid pixels in stack a
  image2d mapb;                          ///< map containing valid pixels in stack b

  std::vector< psoparticle* > particles; ///< list particles for each slice in the stack

  std::vector< transcache* > tcache;     ///< transformation score cache for each slice

  int parslice;                          ///< number of particles pre image

  int largestcluster; ///< largest cluster id
  int goodclusters;   ///< flag to signal if clusters are "good"
  int numberclusters; ///< number of clusters to use
  double clusterentropy; ///< entropy of the clusters

  box searchspace;       ///< search space

  clustermanager *clusman;

  int clusteritr;     ///< pso iterations done with current number of clusters

  psotracker *psostat;
};

#endif
