#include "pso.h"

bool particle_cmp( psoparticle a, psoparticle b );

bool particle_cmp( psoparticle a, psoparticle b )
{
  double as, bs;
  as = pow( a.vx, 2) + pow(a.vy, 2);
  bs = pow( b.vx, 2) + pow(b.vy, 2);
  if( as < bs )
    return true;
  return false;
}

double myrand()
{
  double s;
  s =((double) (random() % 10000 ))/10000.0;
  if( s > 1 || s < 0 )
    std::cout << "ERROR rand\n";
  return s;
}

void *psothread_start( void *d )
{
  psothreadd *z;
  z = ( psothreadd *) d;
  pthread_psorun( z );
  return NULL;
}

int pthread_psobest( psothreadd *d, psoparticle *p,
                     std::vector< psoparticle > *b,
                     int s, int np )
{
  (*b)[s].bscore = 0;
  for( int i = 0; i < np; i++ )
    {
      if( p[i].cscore > (*b)[s].bscore )
        {
          (*b)[s].bscore = p[i].cscore;
          (*b)[s].bx = round( p[i].x );
          (*b)[s].by = round( p[i].y );
        }
    }
  d = NULL;
  return 0;
}

int pthread_psorun( psothreadd *d )
{
  int flag = 1;
  std::cout << "thread " << d->tnum << " running\n";
  timer mytimer, myitertimer;
  int iter;
  iter = 0;
  d->clusterreset = 1;
  *(d->largestcluster) = 0;
  while( flag == 1 )
    {
      pthread_barrier_wait( d->slicesync );
      if( *(d->exit) == 1 )
        {
          pthread_psocluster( d );
          flag = 0;
        }
      else
        {
	  if( d->tnum == 0 )
	    {
              myitertimer.reset();
	      mytimer.reset();
              myitertimer.start();
	      mytimer.start();
	    }
          // calculate particles
          pthread_psoparticle( d );
          pthread_barrier_wait( d->slicesync );
          // update transform cache
          for( int i = d->tnum; i < d->nslices; i = i + d->totalthreads )
            pthread_psotcache( d, d->tcache[i], d->particles[i] );
          pthread_barrier_wait( d->slicesync );

	  if( d->tnum == 0 )
	    {
	      mytimer.stop();
	      std::cout << "mi finished in: " << mytimer.getseconds() << "s ("
			<< mytimer.getseconds()/( (double)(d->nslices*d->npar)) << "\n";
	    }
          // calculate clustering
          pthread_psocluster( d );
          // find the best score per slice
          for( int i = d->tnum; i < d->nslices; i = i + d->totalthreads )
            pthread_psobest( d, d->particles[i],
                             d->besttrans, i, d->npar );
          // update speeds
          double v = 0;
          int df;
          df = 0;
          // find the largest cluster and calculate cluster entropy
          if( d->tnum == 0 )
            {
              int lc;
              lc = 0;
              for( int i = 0; i < *(d->numclusters); i++ )
                {
                  if( (*(d->kcenters))[lc].cscore <
                      (*(d->kcenters))[i].cscore )
                    lc = i;
                  v = v + (*(d->kcenters))[i].cscore;
                }
              if( (int)v != d->npar * d->nslices )
                {
                  std::cout << "ERROR mismatch cluster counts\n"
                            << v << " expecting " 
                            << d->npar * d->nslices << "\n";
                  exit(0);
                }
              *(d->clusterentropy) = 0;
              *(d->largestcluster) = lc;
              
              std::cout << "-" << d->alnname << "-\n";
                            
              double p;
              for( int i = 0; i < *(d->numclusters); i++ )
                {
                  p = (*(d->kcenters))[i].cscore / v;
                  if( p < 0.025 )
                    df = 1;
                  if( p > 0 )
                    *(d->clusterentropy) = *(d->clusterentropy) + p*log(p);
                }

              //print the status 
              pthread_printpsostat( d, lc, iter, v );
            }
          
          pthread_barrier_wait(d->slicesync);
          pthread_psoupdate( d );
          pthread_barrier_wait( d->slicesync );
          if( d->tnum == 0 )
            {
              double avgspeed;
              avgspeed = 0;
              double c;
              c = 0;
              std::cout << "tcache sizes: ";
              for( int i = 0; i < d->nslices; i++ )
                {
                  std::cout << (d->tcache[i])->size() << " ";
                  for( int j = 0; j < d->npar; j++)
                    {
                      avgspeed = avgspeed + pow( d->particles[i][j].vx, 2 )
                        + pow( d->particles[i][j].vy, 2 );
                      c = c + 1;
                    }
                }
              std::cout << "\n";
              avgspeed = avgspeed / c;
              myitertimer.stop();
              std::cout << iter << "/" << 500 << " average speed: " << avgspeed 
                        << " finished in " << myitertimer.getseconds() << "s\n";

              if( avgspeed < 100 || iter > 500)
                *(d->exit) = 1;
              else if( df == 1 )
                {
                  *(d->numclusters) = *(d->numclusters) - 1;
                  d->clusterreset = 1;
                }
              iter++;
            }
        }
      pthread_barrier_wait( d->slicesync );
    }
  pthread_barrier_wait( d->gsync );
  return 0;
}

