#include "emcluster.h"

emcluster::emcluster()
{
  pi = NULL;
  data = NULL;
  gamma = NULL;
  models = NULL;
}

emcluster::~emcluster()
{
  del();
  delmodels();
}

int emcluster::loaddata(double *d, int size)
{
  del();
  data = new matrix [size];
  datapts = 0;
#ifdef __DEBUG_EM__
  std::cout << "loading data\n"
            << std::flush;
#endif
  for(int i = 0; i < size; i = i + 1)
    {
      data[datapts].init(d+i, 1);
      datapts++;
    }
#ifdef __DEBUG_EM__
  std::cout << "loading done: " << datapts << " datapoints\n"
            << std::flush;
#endif
  params = 1;
  return 0;
}

int emcluster::loaddata(double **d, int size, int nparams, int sample)
{
  del();
  data = new matrix [size];
  datapts = 0;
  for(int i = 0; i < size; i = i + sample)
    {
      data[datapts].init(d[i], nparams);
      datapts++;
    }

  //  datapts = size;
  params = nparams;
  return 0;
}

int emcluster::loaddata(std::vector<double*> &d, int nparams, int sample)
{
  del();
  data = new matrix [d.size()];
  datapts = 0;
  for(unsigned int i = 0; i < d.size(); i = i + sample)
    {
      data[datapts].init(d[i], nparams);
      datapts++;
    }

  params = nparams;
  return 0;
}

int emcluster::del()
{
  if(data != NULL)
    {
      delete [] data;
    }
  data = NULL;
  return 0;
}

int emcluster::estep()
{
  double p[nmodels];
  for(int n = 0; n < datapts; n++)
    {
      double d;
      d = 0;
      for(int k = 0; k < nmodels; k++)
        {
          p[k] = models[k].prob(data[n]);
          if(p[k] < 0)
            return -1;
          d = d + pi[k]*p[k];
        }
      if(d == 0)
        d = 0.000001;
      for(int k = 0; k < nmodels; k++)
        gamma[n][k] = pi[k] * p[k] / d;
    }
  return 0;
}

int emcluster::mstep()
{
  double Nk[nmodels];
  matrix tmp;
  matrix tmp2, tmpc, tmpm;
  for(int i = 0; i < nmodels; i++)
    {
      Nk[i] = 0;
      for(int n = 0; n < datapts; n++)
        Nk[i] = Nk[i] + gamma[n][i];
      if(Nk[i] < 0.000001)
        Nk[i] = 0.000001;
      pi[i] = Nk[i]/datapts;
    }

  for(int k = 0; k < nmodels; k++)
    {
      zeromodel(models[k]);
      models[k].getmean(tmp2);
      for(int n = 0; n < datapts; n++)
        {
          //tmp =  data[n] * gamma[n][k];
          tmp2.addmult( data[n], gamma[n][k] );
        }
      tmp2.divide( Nk[k] );
      if( tmp2.notgood() )
        std::cout << "model " << k << " mean not good\n"
                  << std::flush;
      models[k].setmean(tmp2);
    }

  for(int k = 0; k < nmodels; k++)
    {
      models[k].getvariance(tmpc);
      models[k].getmean(tmpm);
      for(int n = 0; n < datapts; n++)
        {
          tmp = data[n] - tmpm;
          //tmp2 = tmp * tmp.transpose();
          //tmp2 = tmp2 * gamma[n][k];
          tmpc.addtransposemult( tmp, gamma[n][k] );
        }
      tmpc.divide( Nk[k] );

      if( tmpc.notgood() )
        {
          std::cout << "model " << k << " variance not good\n"
                    << std::flush;
          exit(0);
        }

      double v;
      v = tmpc.maxdiag();
      if( v > 10000 )
        {
#ifdef __DEBUG_EM__
          std::cout << "max variance " << v << "\n" << std::flush;
#endif
          v = 10000.0 / v;
          tmpc.multiply( v );
        }
      models[k].setvariance( tmpc );
    }
  return 0;
}

long double emcluster::prob( double *d )
{
  long double res = 0;
  for(int k = 0; k < nmodels; k++)
    {
      long double tmp1;
      tmp1 = models[k].prob( d );
      if( isnan(tmp1) || isinf(tmp1) )
        {
          std::cout << "ERROR prob return nan\n";
          std::cout << models[k] << "\n";
          exit(0);
        }
      if(tmp1 < 0)
        tmp1 = pow(10,-200);
      else
        tmp1 = pi[k] * tmp1;
      if( tmp1 > res )
        res = tmp1;
    }
  return res;
}

