#include "GLCP.h"
#include "dbimage.h"
#include "bdbenv.h"
#include "dbalnimage.h"
#include "dbtmpfeature.h"
#include <tclap/CmdLine.h>
#include "naivebayes.h"
#include "emcluster.h"
#include <string>

int trainnb(std::vector< double* > &samples, int fsz, std::vector< naivebayes* > &nbs);

long double nbsprob( double* samples, std::vector< naivebayes* > &nbs );

int getsamples( dbtmpfeature &dbtf, GLCP::mask &mymask,
                std::vector< double* > &sampels, int &numfeatures);

int checkcorrect( std::vector< naivebayes* > &nba, 
                  std::vector< naivebayes* > &nbb,
                  dbtmpfeature &dbtf, GLCP::mask &mymask,
                  int &correct, int &incorrect );

long double nbsprob( std::vector<double> &samples,
                     std::vector< naivebayes* > &nbs );

int classifyimage( image2d &target, Magick::Image &mim,
                   GLCP::mask &mymask,
                   std::vector< naivebayes* > &nba, 
                   std::vector< naivebayes* > &nbb, std::string outim );

long double nbsprob( std::vector<double> &samples,
                     std::vector< naivebayes* > &nbs )
{
  long double res;
  res = -10e200;
  for( unsigned i = 0; i < nbs.size(); i++ )
    {
      long double z;
      if( nbs[i] != NULL )
        {
          z = nbs[i]->prob( samples );
          if( z > res )
            res = z;
        }
    }
  return res;
}

long double nbsprob( double* samples, std::vector< naivebayes* > &nbs )
{
  long double res;
  res = -10e200;
  for( unsigned i = 0; i < nbs.size(); i++ )
    {
      long double z;
      if( nbs[i] != NULL )
        {
          z = nbs[i]->prob( samples );
          if( z > res )
            res = z;
        }
    }
  return res;
}

int trainnb(std::vector< double* > &samples, int fsz, std::vector< naivebayes* > &nbs)
{
  kmeans kms;
  std::cout << "training with " << fsz << " features\n";
  kms.loadnocopy( samples, fsz );
  kms.cluster( 2 );
  nbs.resize( 2 );
  for( int i = 0; i < 2; i++ )
    {
      std::vector< double* > tset;
      tset.clear();
      for( unsigned z = 0; z < samples.size(); z++ )
        if( kms.classify( samples[z] ) == i )
          tset.push_back( samples[z] );
      if( tset.size() > 10 )
        {
          nbs[i] = new naivebayes;
          std::cout << "training with " << tset.size() << " points\n";
          nbs[i]->train( tset, fsz );
          nbs[i]->print();
        }
      else
        nbs[i] = NULL;
    }
  return 0;
}

int getsamples( dbtmpfeature &dbtf, GLCP::mask &mymask,
                std::vector< double* > &samples,
                int &numfeatures )
{
  GLCP myglcp;
  char *buff;
  int bufsz;
  unsigned int maxid;
  maxid = dbtf.maxid();
  dbtf.getrecsize( bufsz );
  buff = new char[bufsz];
  std::vector< double > meanvar; // keep track of average mean
  std::vector< double > meanmean;
  for( unsigned i = 1; i < maxid; i++ )
    {
      dbtf.get( i, buff );
      myglcp.del();
      myglcp.deserialize( buff, bufsz );
      std::vector< double > tfm, tfv;
      tfm.clear();
      myglcp.meanthetadelta( mymask, tfm );
      if( i > 1 && numfeatures != (int)tfm.size() )
        {
          std::cout << "ERROR feature size not consistant\n";
          exit(0);
        }
      numfeatures = (int)tfm.size();
      double *z;
      z = new double[ tfm.size() ];
      for( unsigned j = 0; j < tfm.size(); j++ )
        z[j] = tfm[j];
      samples.push_back( z );
      tfv.clear();
      myglcp.varthetadelta( mymask, tfv );
      if( i == 1 )
        {
          meanvar = tfv;
          meanmean = tfm;
        }
      else
        {
          for( unsigned int z = 0; z < tfv.size(); z++ )
            meanvar[z] = meanvar[z] + tfv[z];
          for( unsigned int z = 0; z < tfm.size(); z++ )
            meanmean[z] = meanmean[z] + tfm[z];
            
        }
      tfm.clear();
      tfv.clear();
    }
  std::cout << "mean mean: ";
  for( unsigned i = 0; i < meanmean.size(); i++ )
    std::cout << meanmean[i]/ ( (double)(maxid -1) ) << " ";
  std::cout << "\n";
  std::cout << "mean variance: ";
  for( unsigned i = 0; i < meanvar.size(); i++ )
    std::cout << meanvar[i]/ ( (double)(maxid -1) ) << " ";
  std::cout << "\n";
  delete [] buff;
  return 0;
}