int pthread_printpsostat( psothreadd *d, int lc, int iter, int v )
{

  std::vector< psotrack_cluster > kmeans;
  kmeans.clear();
  double p;
  std::cout << "cluster  size  center  best score location  best score\n";

  for( int i = 0; i < *(d->numclusters); i++ )
    {
      p = (*(d->kcenters))[i].cscore / v;
#ifdef __DEBUG_PSO_ISNAN__
      if( isnan((*(d->kcenters))[i].x ) )
        {
          std::cout << "cluster center is nan\n";
          exit(0);
        }
#endif
      psotrack_cluster ptc;
      ptc.x = (*(d->kcenters))[i].x;
      ptc.y = (*(d->kcenters))[i].y;
      kmeans.push_back( ptc );
      if( lc == i )
        std::cout << "*";
      else
        std::cout << " ";
      std::cout << i << " " << p << "\t" << (*(d->kcenters))[i].x
                << " " << (*(d->kcenters))[i].y << "\t"
                << (*(d->kcenters))[i].bx << " "
                << (*(d->kcenters))[i].by << " "
                << (*(d->kcenters))[i].bscore << " | "
                << (*(d->kcenters))[i].vx << " "
                << (*(d->kcenters))[i].vy << "\n";
    }
  std::cout << "entropy: " << *(d->clusterentropy) << "\n";
  if( d->statdb != NULL )
    {
#ifdef __DEBUG_PSO_ISNAN__
      for( unsigned ii = 0; ii < d->parlist.size(); ii++ )
        {
          if( isnan(d->parlist[ii]->x) || isnan(d->parlist[ii]->y) ||
              isnan(d->parlist[ii]->vx) || isnan(d->parlist[ii]->vy) )
            {
              std::cout << "pso track isnan\n";
              exit(0);
            }
        }
#endif
      d->statdb->save_iterationdata( iter, d->parlist, kmeans,
                                     *(d->cachestat),
                                     d->newtrans );
    }
  return 0;              
}

int pthread_psocluster( psothreadd *d )
{
  if( d->nslices != (int) d->particles.size() )
    {
      std::cout << "ERROR nslice not match\n";
      exit(0);
    }
  if( d->tnum == 0 )
    {
      *(d->clusflag) = 1;
      // reset only when siginaled
      if( d->clusterreset == 1 )
        {
          std::cout << "resetting cluster\n";
          for( int i = 0; i < *(d->numclusters); i++ )
            {
              int s, p;
              s = random()%d->nslices;
              p = random()%d->npar;
              psoparticle *tvec;
              tvec = d->particles[s];
              (*(d->kcenters))[i].x = tvec[p].x;
              (*(d->kcenters))[i].y = tvec[p].y;
            }
          d->clusterreset = 0;
        }
      // we prefer the largest cluster to be at the top this prevents the cluster from being deleted
      else if( *(d->largestcluster) != 0 )
        {
          psoparticle tpsop;
          tpsop = (*(d->kcenters))[ *(d->largestcluster) ];
          (*(d->kcenters))[ *(d->largestcluster) ] = (*(d->kcenters))[ 0 ];
          (*(d->kcenters))[ 0 ] = tpsop;
        }
    }
  int iter = 0;
  pthread_barrier_wait( d->slicesync );
#ifdef __DEBUG_THEAD_CLUSTER__
  std::cout << "thread: " << d->tnum << " flag: "
            << *(d->clusflag) << "\n";
#endif
  int ret;
  while( *(d->clusflag) == 1 && iter < 100) 
    {
      //assign clusters
      pthread_barrier_wait( d->slicesync );
      if( d->tnum == 0 )
        *(d->clusflag) = 0;
      pthread_barrier_wait( d->slicesync );
#ifdef __DEBUG_THEAD_CLUSTER__
      std::cout << "thread: " << d->tnum << " assigning clusters "
                << d->parnumlow << " " << d->parnumhigh << " "
                << *(d->numclusters) << "\n";
#endif
      ret = pthread_psocluster( d, d->kcenters, *(d->numclusters) );
      if( ret == 1 )
	*(d->clusflag) = ret;
      pthread_barrier_wait( d->slicesync );

#ifdef __DEBUG_THEAD_CLUSTER__
      std::cout << "thread: " << d->tnum << " moving clusters\n";
#endif
      // calculate centers
      pthread_psoclustermove( d, (d->tkcenters)[d->tnum], *(d->numclusters) );
      pthread_barrier_wait( d->slicesync );

      for( int i = d->tnum; i < *(d->numclusters); i = i + d->totalthreads )
        {
          psoparticle *myc;
          myc = &((*(d->kcenters))[i]);
          myc->x = myc->y = myc->cscore = 0;
          for( int t = 0; t < NUMTHREADS; t++ )
            {
              myc->x = myc->x + (d->tkcenters[t])[i].x;
              myc->y = myc->y + (d->tkcenters[t])[i].y;
              myc->cscore = myc->cscore + (d->tkcenters[t])[i].cscore;
            }
          if( myc->cscore > 0 )
            {
              myc->x = myc->x/myc->cscore;
              myc->y = myc->y/myc->cscore;
            }
        }
      
#ifdef __DEBUG_THEAD_CLUSTER__
          std::cout << "thread: " << d->tnum << " flag: "
                    << *(d->clusflag) << "\n";
#endif
      
      pthread_barrier_wait( d->slicesync );
      iter++;
    }

  pthread_psoclusterbests( d, (d->tkcenters)[d->tnum], *(d->numclusters) );
  pthread_barrier_wait( d->slicesync );
  for( int i = d->tnum; i < *(d->numclusters); i = i + d->totalthreads )
    {
      psoparticle *myc;
      myc = &((*(d->kcenters))[i]);
      myc->bx = myc->by = myc->bscore = 0;
      myc->vx = myc->vy = 0;
      for( int t = 0; t < NUMTHREADS; t++ )
        {
          if( myc->bscore < (d->tkcenters[t])[i].bscore )
            {
              myc->bscore = (d->tkcenters[t])[i].bscore;
              myc->bx = (d->tkcenters[t])[i].bx;
              myc->by = (d->tkcenters[t])[i].by;
            }
          myc->vx = myc->vx + (d->tkcenters[t])[i].vx;
          myc->vy = myc->vy + (d->tkcenters[t])[i].vy;
        }
      myc->vx = myc->vx / ( (double) myc->cscore );
      myc->vy = myc->vy / ( (double) myc->cscore );
    }
  pthread_barrier_wait( d->slicesync );

  return 0;
}

