#include "kldivergence.h"

double kldivergence::calc( std::vector< double > p,
                           std::vector< double > q,
                           double delta )
{
  genprobs( p, q, delta );
  double res;
  res = 0;
  for( int i = 0; i < steps; i++ )
    {
      if( probp[i] > 0 )
        {
          res = res + probp[i]*log(probp[i]/proba[i]);
        }
    }
  for( int i = 0; i < steps; i++ )
    {
      if( probq[i] > 0 )
        {
          res = res + probq[i]*log(probq[i]/proba[i]);
        }
    }
  return res/2.0;
}


int kldivergence::genprobs( std::vector<double> &p,
                            std::vector<double> &q,
                            double delta )
{
#ifdef __DEBUG__
  std::cout << "generating probabilities\n";
#endif
  min = p[0];
  max = q[0];
  for( unsigned i = 0; i < p.size(); i++ )
    {
      if( min > p[i] )
        min = p[i];
      if( max < p[i] )
        max = p[i];
    }
  for( unsigned i = 0; i < q.size(); i++ )
    {
      if( min > q[i] )
        min = q[i];
      if( max < q[i] )
        max = q[i];
    }
#ifdef __DEBUG__
  std::cout << "min max " << min << " " << max << "\n";
#endif
  steps = ceil( ( max - min ) / delta );
  scale =  ( (double) steps - 1 ) / (max - min);

  for( unsigned i = 0; i < p.size(); i++ )
    p[i] = scale*(p[i]-min);
  for( unsigned i = 0; i < q.size(); i++ )
    q[i] = scale*(q[i]-min);
#ifdef __DEBUG__
  std::cout << "number of steps " << steps << "\n";
#endif
  probp.resize( steps );
  probq.resize( steps );
  proba.resize( steps );
  for( int i = 0; i < steps; i++ )
    probq[i] = probp[i] = proba[i] = 0;

  calcprob( p, probp );
  calcprob( q, probq );

  double t;
  t = 0;
  for( int i = 0; i < steps; i++ )
    {
      t = t + probq[i];
    }
  if( fabs( t - 1 ) > 0.000001 )
    {
      std::cout << "ERROR probq failed\n";
      std::cout << t << "\n";
      exit(0);
    }
  t = 0;
  for( int i = 0; i < steps; i++ )
    {
      t = t + probp[i];
    }
  if( fabs( t - 1 ) > 0.000001 )
    {
      std::cout << "ERROR probp failed\n";
      std::cout << t << "\n";
      exit(0);
    }
  for( int i = 0; i < steps; i++ )
    proba[i] = probq[i]*0.5+probp[i]*0.5;
  return 0;
}

int kldivergence::calcprob( std::vector< double > &d, 
                            std::vector< double > &pmf )
{
  for( unsigned i = 0; i < pmf.size(); i++ )
    pmf[i] = 0;
  for( unsigned i = 0; i < d.size(); i++ )
    {
      int x;
      double iptr,fptr;
      fptr = modf( d[i], &iptr );
      x = (int)iptr;
      if( x >= pmf.size() )
        {
          std::cout << "ERROR over flow\n";
          exit(0);
        }
      pmf[x] = pmf[x] + ( 1.0 - fptr );
      x = x + 1;
      if( x < pmf.size() )
        pmf[x] = pmf[x] + ( fptr );
    }
  double z;
  z = 0;
  for( unsigned i = 0; i < pmf.size(); i++ )
    z = z + pmf[i];
  for( unsigned i = 0; i < pmf.size(); i++ )
    pmf[i] = pmf[i]/z;
  return 0;
}
