#ifndef _NOMAGICK_
#include <Magick++.h>
#endif

#include <iostream>
#include <string>
#include <fstream>
#include <math.h>
#include "imagealign.h"
#include "image2d.h"
#include <fftw3.h>
#include <tclap/CmdLine.h>
#include "bdbenv.h"
#include "dbimage.h"
#include "pso.h"
#include "histequalize.h"
#include "tile.h"
#include "tileimage.h"
#include "alignimage.h"
#include "expat.h"
#include "common.h"

struct alntrans
{
  int a;
  int b;
  double score;
  int tx;
  int ty;
};

struct xmldata
{
  std::vector< int > tiles;
  std::vector< alntrans > connections;
  std::vector< std::string > tmpstr;
  std::vector< std::string > tag;
  alntrans curtrans;
};

int addtile( int x, xmldata* d);

int readfile( std::string xmlout, 
              std::vector< tile * > &tiles,
              std::vector< alignimage::setconn2 > &conn,
              dbalnimage *adb );

int dotout( std::vector< tile * > &tiles,
            std::vector< alignimage::setconn2 > &conn );

void xmldata_handle(void *userData,
                    const XML_Char *s,
                    int len);

void xmlstart_handle( void *userData,
                      const XML_Char *name,
                      const XML_Char **atts );

void xmlend_handle( void *data, const char *el );

int createconnections( xmldata &xmld, std::vector< tile* > &tiles, 
                       std::vector< alignimage::setconn2 > &conn,
                       dbalnimage *adb );

int get_edges( int it1, int it2,
               std::vector< alntrans > &connections,
               std::vector< alntrans > &out );

int get_transform( std::vector< alntrans > &connections,
                   double sd, double count,
                   alignimage::setconn2 &sc );

int calc_avg_sd( std::vector< alntrans > &connections,
                 double &avg, double &sd, double &count);

int addtile( int x, xmldata* d)
{
  for( unsigned i = 0; i < d->tiles.size(); i++ )
    if( d->tiles[i] == x )
      return 0;
  d->tiles.push_back( x ); 
  return 0;
}

void xmldata_handle(void *userData,
                    const XML_Char *s,
                    int len)
{
  char *buffer;
  buffer = new char[len+1];
  buffer[len] = 0;
  memcpy( buffer, s, len );
  xmldata *x;
  x = (xmldata*)userData;
  std::string ts;
  ts = buffer;
  stripstr( ts );
  x->tmpstr[ x->tmpstr.size()-1 ].append( ts );
  delete [] buffer;
}

void xmlstart_handle( void *userData,
                      const XML_Char *name,
                      const XML_Char **atts )
{
  xmldata *x;
  x = (xmldata*)userData;
  x->tmpstr.push_back( "" );
  x->tag.push_back( name );
  if( strcmp(name, "alignment") == 0 )
    {
      x->curtrans.a = x->curtrans.b = -1;
    }
  else if( strcmp( name, "transform" ) == 0 )
    {
      int i = 0;
      while( atts[i] != NULL )
        {
          if( strcmp(atts[i], "x") == 0 )
            x->curtrans.tx = strtol(atts[i+1], NULL, 10);
          else if(strcmp(atts[i], "y") == 0 )
            x->curtrans.ty = strtol(atts[i+1], NULL, 10);
          i = i + 2;
        }
    }
}

void xmlend_handle( void *data, const char *el )
{
  xmldata *x;
  x = (xmldata*)data;
  // create rows
  if( strcmp( el, "tilea" ) == 0 )
    {
      std::vector< int > nums;
      converttoint( x->tmpstr.back(), nums );
      x->curtrans.a = nums[0];
      addtile( nums[0], x );
    }
  else if( strcmp( el, "tileb" ) == 0 )
    {
      std::vector< int > nums;
      converttoint( x->tmpstr.back(), nums );
      x->curtrans.b = nums[0];
      addtile( nums[0], x );
    }
  else if( strcmp( el, "transform" ) == 0 )
    {
      x->curtrans.score = strtod( x->tmpstr[ x->tmpstr.size()-1 ].c_str(), NULL );
      x->connections.push_back( x->curtrans );
    }
  x->tmpstr.pop_back();
  x->tag.pop_back();
}

int dotout( std::vector< tile * > &tiles,
            std::vector< alignimage::setconn2 > &conn )
{
  std::cout << "digraph align{\nrankdir=LR\n";
  for( unsigned i = 0; i < tiles.size(); i++ )
    {
      std::cout << "N" << i 
                << " [shape=\"box\" label=\""
                << tiles[i]->getaid() << "\"];\n";
    }
  for( unsigned i = 0; i < conn.size(); i++ )
    {
      std::cout << "N" << conn[i].a << " -> N"
                << conn[i].b << " [ label=\""
                << conn[i].con.lowtx << " "
                << conn[i].con.lowty << "\"];\n";
    }
  std::cout << "}\n";
  return 0;
}