int pthread_psoclusterbests( psothreadd *d,
                             std::vector<psoparticle> &kc,
                             int numclusters )
{
  if( (int)kc.size() < numclusters )
    {
      psoparticle tkc;
      tkc.x = tkc.y = tkc.cscore = 0;
      tkc.bx = tkc.by = tkc.bscore = 0;
      kc.resize( numclusters, tkc );
    }
  for( unsigned int i = 0; i < kc.size(); i++ )
    kc[i].bx = kc[i].by = kc[i].vx = kc[i].vy = kc[i].bscore = 0;

  for( int i = 0; i < d->nslices; i++ )
    {
      for( int j = d->parnumlow; j < d->parnumhigh ; j++)
        {
          int c;
          c = d->particles[i][j].cnum;
          kc[c].vx = kc[c].vx + 
            pow( (*(d->kcenters))[c].x - d->particles[i][j].x, 2 );
          kc[c].vy = kc[c].vy + 
            pow( (*(d->kcenters))[c].y - d->particles[i][j].y, 2 );

#ifdef __DEBUG_PSO_ISNAN__
          if( isnan(d->particles[i][j].x) ||
              isnan(d->particles[i][j].y) )
            {
              std::cout << "ERROR in cluster move particle location is nan\n";
              exit(0);
            }
#endif
          if( kc[c].bscore < d->particles[i][j].cscore )
            {
              kc[c].bscore = d->particles[i][j].cscore;
              kc[c].bx = d->particles[i][j].x;
              kc[c].by = d->particles[i][j].y;
            }
        }
    }
    
#ifdef __DEBUG_PSO_ISNAN__
  if( isnan( kc->x ) || isnan( kc->y ) )
    {
      std::cout << "ERROR in cluster move cluster cneter = 0 \n";
      exit(0);
    }
#endif

  return 0;
}

int pthread_psoupdate( psoparticle *p, std::vector< psoparticle> *kc,
                       std::vector< psoparticle > *bs, psoparticle *lc,
                       double lgc, psothreadd *d )
{
  p->vx = p->vx/4.0 + 1.5*myrand()*( p->bx - p->x )
    + 1.0*myrand()*( (*bs)[ p->imnum ].bx - p->x )
    + lgc*myrand()*( lc->x - p->x );

  p->vy = p->vy/4.0 + 1.5*myrand()*( p->by - p->y )
    + 1.0*myrand()*( (*bs)[ p->imnum ].by - p->y )
    + lgc*myrand()*( lc->y - p->y );
  
  if( p->vx < -32 )
    p->vx = -32;
  if( p->vy < -32 )
    p->vy = -32;
  if( p->vx > 32 )
    p->vx = 32;
  if( p->vy > 32 )
    p->vy = 32;
  
  p->x = p->vx + p->x;
  p->y = p->vy + p->y;

  // pick new center when it is out of bounds
  while( p->x < d->searchspace.lowx || p->x > d->searchspace.highx 
         || p->y < d->searchspace.lowy || p->y > d->searchspace.highy )
    {
      int newcenter;
      newcenter = random() % kc->size();
      p->x = (*kc)[newcenter].x + (myrand()-0.5)*10.0;
      p->y = (*kc)[newcenter].y + (myrand()-0.5)*10.0;
    }
#ifdef __DEBUG_PSO_ISNAN__
  if( isnan(p->vy) )
    {
      std::cout << "particle vy nan\n";
      exit(0);
    }
  if( isnan(p->vx) )
    {
      std::cout << "particle vx nan\n";
      exit(0);
    }
  if( isnan(p->x) )
    {
      std::cout << "particle x nan\n";
      exit(0);
    }
#endif
  return 0;
}

