#include "kmeans.h"

kmeans::kmeans()
{
  nocopy = 0;
}

kmeans::~kmeans()
{
  clear();
  clearclusters();
}

int kmeans::load( double *pts, int ndata)
{
  clear();
  clearclusters();
  nocopy = 0;
  for( int i = 0; i < ndata; i++ )
    {
      datapts.push_back( new double [2] );
      datapts_class.push_back(0);
      datapts[i][0] = pts[i];
    }
  npts = ndata;
  nparams = 1;
  return 0;
}

int kmeans::load( double **pts, int ndata,
                  int np)
{
  clear();
  clearclusters();
  nocopy = 0;
  for( int i = 0; i < ndata; i++ )
    {
      datapts.push_back(new double [np]);
      datapts_class.push_back(0);
      //      double *d;
      memcpy( datapts.back(), pts[i], 
              sizeof(double)*np );
    }
  npts = ndata;
  nparams = np;
  return 0;
}

int kmeans::loadnocopy( std::vector< double* > &pts, int np )
{
  clear();
  clearclusters();
  nocopy = 1;
  datapts = pts;
  npts = (int) pts.size();
  datapts_class.resize( npts, 0 );
  nparams = np;
  return 0;
}

int kmeans::clear()
{
  if( nocopy == 0 )
    {
      for( unsigned i = 0; i < datapts.size(); i++ )
        delete [] datapts[i];
    }
  datapts.clear();
  datapts_class.clear();
  return 0;
}

int kmeans::clearclusters()
{
  for( unsigned i = 0; i < centers.size(); i++ )
    delete [] centers[i];
  centers.clear();
  return 0;
}

int kmeans::group()
{
  double d, md;

  int change;
  change = 0;
  for( unsigned int i = 0; i < datapts.size(); i++ )
    {
      md = distance( datapts[i], centers[0] );
      int oldc;
      oldc = datapts_class[i];
      datapts_class[i] = 0;
      for( unsigned int j = 1; j < centers.size(); j++ )
        {
          d = distance( datapts[i], centers[j] );
          if( d < md )
            {
              md = d;
              datapts_class[i] = j;
            }
        }
      if( oldc != datapts_class[i] )
        change++;
    }

  return change;
}

int kmeans::recalcenter()
{
  for( unsigned i = 0; i < centers.size(); i++ )
    {
      for( int j = 0; j < nparams; j++ )
        centers[i][j] = 0;
      csize[i] = 0;
    }
  for( unsigned int i = 0; i < datapts.size(); i++ )
    {
      add( centers[ datapts_class[i] ],
           datapts[i] );
      csize[ datapts_class[i] ]++;
    }
#ifdef __DEBUG_KMEANS__
  std::cout << "cluster size: ";
#endif
  for( unsigned int i = 0; i < centers.size(); i++ )
    {
      if( csize[i] == 0 )
        {
          int p;
          csize[i] = 1;
          p = random() % npts;
          memcpy( centers[i], datapts[p], sizeof(double)*nparams );
        }
      else
        divide( centers[i], csize[i] );
#ifdef __DEBUG_KMEANS__
      std::cout << csize[i] << " ";
#endif
    }
#ifdef __DEBUG_KMEANS__
  std::cout << "\n";
    
#endif
  return 0;
}

double kmeans::distance( double *a, double *b )
{
  double d = 0;
  for( int i = 0; i < nparams; i++ )
    d = d + pow( a[i] - b[i], 2 );
  return d;
}

int kmeans::add( double *a, double *b )
{
  for( int i = 0; i < nparams; i++ )
    a[i] += b[i];
  return 0;
}

int kmeans::divide( double *a, double v )
{
  for( int i = 0; i < nparams; i++ )
    a[i] = a[i] / v;
  return 0;
}

int kmeans::clustersize( int c )
{
  return csize[c];
}

int kmeans::cluster( int nc )
{
#ifdef __DEBUG_KMEANS__
  std::cout << "clearing...\n";
#endif
  clearclusters();
#ifdef __DEBUG_KMEANS__
  std::cout << "init...\n";
#endif
  for( int i = 0 ; i < nc; i++ )
    {
      centers.push_back( new double [nparams] );
      csize.push_back(0);
    }
  for( int i = 0; i < npts; i++ )
    datapts_class[i] = 0;
#ifdef __DEBUG_KMEANS__
  std::cout << "random init...\n";
#endif
  for( unsigned i = 0; i < centers.size(); i++ )
    {
      int p;
      p = random() % npts;
      memcpy( centers[i], datapts[p], sizeof(double)*nparams );
    }
  int flag = 5;
#ifdef __DEBUG_KMEANS__
  std::cout << "begining run\n";
#endif
  int iter;
  iter = 0;
#ifdef __QUICK_TRAIN__
  int traincycles = 2;
#else
  int traincycles = 500;
#endif
  while( flag > 0 && iter < traincycles )
   {
      iter++;
      int p;
#ifdef __DEBUG_KMEANS__
      std::cout << "grouping...\n";
#endif
      p = group();
#ifdef __DEBUG_KMEANS__
      std::cout << "recalc...\n";
#endif
      recalcenter();
      if( p < 2 )
        flag--;
    }

  for( unsigned i = 0; i < centers.size(); i++ )
    {
      for( int j = 0; j < nparams; j++ )
        {
          if( isnan( centers[i][j] ) )
            {
              std::cout << "ERROR got nan as final in kmeans\n";
              std::cout << i << "\n";
              exit(0);
            }
        }
    }

#ifdef __DEBUG_KMEANS__
  std::cout << "clustering finished in " << iter << " iterations\n";
#endif
  return 0;
}

int kmeans::classify( double *d )
{
  double md, td;
  int mc;
  md = distance( d, centers[0] );
  mc = 0;
  for( unsigned int j = 1; j < centers.size(); j++ )
    {
      td = distance( d, centers[j] );
      if( td < md )
        {
          md = td;
          mc = j;
        }
    }
  return mc;
}