int readfile( std::string xmlout, 
              std::vector< tile * > &tiles,
              std::vector< alignimage::setconn2 > &conn,
              dbalnimage *adb )
{
  conn.clear();
  tiles.clear();
  XML_Parser p = XML_ParserCreate(NULL);

  XML_SetElementHandler(p, xmlstart_handle,
                        xmlend_handle);
  XML_SetCharacterDataHandler( p,
                               xmldata_handle );
  xmldata xmld;
  FILE *fp;
  XML_SetUserData( p, &xmld );
  fp = fopen( xmlout.c_str(), "r" );
  char buf[1024];
  while(1) 
    {
      int done;
      int len;
      
      len = fread(buf, 1, 1024, fp );
      done = feof(fp);
      
      if (! XML_Parse(p, buf, len, done)) {
        fprintf(stderr, "Parse error at line %d:\n%s\n",
                (int)XML_GetCurrentLineNumber(p),
                XML_ErrorString( XML_GetErrorCode(p)) );
        exit(-1);
      }

      if (done)
        break;
    }
  fclose(fp);
  XML_ParserFree( p );
  createconnections( xmld, tiles, conn, adb );
  return 0;
}

int createconnections( xmldata &xmld, std::vector< tile* > &tiles, 
                       std::vector< alignimage::setconn2 > &conn,
                       dbalnimage *adb )
{
  std::map< int, int > mapping;
  conn.clear();
  tiles.clear();
  for( std::vector< int >::iterator it = xmld.tiles.begin();
       it != xmld.tiles.end(); it ++ )
    {
      tile *t;
      t = new tile;
      t->init( *it, adb );
      mapping[ *it ] = tiles.size(); // save the tile coor to array coor mapping
      tiles.push_back(t);
    }
  for( std::vector< int >::iterator it1 = xmld.tiles.begin();
       it1 != xmld.tiles.end(); it1 ++ )
    {
      for( std::vector< int >::iterator it2 = xmld.tiles.begin();
           it2 != xmld.tiles.end(); it2 ++ )
        {
          double avg, sd;
          double count;
          int ret;
          std::vector< alntrans > filterededges;
          get_edges( *it1, *it2, xmld.connections, filterededges );
          if( filterededges.size() > 0 )
            {
              ret = calc_avg_sd( filterededges, avg, sd, count ); 
              if( ret == 0 )
                {
                  alignimage::setconn2 sc;
                  ret = get_transform( filterededges, sd, count, sc );
                  if( ret == 0 )
                    {
                      sc.a = mapping[ *it1 ];
                      sc.b = mapping[ *it2 ];
                      conn.push_back( sc );
                    }
                }
            }
        }
    }

  return 0;
}

int get_edges( int it1, int it2,
               std::vector< alntrans > &connections,
               std::vector< alntrans > &out )
{
  out.clear();
  for( unsigned i = 0; i < connections.size(); i++ )
    if( connections[i].a == it1 && connections[i].b == it2 )
      out.push_back( connections[i] );
  return 0;
}

int get_transform( std::vector< alntrans > &connections,
                   double sd, double count,
                   alignimage::setconn2 &sc )
{
  double max;
  int maxid;
  max = 0;
  maxid = -1;
  for( unsigned i = 0; i < connections.size(); i++ )
    {
      if( connections[i].score > max )
        {
          max = connections[i].score;
          maxid = i;
        }
    }
  // if there is a valid
  if( maxid != -1 )
    {
      int candidates;
      candidates = 0;
      for( unsigned i = 0; i < connections.size(); i++ )
        if( connections[i].score >= max - sd )
          candidates++;
      if( candidates == 1 || count == 1 )
        {
          sc.con.lowtx = sc.con.hightx = connections[maxid].tx;
          sc.con.lowty = sc.con.highty = connections[maxid].ty;
          return 0;
        }
    }
  return -1;
}

int calc_avg_sd( std::vector< alntrans > &connections,
                 double &avg, double &sd, double &count)
{
  avg = sd = count = 0;
  for( unsigned i = 0; i < connections.size(); i++ )
    {
      avg = avg + connections[i].score;
      count = count + 1;
    }
  if( count > 0 )
    {
      avg = avg / count;
      for( unsigned i = 0; i < connections.size(); i++ )
        {
          sd = sd + pow( avg - connections[i].score, 2);
          count = count + 1;
        }
      sd = sqrt( sd/count );
      return 0;
    }
  return -1;
}

int main( int argc, char *argv[] )
{
  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<std::string> xmlout( "x", "xml", "intput xml file", true, "", "xml file");

#ifndef _NOMAGICK_
      Magick::Image mim;
#endif

      cmd.add(dbdir);
      cmd.add(xmlout);

      cmd.parse(argc, argv);

      bdbenv bdb;
      dbimage imstore;
      dbalnimage dba;

      bdb.open( dbdir.getValue() );
      imstore.init( bdb.envptr() );
      dba.init( bdb.envptr() );
      bdb.adddb( &imstore );
      bdb.adddb( &dba );
      bdb.cleanlog();
      alignimage aim;
      aim.init( &dba );

      std::vector< tile * > tiles;
      std::vector< alignimage::setconn2 > conn;

      readfile( xmlout.getValue(), 
                tiles, conn, &dba );
      dotout( tiles, conn );
      aim.mergesets( tiles, conn, 1 );
      for( unsigned i = 0; i < tiles.size(); i++ )
        delete tiles[i];
      dba.close();
      imstore.close();
      bdb.close();
    }
  catch (TCLAP::ArgException &e)
    {
      std::cout << "ERROR: " << e.error() 
                << " for arg " << e.argId() << "\n";
      return 0;
    }
  return 0;
}