int pthread_psoupdate( psothreadd *d )
{
  double pp;
  double z;
  z = 1.0/( (double)(*(d->numclusters)) );
  if( *(d->numclusters) == 1 )
    pp = 1;
  else
    pp = ( 1.0-((*(d->clusterentropy))/(log( z ))) )*4;
  if( d->tnum == 0 )
    std::cout << "contribution: " << pp << "\n";
  for( int i = 0; i < d->nslices; i++ )
    {
      for( int j = d->parnumlow; j < d->parnumhigh; j++)
        {
          pthread_psoupdate( &(d->particles[i][j]), d->kcenters,
                             d->besttrans,
                             &((*(d->kcenters))[ *(d->largestcluster) ]),
			     pp, d);
        }
    }
  return 0;
}

int pthread_psoparticle( psothreadd *d )
{
  for( int i = 0; i < d->nslices; i++ )
    {
      for( int j = d->parnumlow; j < d->parnumhigh; j++)
        {
          double s;
          transformation ts;
          ts.transx = (int)round(d->particles[i][j].x);
          ts.transy = (int)round(d->particles[i][j].y);
          int ret;
          ret = d->tcache[i]->get( ts );
          if( ret == -1 )
            {
              d->particles[i][j].newpt = 1;
              s = d->mi->miscore( *(d->imstacka[i]), *(d->imstackb[i]),
                                  (*d->mapa), *(d->mapb), (int)ts.transx, (int)ts.transy );
            }
          else
            {
              d->particles[i][j].newpt = 0;
              s = ts.score;
            }
          d->particles[i][j].cscore = s;
          if( d->particles[i][j].bscore < s )
	    {
	      d->particles[i][j].bscore = s;
	      d->particles[i][j].bx = round( d->particles[i][j].x );
	      d->particles[i][j].by = round( d->particles[i][j].y );
	    }
        }
      if( i % 4 == 0 )
        pthread_barrier_wait( d->slicesync );
    }
  return 0;
}

int pthread_psotcache( psothreadd *d, transcache *c, psoparticle *pl )
{
  if( d->newtrans != NULL )
    d->newtrans[pl[0].imnum].clear();
  for( int i = 0; i < d->npar; i++ )
    {
      if( pl->newpt == 1 )
        {
          transformation ts;
          ts.transx = round(pl[i].x);
          ts.transy = round(pl[i].y);
          ts.score = pl[i].cscore;
          if( d->newtrans != NULL )
            d->newtrans[pl[i].imnum].push_back( ts );
          c->add( ts );
        }
    }
  (*(d->cachestat))[ pl[0].imnum ].size = c->size();
  (*(d->cachestat))[ pl[0].imnum ].hit = c->hitfrequency();
  c->resetfreq();
  return 0;
}

int pthread_psocluster( psothreadd *d, std::vector< psoparticle > *kc, int nclus )
{
  int ret;
  ret  = 0;
  for( int i = 0; i < d->nslices; i++ )
    {
      for( int j = d->parnumlow; j < d->parnumhigh; j++)
        {
          int k;
          k = pthread_clusterclosest( d->particles[i][j], kc, nclus );
          if( d->particles[i][j].cnum != k )
            {
	      ret = 1;
              d->particles[i][j].cnum = k;
            }
        }
    }
  return ret;
}

int pthread_clusterclosest( psoparticle &p, std::vector< psoparticle > *kc, int nclus )
{
  int k = 0;
  double dis;
#ifdef __DEBUG_PSO_ISNAN__
  if( isnan( (*kc)[0].x ) || isnan( (*kc)[0].y ) )
    {
      std::cout << "ERROR in find closest cluster center is nan\n";
      exit(0);
    }
#endif
  dis = pow((*kc)[0].x - p.x, 2 ) + pow((*kc)[0].y - p.y, 2 );
  for( int i = 1; i < nclus; i++ )
    {
      double t;
#ifdef __DEBUG_PSO_ISNAN__
      if( isnan( (*kc)[i].x ) || isnan( (*kc)[i].y ) )
	{
	  std::cout << "ERROR in find closest cluster center is nan\n";
	  exit(0);
	}
#endif
      t = pow((*kc)[i].x - p.x, 2 ) + pow((*kc)[i].y - p.y, 2 );
      if( t < dis )
        {
          k = i;
          dis = t;
        }
    }
  return k;
}

int pthread_psoclustermove( psothreadd *d, std::vector<psoparticle> &kc,
                            int numclusters )
{
  if( (int)kc.size() < numclusters )
    {
      psoparticle tkc;
      tkc.x = tkc.y = tkc.cscore = 0;
      tkc.bx = tkc.by = tkc.bscore = 0;
      kc.resize( numclusters, tkc );
    }
  for( unsigned int i = 0; i < kc.size(); i++ )
    {
      kc[i].x = kc[i].y = kc[i].cscore = 0;
    }

  for( int i = 0; i < d->nslices; i++ )
    {
      for( int j = d->parnumlow; j < d->parnumhigh ; j++)
        {
          int c;
          c = d->particles[i][j].cnum;
#ifdef __DEBUG_PSO_ISNAN__
          if( isnan(d->particles[i][j].x) ||
              isnan(d->particles[i][j].y) )
            {
              std::cout << "ERROR in cluster move particle location is nan\n";
              exit(0);
            }
#endif
          kc[c].x = kc[c].x + d->particles[i][j].x;
          kc[c].y = kc[c].y + d->particles[i][j].y;
          kc[c].cscore = kc[c].cscore + 1;
        }
    }
    
#ifdef __DEBUG_PSO_ISNAN__
  if( isnan( kc->x ) || isnan( kc->y ) )
    {
      std::cout << "ERROR in cluster move cluster cneter = 0 \n";
      exit(0);
    }
#endif
  return 0;
}