int checkcorrect( std::vector< naivebayes* > &nba, 
                  std::vector< naivebayes* > &nbb,
                  dbtmpfeature &dbtf, GLCP::mask &mymask,
                  int &correct, int &incorrect )
{
  GLCP myglcp;
  char *buff;
  int bufsz;
  unsigned int maxid;
  maxid = dbtf.maxid();
  dbtf.getrecsize( bufsz );
  buff = new char[bufsz];
  int numfeatures;
  numfeatures = 0;
  for( unsigned i = 1; i < maxid; i++ )
    {
      dbtf.get( i, buff );
      myglcp.del();
      myglcp.deserialize( buff, bufsz );
      std::vector< double > tf;
      tf.clear();
      myglcp.meanthetadelta( mymask, tf );
      if( i > 1 && numfeatures != (int)tf.size() )
        {
          std::cout << "ERROR feature size not consistant\n";
          exit(0);
        }
      numfeatures = (int)tf.size();
      double *z;
      z = new double[ tf.size() ];
      for( unsigned j = 0; j < tf.size(); j++ )
        z[j] = tf[j];
      long double sa, sb;
      sa = nbsprob( z, nba );
      sb = nbsprob( z, nbb );
      if( sa >= sb )
        correct++;
      else
        incorrect++;
      delete [] z;
    }
  delete [] buff;
  return 0;
}

int classifyimage( image2d &target, Magick::Image &mim, GLCP::mask &mymask,
                   std::vector< naivebayes* > &nba, 
                   std::vector< naivebayes* > &nbb, std::string outim )
{
  GLCP ss;
  std::cout << "starting classification: " << target.getheight() 
            << "x" << target.getwidth() << "\n";
  for( int i = 7; i < target.getheight(); i = i + 8 )
    {
      for( int j = 7; j < target.getwidth(); j = j + 8 )
        {
          if( ss.init( &target, i, j, 14 ) == 0 )
            {
              Magick::ColorRGB rgb("red");
              std::vector< double > myts;
              myts.clear();
              ss.meanthetadelta( mymask, myts );
              double g, b;
              g =  nbsprob( myts, nba );
              b =  nbsprob( myts, nbb );
              if( b > g  )
                rgb.red(0);
              else
                rgb.red(1);
              mim.pixelColor( j, i, rgb);
            }
        }
      std::cout << "." << std::flush;
    }
  std::cout << "\n";
  mim.write( outim.c_str() );
  return 0;
}