long double emcluster::loglikely()
{
  long double res = 0;

  long double tmp, tmp1;
  for(int n = 0; n < datapts; n++)
    {
      tmp = 0;
      for(int k = 0; k < nmodels; k++)
        {
          tmp1 = models[k].prob( data[n] );
          if( isnan(tmp1) || isinf(tmp1) )
            {
              std::cout << "ERROR prob return nan\n";
              std::cout << models[k] << "\n";
              exit(0);
            }
          if(tmp1 < 0)
            return -pow(10, 200);
          tmp = tmp + pi[k] * tmp1;
          if( isnan( tmp ) )
            {
              std::cout << "ERROR tmp is nana\n";
              exit(0);
            }
        }
      if( tmp == 0 )
        {
          std::cout << "ERROR overall tmp is zero for datapoint\n";
          std::cout << "datapoint: " << n << "\n";
          for( int z = 0; z < models[0].size(); z++) 
            std::cout << data[n].get(z, 0) << " ";
          std::cout << "\n";
          tmp = pow(10,-200);
        }
      res = res + log(tmp);
    }
 
  return res;
}

int emcluster::bicrun(int kmin, int kmax)
{
  long double highscore = 0;
  int highk = 0;
  std::cout << "number of data points: " << datapts << "\n";
  for(int i = kmin; i <= kmax; i++)
    {
      long double ts;
      std::cout << "trying " << i << " mixtures : " << std::flush;
      //run multiple searches to avoid local mins
      for(int j = 0; j < 1; j++)
        {
          ts = run(i);
          //ts = ts - 0.5 * (i*(params*params + params)) * log(datapts);  //BIC
          ts = 2*ts-2*(i*(params*params+params))*log(datapts);  //AIC
          std::cout << ts;
          if( (highk == 0 || highscore < ts) && isinf(ts) == 0 )
            {
              std::cout << "*";
              highscore = ts;
              highk = i;
            }
          std::cout << " " << std::flush;
        }
      std::cout << "\n";
    }
  return highk;
}

long double emcluster::run(int z)
{
  long double ll;
  long double llo;
  matrix mt;
#ifdef __DEBUG_EM__
  std::cout << "init... " << z << "\n"
            << std::flush;
#endif
  initmodels( z, mt );
#ifdef __DEBUG_EM__
  std::cout << "calculating first loglikely\n" << std::flush;
#endif
  ll = loglikely();
  int flag = 2;
#ifdef __DEBUG_EM__
  std::cout << "starting to run " << ll << "\n"
            << std::flush;
#endif
  int iter = 0;
  timer mytimer;
#ifdef __QUICK_TRAIN__
  int traincycles = 1;
#else
  int traincycles = 500;
#endif
  while(flag > 0 && iter < traincycles)
    {
#ifdef __DEBUG_EM__
      std::cout << "estep\n" << std::flush;
#endif
      mytimer.reset();
      mytimer.start();
      if(estep() != 0)
        {
          std::cout << "estep failed\n";
          break;
        }
      mytimer.stop();
#ifdef __DEBUG_EM__
      std::cout << "estep finished in: " << mytimer.getseconds() << "s\n"
                << std::flush;
      std::cout << "mstep\n" << std::flush;
#endif
      mytimer.reset();
      mytimer.start();
      mstep();
      mytimer.stop();
#ifdef __DEBUG_EM__
      std::cout << "mstep finished in: " << mytimer.getseconds() << "s\n" 
                << std::flush;
#endif
      int mf;
      mf = 0;
#ifdef __DEBUG_EM__
      std::cout << "checking covariance\n";
#endif
      for( int i = 0; i < nmodels; i++ )
        {
          models[i].fixcovariance();
          if( models[i].isgood() == 0 )
            {
#ifdef __DEBUG_EM__
              std::cout << "reseting model " << i << "\n"
                        << std::flush;
#endif
              mf = 1;
              models[i].setmean( data[ random()%datapts ] );
              models[i].setvariance( mt );
              pi[i] = 1;
            }
          if( isnan(pi[i]) || isinf(pi[i]) || pi[i] == 0 )
            {
              std::cout << "pi[" << i << "] = 0\n"
                        << std::flush;
              pi[i] = 1;
            }
        }
#ifdef __DEBUG_EM__
      std::cout << "checking exit\n"
                << std::flush;
#endif
      if( mf == 0 )
        {
          llo = ll;

          ll = loglikely();
          if(isnan(ll))
            {
              flag = 0;
              std::cout << "ERROR log likely failed\n"
                        << std::flush;
            }
          //      printmodels();
#ifdef __DEBUG_EM__
          std::cout << llo << " -> " << ll << "\npi: ";
          for( int i = 0; i < nmodels; i++ )
            std::cout << pi[i] << " ";
          std::cout << "\n";
#endif
          if( ll - llo < 0.001)
            {
#ifdef __DEBUG_EM__
              std::cout << "..\n";
#endif
              flag--;
            }
        }
      iter++;
    }
#ifdef __DEBUG_EM__
  std::cout << "finished in " << iter << " iterations\n"
            << std::flush;
#endif
  return ll;
}