pso::pso()
{
}

pso::~pso()
{
  del();
}

int pso::del()
{
  for( unsigned i = 0; i < particles.size(); i++)
    {
      delete [] particles[i];
    }
  for( unsigned i = 0; i < tcache.size(); i++)
    delete tcache[i];
  tcache.clear();
  particles.clear();
  besttrans.clear();
  return 0;
}

int pso::getresult( double &tx, double &ty, double &score )
{
  tx = kcenters[largestcluster].bx;
  ty = kcenters[largestcluster].by;
  score = kcenters[largestcluster].bscore;
  return 0;
}

int pso::getresult( std::vector< transformation > &trans )
{
  transformation t, t1;
  double total;
  total = parslice*imstacka.size();
  trans.clear();
  std::cout << "getting " << numberclusters << " results\n";
  for( int i = 0; i < numberclusters; i++ )
    {
      t.transx = kcenters[i].x;
      t.transy = kcenters[i].y;
      t.x_var = kcenters[i].vx;
      t.y_var = kcenters[i].vy;
      t.score = kcenters[i].cscore;
      std::cout << t.transx << " " << t.transy << " " << t.score << "\n";
      trans.push_back( t );
    }
  std::sort( trans.begin(), trans.end() );
  return 0;
}

int pso::initparticles( int npar )
{
  del();
  psoparticle p;
  int xwindow;
  int ywindow;
  std::cout << "search window: " 
            << searchspace.lowx << " " 
            << searchspace.lowy << " "
            << searchspace.highx << " " 
            << searchspace.highy << "\n";
  xwindow = (int)( searchspace.highx - searchspace.lowx );
  ywindow = (int)( searchspace.highy - searchspace.lowy );
  std::cout << "\t" << xwindow << " " << ywindow << "\n";

  for( unsigned s = 0; s < imstacka.size(); s++ )
    {
      particles.push_back( new psoparticle[ npar ] );
      tcache.push_back( new transcache );
      besttrans.push_back( p );
      besttrans[s].bscore = 0;
      for( int i = 0; i < npar; i++ )
        {
          particles[s][i].x = random()%xwindow + searchspace.lowx;
          particles[s][i].y = random()%ywindow + searchspace.lowy;

          particles[s][i].vx = random()%10-5;
          particles[s][i].vy = random()%10-5;

          particles[s][i].bx = 0;
          particles[s][i].by = 0;
          particles[s][i].bscore = 0;
          particles[s][i].imnum = s;
          particles[s][i].imparnum = i;
        }
    }
  bestx = besty = bscore = 0;
  return 0;
}

double pso::rand()
{
  double s;
  s =((double) (random() % 10000 ))/10000.0;
  if( s > 1 || s < 0 )
    std::cout << "ERROR rand\n";
  return s;
}

int pso::init( std::vector< image2d* > ima, std::vector< image2d* > imb,
               image2d &ma, image2d &mb)
{
  del();
  imstacka = ima;
  imstackb = imb;
  mapa.load(&ma);
  mapb.load(&mb);
  return 0;
}

int pso::distributework()
{
  for( unsigned s = 0; s < particles.size(); s++ )
    {
      for( int i = 0; i < parslice; i++ )
        {
          distributeparticle( i, s );
        }
    }
  clusman->runjoblist();
  std::vector< workunit* > wus;
  clusman->getfinished( wus );
  for( unsigned i = 0; i < wus.size(); i++ )
    {
      int p, s, tx, ty;
      wus[i]->getinput( tx, ty, s, p );
      double score;
      wus[i]->getresult( score );
      particles[s][p].cscore = score;
      delete wus[i];
    }
  wus.clear();
  for( unsigned s = 0; s < particles.size(); s++ )
    {
      for( int p = 0; p < parslice; p++ )
        {
          if( particles[s][p].bscore < 
              particles[s][p].cscore )
            {
              particles[s][p].bscore = particles[s][p].cscore;
              particles[s][p].bx = round( particles[s][p].x );
              particles[s][p].by = round( particles[s][p].y );
            }
          if( besttrans[s].bscore < particles[s][p].cscore )
            {
              besttrans[s].bscore = particles[s][p].cscore;
              besttrans[s].bx = round( particles[s][p].x );
              besttrans[s].by = round( particles[s][p].y );
            }
          if( bscore < particles[s][p].cscore )
            {
              bscore = particles[s][p].cscore;
              bestx = round(particles[s][p].x);
              besty = round(particles[s][p].y);
            }
        }
    }
  return 0;
}

