#include "naivebayes.h"

naivebayes::naivebayes()
{
}

naivebayes::~naivebayes()
{
}

int naivebayes::train( std::vector<double*> &samples, int fs )
{
  nbmodel.resize( fs );
  featuresize = fs;
  trainingsize = samples.size();
  for( int i = 0; i < fs; i++ )
    {
      nbmodel[i].mean = nbmodel[i].variance = 0;
    }
  calcmean( samples );
  calcvariance( samples );
  return 0;
}

int naivebayes::calcmean( std::vector<double*> &d )
{
  for( int i = 0; i < trainingsize; i++ )
    {
      for( int j = 0; j < featuresize; j++)
        {
          nbmodel[j].mean = 
            nbmodel[j].mean + d[i][j];
        }
    }
  for( int j = 0; j < featuresize; j++)
    {
      nbmodel[j].mean = 
        nbmodel[j].mean/( (double)trainingsize );
    }
  return 0;
}

int naivebayes::calcvariance( std::vector<double*> &d )
{
  for( int i = 0; i < trainingsize; i++ )
    for( int j = 0; j < featuresize; j++)
      nbmodel[j].variance = nbmodel[j].variance + 
        pow( nbmodel[j].mean - d[i][j], 2 );
  for( int j = 0; j < featuresize; j++)
    nbmodel[j].variance = 
      nbmodel[j].variance/( (double)trainingsize - 1.0 );
  for( int j = 0; j < featuresize; j++)
    nbmodel[j].prefix = 1.0/sqrt( nbmodel[j].variance*2.0*PI );
  return 0;
}

long double naivebayes::prob( double *sample )
{
  long double res;
  res = 0;
  for( int i = 0; i < featuresize; i++)
    {
      res = res + log( nbmodel[i].prefix*exp( -pow(sample[i] - nbmodel[i].mean, 2)/(2.0*nbmodel[i].variance )));
    }
  return res;
}

long double naivebayes::prob( std::vector<double> &sample )
{
  long double res;
  res = 0;
  for( int i = 0; i < featuresize; i++)
    {
      res = res + log( nbmodel[i].prefix*exp( -pow(sample[i] - nbmodel[i].mean, 2)/(2.0*nbmodel[i].variance )));
    }
  return res;
}

int naivebayes::print( )
{
  long double res;
  res = 1;
  for( int i = 0; i < featuresize; i++)
    std::cout << "====\n" << "  " << nbmodel[i].mean 
              << "\n  " << nbmodel[i].variance << "\n";
  return res;
}