int emcluster::delmodels()
{
  if( pi != NULL )
    delete [] pi;
  if( gamma != NULL )
    {
      for(int i = 0; i < datapts; i++)
        delete [] gamma[i];
      delete [] gamma;
    }
  if(models != NULL)
    delete [] models;
  pi = NULL;
  gamma = NULL;
  models = NULL;
  return 0;
}

int emcluster::initmodels(int z, matrix &vmt)
{
  delmodels();
  srand(time(NULL));
  nmodels = z;
  pi = new double [z];
  models = new gaussianmodel [z];
  for(int i = 0; i < z; i++)
    {
      models[i].init(params);
      zeromodel(models[i]);
    }
  gamma = new double * [datapts];
  for(int i = 0; i < datapts; i++)
    gamma[i] = new double [z];

  matrix mx;
  matrix mt;
  mx = data[0];
  for(int i = 1; i < datapts; i++)
    mx.add( data[i] );
  mx = mx * ( 1.0 / (double) datapts );

  matrix sd;
  sd.init( mx.rows(), mx.cols() );
  sd.zero();
  for(int i = 0; i < datapts; i++)
    {
      mt = mx - data[i];
      for( int j = 0; j < mx.rows(); j++ )
        {
          double z;
          z = sd.get(j, 0) + pow(mt.get(j,0), 2);
          sd.set(z, j, 0);
        }
    }

  sd.divide( datapts );

  std::vector< double* > kmdpts;
  for( int i = 0; i < datapts; i++ )
    kmdpts.push_back( data[i].dataptr() );
  kmeans kms;
#ifdef __DEBUG_EM__
  std::cout << "init kmeans...\n"
            << std::flush;
#endif
  kms.loadnocopy( kmdpts, params );
#ifdef __DEBUG_EM__
  std::cout << "kmeans clustering\n"
            << std::flush;
#endif
  kms.cluster( nmodels );

  mt.initidentity( 1, params );
  /*
  for( int i = 0; i < params; i++ )
    {
      mt.set( sd.get(i, 0), i, i );
    }
  */

#ifdef __DEBUG_EM__
  std::cout << "variance:\n" << mt << "\n";
#endif

  for(int i = 0; i < z; i++)
    {
      matrix mm;
      mm.init( kms.center(i), params );
#ifdef __DEBUG_EM__
      std::cout << "mean: " << i << "\n"
                << std::flush;
      std::cout << mm << "\n"
                << std::flush;
#endif
      models[i].setmean( mm );
      models[i].setvariance( mt );
    }
#ifdef __DEBUG_EM__
  std::cout << "cluster size:";
  for(int i = 0; i < z; i++)
    {
      std::cout << kms.clustersize(i) << " ";
    }
  std::cout << "\n";
#endif
  double div;
  div = 0;
  for(int i = 0; i < z; i++)
    {
      pi[i] = kms.clustersize( i );
      div = div + pi[i];
    }
  for(int i = 0; i < z; i++)
    pi[i] = pi[i]/div;

  vmt = mt;
#ifdef __DEBUG_EM__
  std::cout << "init finished\n";
#endif
  return 0;
}

int emcluster::zeromodel( gaussianmodel &m )
{
  m.zero();
  return 0;
}

std::vector< gaussianmodel* > emcluster::getmodels()
{
  std::vector< gaussianmodel* > ml;
  for(int i = 0; i < nmodels; i++)
    {
      ml.push_back( new gaussianmodel );
      matrix m, c;
      models[i].getmean(m);
      models[i].getvariance(c);
      ml[i]->init( params );
      ml[i]->setmean(m);
      ml[i]->setvariance(c);
    }
  return ml;
}

int emcluster::printmodels()
{
  for(int i = 0; i < nmodels; i++)
    {
      matrix m, c;
      models[i].getmean(m);
      models[i].getvariance(c);
      std::cout << "===========================\n";
      std::cout << "pi: " << pi[i] << " " << pi[i]*datapts << "\n";
      std::cout << m << "\n";
      std::cout << c << "\n";
      std::cout << "---------------------------\n"
                << std::flush;
    }
  return 0;
}