int pso::mpoptimize( int npar, double lowx, double lowy,
                     double highx, double highy, std::string name,
                     psotracker *pt)
{
  psostat = pt;
  searchspace.lowx = lowx;
  searchspace.lowy = lowy;
  searchspace.highx = highx;
  searchspace.highy = highy;
  parslice = npar;
  kcenters.clear();
  for( int i = 0; i < 10; i++ )
    {
      psoparticle tp;
      kcenters.push_back( tp );
    }
  numberclusters = 10;
  srandom( time( NULL ) );
  initparticles( npar );
  std::cout << "particles inited\n";
  //initilize thread data
  std::vector< psothreadd* > td;
  int exit;
  exit = 0;
  int clusflag;
  clusflag = 0;
  pthread_barrier_t globalb, localb;

  pthread_barrier_init( &globalb, NULL, NUMTHREADS+1 );
  pthread_barrier_init( &localb, NULL, NUMTHREADS );
  int secsize;
  secsize = npar/NUMTHREADS;

  std::vector< psoparticle* > parlist;
  for( unsigned s = 0; s < particles.size(); s++ )
    for( int p = 0; p < parslice; p++ )
      parlist.push_back( &(particles[s][p]) );

  std::vector< transformation > *newtrans;
  newtrans = new std::vector< transformation >[particles.size()];
  std::vector< psotrack_cachestat > cachestat;
  for( unsigned s = 0; s < particles.size(); s++ )
    {
      psotrack_cachestat pc;
      pc.psoslice = s;
      pc.size = 0;
      pc.hit = 0;
      cachestat.push_back( pc );
    }
  std::vector< psoparticle > tkcenters[NUMTHREADS];
  for( int i = 0; i < NUMTHREADS; i++ )
    {
      td.push_back( new psothreadd );
      td[i]->alnname = name;
      td[i]->tnum = i;
      td[i]->totalthreads = NUMTHREADS;

      td[i]->imstacka = imstacka;
      td[i]->imstackb = imstackb;

      td[i]->slicesync = &localb;
      td[i]->gsync = &globalb;

      td[i]->particles = particles;
      td[i]->kcenters = &kcenters;
      td[i]->tkcenters = tkcenters;
      td[i]->besttrans = &besttrans;
      td[i]->tcache = tcache;

      td[i]->mapa = &mapa;
      td[i]->mapb = &mapb;

      td[i]->parnumlow = secsize*i;
      td[i]->parnumhigh = secsize*(i+1);
      if( i == (NUMTHREADS-1) )
        td[i]->parnumhigh = npar;
      else if( td[i]->parnumhigh > parslice )
        td[i]->parnumhigh = parslice;
      std::cout << td[i]->parnumlow << ":" << td[i]->parnumhigh << "\n";
      td[i]->nslices = (int)imstacka.size();
      td[i]->npar = parslice;

      td[i]->clusflag = &clusflag;
      td[i]->numclusters = &numberclusters;
      td[i]->clusterentropy = &clusterentropy;
      td[i]->largestcluster = &largestcluster;

      td[i]->mi = new mutualinformation;
      td[i]->searchspace = searchspace;
      td[i]->exit = &exit;

      td[i]->statdb = psostat;
      td[i]->parlist = parlist;
      td[i]->newtrans = newtrans;
      td[i]->cachestat = &cachestat;
      std::cout << "thread " << i << " data set\n";
    }

  // start threads
  for( int i = 0; i < NUMTHREADS; i++ )
    {
      int ret;
      ret = pthread_create( &(td[i]->tid), NULL, psothread_start, td[i] );
      std::cout << "started thread: " << i << " return: " << ret << "\n";
    }
  std::cout << "threads started\n";
  void *p;
  pthread_barrier_wait( &globalb );
  for( int i = 0; i < NUMTHREADS; i++ )
    pthread_join( td[i]->tid, &p );

  // clean up thread data
  for( int i = 0; i < NUMTHREADS; i++ )
    {
      delete td[i]->mi;
      delete td[i];
    }
  delete [] newtrans;
  pthread_barrier_destroy( &localb );
  pthread_barrier_destroy( &globalb );
  return 0;
}