int main( int argc, char *argv[] )
{

#ifndef _NOMAGICK_
  try
    {
      ///set up the command line parsing
      TCLAP::CmdLine cmd("Image alignment program", ' ', "0.001");

      TCLAP::ValueArg<std::string> dbdir("d", "db", "database dirctory", true,
                                         "", "databaase");

      TCLAP::ValueArg<int> aid("a", "aid", "tile id", true, 0, "id");
      TCLAP::ValueArg<int> si("i", "id", "image id", true, 0, "0" );
      TCLAP::ValueArg<int> ch("c", "channel", "channel", true, 0, "0" );

      TCLAP::ValueArg< bool > fmean("m", "mean", "mean flag", false,
                                    true, "true/false" );
      TCLAP::ValueArg< bool > fentropy("e", "entropy", "entropy flag", false,
                                       true, "true/false" );
      TCLAP::ValueArg< bool > fcontrast("t", "contrast", "contrast flag", false,
                                        true, "true/false" );
      TCLAP::ValueArg< bool > fcorrelation("r", "correlation", "correlation flag",
                                           false, true, "true/false" );
      TCLAP::ValueArg< bool > fpower("p", "power", "power flag", false,
                                     true, "true/false" );

      cmd.add(dbdir);
      cmd.add(aid);
      cmd.add(si);
      cmd.add(ch);
      cmd.add( fmean );
      cmd.add( fentropy );
      cmd.add( fcontrast );
      cmd.add( fcorrelation );
      cmd.add( fpower );

      bdbenv bdb;
      dbtmpfeature dbtf_g, dbtf_b;
      dbalnimage dba;

      cmd.parse(argc, argv);

      GLCP::mask mymask;
      if( fmean.getValue() == true )
        mymask.mean = 1;
      else
        mymask.mean = 0;

      if( fentropy.getValue() == true )
        mymask.entropy = 1;
      else
        mymask.entropy = 0;

      if( fcontrast.getValue() == true )
        mymask.contrast = 1;
      else
        mymask.contrast = 0;

      if( fcorrelation.getValue() == true )
        mymask.correlation = 1;
      else
        mymask.correlation = 0;

      if( fpower.getValue() == true )
        mymask.power = 1;
      else
        mymask.power = 0;

      dbtf_g.setname( "good.db" );
      dbtf_b.setname( "bad.db" );

      bdb.open( dbdir.getValue() );
      dbtf_g.init( bdb.envptr(), 0 );
      dbtf_b.init( bdb.envptr(), 0 );
      dba.init( bdb.envptr() );

      bdb.adddb( &dba );
      bdb.adddb( &dbtf_g );
      bdb.adddb( &dbtf_b );
      bdb.cleanlog();

      std::vector< double* > samples;

      int numfeatures;
      samples.clear();
      getsamples( dbtf_g, mymask, samples, numfeatures );
      //////////////////////////////////////////////////////////////

      std::vector< naivebayes* > nb_g, nb_b;
      trainnb( samples, numfeatures, nb_g );

      for( unsigned int i = 0; i < samples.size(); i++ )
        delete [] samples[i];
      samples.clear();

      getsamples( dbtf_b, mymask, samples, numfeatures );

      trainnb( samples, numfeatures, nb_b );

      for( unsigned int i = 0; i < samples.size(); i++ )
        delete [] samples[i];

      int gg, gb, bg, bb;
      gg = gb = bg = bb = 0;
      checkcorrect( nb_g, nb_b, dbtf_g, mymask, gg, gb );
      checkcorrect( nb_b, nb_g, dbtf_b, mymask, bg, bb );
      std::cout << gg << " : " << gb << "\n";
      std::cout << bg << " : " << bb << "\n";

      int myaid;
      std::cout << "checking aid\n";
      myaid = aid.getValue();

      if( dba.checkaid( myaid ) == 1 )
        {
          std::set<int> channels;
          std::cout << "getting channels\n";
          dba.get_channels(myaid, channels );
          std::cout << "checking for valid channel\n";
          int mych;
          mych = ch.getValue();

          if( channels.count( mych ) > 0 )
            {
              image2d im;
              dba.get_aligned_image( myaid, si.getValue(), mych, &im );
              im.normalize( 0, 256.0/4096.0 );
              GLCP ss;
              int sz;
              Magick::Image mim;
              im.getimage( mim );
              classifyimage( im, mim, mymask, nb_g, nb_b, "output.png" );
            }
        }

      dba.close();
      dbtf_g.close();
      dbtf_b.close();
      bdb.close();
    }
  catch (TCLAP::ArgException &e)
    {
      std::cout << "ERROR: " << e.error() 
                << " for arg " << e.argId() << "\n";
      return 0;
    }
#else
  std::cout << "compiled with no image magick\n";
#endif
  return 0;
}