int pso::optimize( int npar, double lowx, double lowy,
                   double highx, double highy, clustermanager *cm)
{
  searchspace.lowx = lowx;
  searchspace.lowy = lowy;
  searchspace.highx = highx;
  searchspace.highy = highy;

  clusman = cm;
  if( cm != NULL )
    cm->setdata( imstacka, imstackb, &mapa, &mapb );    

  parslice = npar;
  kcenters.clear();
  psoparticle tp;
  int flag = 5;
  for( int i = 0; i < 10; i++ )
    kcenters.push_back(tp);
  numberclusters = 10;
  srandom( time( NULL ) );
  initparticles( npar );
  std::vector< double > medianspeed;
  medianspeed.clear();
  timer mytimer;
  clusteritr = 0;
  for( unsigned k = 0; k < besttrans.size(); k++ )
    medianspeed.push_back(0);
  for( int itr = 0; itr < 250; itr++ )
    {
      std::cout << "calc particle\n";
      mytimer.reset();
      mytimer.start();
      if( cm == NULL )
        {
          for( unsigned s = 0; s < particles.size(); s++ )
            {
              for( int i = 0; i < npar; i++ )
                {
                  calcparticle((particles[s])+i, s);
                }
            }
        }
      else
        {
          distributework();
        }
      mytimer.stop();
      std::cout << "mutual information calculations done in " 
                << mytimer.getseconds() << "s (" 
                << ((double)(mytimer.getseconds()))/((double)(particles.size()*npar))
                << ")s\n";
      kmeans();
      findbestcluster();
      for( unsigned s = 0; s < particles.size(); s++ )
        {
          for( int i = 0; i < npar; i++ )
            {
              updateparticle((particles[s])+i);
            }
        }
      std::cout << "itr: " << itr << "\n";
      for( unsigned k = 0; k < kcenters.size(); k++ )
        {
          if( (int)k == largestcluster )
            std::cout << "*";
          std::cout << "\t" << k << ":" 
                    << kcenters[k].bx << " " 
                    << kcenters[k].by
                    << " : " 
                    << kcenters[k].bscore 
                    << " - " << kcenters[k].x 
                    << " " << kcenters[k].y 
                    << " - " << kcenters[k].cscore 
                    << "\n";
        }
      double z, pp;
      z = 1.0/( (double)kcenters.size() );
      pp = ( 1.0-(clusterentropy/(log( z ))) )*4;
      std::cout << "entropy: " << clusterentropy << "\n";
      std::cout << "cluster contributeion: " << pp << "\n";
      for( unsigned k = 0; k < besttrans.size(); k++ )
        std::sort( particles[k], particles[k]+npar, particle_cmp );
      
      for( unsigned s = 0; s < tcache.size(); s++ )
        {
          medianspeed[s] = pow( particles[s][(int)(npar * 0.6)].vx, 2 ) + 
            pow( particles[s][(int)(npar * 0.6)].vy, 2 );
        }
      std::sort( medianspeed.begin(), medianspeed.end() );
      std::cout << "convergance check: " << medianspeed[ (int)( tcache.size()*0.6) ] << "\n";
      if( medianspeed[ (int)( tcache.size()*0.6) ] < 1 && goodclusters == 1 )
        flag--;
      else
        flag = 10;
      if( flag < 0 )
        break;
      clusteritr++;
    }
  std::cout << "search space covered:\n";
  double spacesize;
  spacesize = (highx-lowx)*(highy-lowy);
  for( unsigned i = 0; i < tcache.size(); i++ )
    std::cout << i << " : "
              << ( (double)tcache[i]->size() )/spacesize << "\n";

  return 0;
}

int pso::findbestcluster()
{
  int b = 0;
  double bs = kcenters[0].cscore;
  int nunity = 0;
  clusterentropy = 0;
  double total;
  total = parslice*imstacka.size();
  std::cout << "finding best cluster\n\t";
  double avgsz = 0;
  double stdvar;
  for( unsigned i = 1; i < kcenters.size(); i++ )
    {
      if( kcenters[i].cscore > bs )
        {
          b = i;
          bs = kcenters[i].cscore;
        }
      if( kcenters[i].cscore < 0.05*tcache.size()*parslice )
        nunity++;
      double p;
      p = kcenters[i].cscore/total;
      if( p > 0 )
        clusterentropy = clusterentropy + p*log(p);
      avgsz = avgsz + kcenters[i].cscore;
    }
  avgsz = avgsz / ( (double)( kcenters.size() ) );
  stdvar = 0;
  for( unsigned i = 0; i < kcenters.size(); i++ )
    stdvar = stdvar + pow( avgsz - kcenters[i].cscore, 2 );
  stdvar = sqrt( stdvar / ((double)( kcenters.size() ) ) );
  goodclusters = 1;
  for( unsigned i = 0; i < kcenters.size(); i++ )
    if( (int)i != b && kcenters[i].cscore > (bs - stdvar) )
      goodclusters = 0;
  std::cout << "\n";
  std::cout << "variance: " << stdvar << " good distribution: " 
            << goodclusters << "\n";
  clusterentropy = -clusterentropy;
  std::cout << "\t" << clusterentropy << "\n";
  if( clusteritr > 5 )
    {
      if( nunity > 0 )
        clusteritr = 0;
      numberclusters = kcenters.size()-nunity;
    }
  largestcluster = b;
  return 0;
}

int pso::distributeparticle( int p, int s )
{
  //double ss;
  mutualinformation ia;
  transformation ts;
  ts.transx = floor(particles[s][p].x);
  ts.transy = floor(particles[s][p].y);
  if( tcache[s]->get( ts ) == -1 )
    {
      workunit *wu;
      wu = new workunit;
      wu->init( (int)ts.transx, (int)ts.transy, s, p );
      clusman->addwork( wu );
    }
  else
    particles[s][p].cscore = ts.score;

  return 0;
}

int pso::calcparticle( psoparticle *p, int s )
{
  double ss;
  mutualinformation ia;
  transformation ts;
  ts.transx = floor(p->x);
  ts.transy = floor(p->y);
  if( tcache[s]->get( ts ) == -1 )
    {
      ss = ia.miscore( *(imstacka[s]), *(imstackb[s]),
                       (int)(p->x), (int)(p->y) );
      ts.score = ss;
      tcache[s]->add( ts );
    }
  else
    ss = ts.score;
  p->cscore = ss;
  if( p->bscore < ss )
    {
      p->bscore = ss;
      p->bx = round(p->x);
      p->by = round(p->y);
    }
  if( besttrans[s].bscore < ss )
    {
      besttrans[s].bscore = ss;
      besttrans[s].bx = floor( p->x );
      besttrans[s].by = floor( p->y );
    }
  if( bscore < ss )
    {
      bscore = ss;
      bestx = round(p->x);
      besty = round(p->y);
    }
  return 0;
}

int pso::updateparticle( psoparticle *p )
{
  int c;
  c = p->cnum;
  double pp;
  double z;
  z = 1.0/( (double)kcenters.size() );
  pp = ( 1.0-(clusterentropy/(-log( z ))) )*4;
  p->vx = p->vx/4.0 + 1.5*rand()*( p->bx - p->x )
    //    + 2.0*rand()*( kcenters[c].bx - p->x )
    + 1.5*rand()*( besttrans[ p->imnum ].bx - p->x )
    + pp*rand()*( kcenters[largestcluster].x - p->x );
  p->vy = p->vy/4.0 + 1.5*rand()*( p->by - p->y )
    //    + 2.0*rand()*( kcenters[c].by - p->y )
    + 1.5*rand()*( besttrans[ p->imnum ].by - p->y )
    + pp*rand()*( kcenters[largestcluster].y - p->y );

  if( p->vx < -32 )
    p->vx = -32;
  if( p->vy < -32 )
    p->vy = -32;
  if( p->vx > 32 )
    p->vx = 32;
  if( p->vy > 32 )
    p->vy = 32;

  p->x = p->vx + p->x;
  p->y = p->vy + p->y;
  if( p->x < searchspace.lowx )
    {
      p->vx = 0;
      p->x = searchspace.lowx;
    }
  if( p->x > searchspace.highx )
    {
      p->vx = 0;
      p->x = searchspace.highx;
    }
  if( p->y < searchspace.lowy )
    {
      p->vy = 0;
      p->y = searchspace.lowy;
    }
  if( p->y > searchspace.highy )
    {
      p->vy = 0;
      p->y = searchspace.highy;
    }
  return 0;
}

int pso::kmeans()
{
  int flag = 1;
  int iter = 0;
  while( (int)kcenters.size() > numberclusters )
    kcenters.pop_back();
  /// randomly assign cluster centers
  for( int i = 0; i < numberclusters; i++ )
    {
      int s, p;
      s = random()%particles.size();
      p = random()%parslice;
      kcenters[i].x = particles[s][p].x;
      kcenters[i].y = particles[s][p].y;
    }
  while( flag == 1 )
    {
      flag = 0;
      for( unsigned i = 0; i < kcenters.size(); i++ )
        {
          kcenters[i].bx = kcenters[i].by 
            = kcenters[i].bscore = 0;
          kcenters[i].cnum = i;
        }
      for( unsigned i = 0; i < particles.size(); i++) 
        {
          for( int j = 0; j < parslice; j++ )
            {
              int cnum;
              cnum = closestcluster( particles[i][j] );
              if( particles[i][j].cnum != cnum )
                flag++;
              particles[i][j].cnum = cnum;
              if( kcenters[cnum].bscore <
                  particles[i][j].bscore )
                {
                  kcenters[cnum].bscore = 
                    particles[i][j].bscore;
                  kcenters[cnum].bx = particles[i][j].bx;
                  kcenters[cnum].by = particles[i][j].by;
                }
            }
        }
      for( unsigned i = 0; i < kcenters.size(); i++ )
        {
          kcenters[i].cscore = 
            kcenters[i].x = 
            kcenters[i].y = 0;
        }
      for( unsigned i = 0; i < particles.size(); i++) 
        {
          for( int j = 0; j < parslice; j++ )
            {
              int cnum;
              cnum = particles[i][j].cnum;
              kcenters[cnum].x += particles[i][j].x;
              kcenters[cnum].y += particles[i][j].y;
              kcenters[cnum].cscore += 1;
            }
        }
      for( unsigned i = 0; i < kcenters.size(); i++ )
        {
          if( kcenters[i].cscore > 0 )
            {
              kcenters[i].x 
                = kcenters[i].x/kcenters[i].cscore;
              kcenters[i].y 
                = kcenters[i].y/kcenters[i].cscore;
            }
          else
            {
              int z, zz;
              z = random()%particles.size();
              zz = random()%parslice;
              kcenters[i].x = particles[z][zz].x;
              kcenters[i].y = particles[z][zz].y;
            }
        }
      if( iter < 2 )
        {
          flag = 1;
        }
      else
        {
          if( flag > 2)
            flag = 1;
          else
            flag = 0;
        }
      if(iter > 100 )
        break;
      iter++;
    }
  return 0;
}

int pso::closestcluster( psoparticle &p )
{
  double bd;
  int res;
  bd = pow( kcenters[0].x - p.x, 2) +
    pow( kcenters[0].y -p.y, 2 );
  res = 0;
  for( unsigned i = 1; i < kcenters.size(); i++ )
    {
      double s;
      s = pow( kcenters[i].x - p.x, 2) +
        pow( kcenters[i].y -p.y, 2 );
      if( bd > s )
        {
          bd = s;
          res = (int)i;
        }
    }
  return res;
}
