// file:        sift-driver.cpp
// author:      Andrea Vedaldi
// description: SIFT command line utility implementation

// AUTORIGTHS

#include<sift.hpp>

#include<string>
#include<iostream>
#include<iomanip>
#include<fstream>
#include<sstream>
#include<algorithm>

extern "C" {
#include<getopt.h>
#if defined (VL_MAC)
#include<libgen.h>
#else
#include<string.h>
#endif
#include<assert.h>
}

using namespace std ;

size_t const not_found = numeric_limits<size_t>::max() - 1 ;

/** @brief Case insensitive character comparison
 **
 ** This predicate returns @c true if @a a and @a b are equal up to
 ** case.
 **
 ** @return predicate value.
 **/
inline
bool ciIsEqual(char a, char b)
{
  return 
    tolower((char unsigned)a) == 
    tolower((char unsigned)b) ;
}

/** @brief Case insensitive extension removal
 **
 ** The function returns @a name with the suffix $a ext removed.  The
 ** suffix is matched case-insensitve.
 **
 ** @return @a name without @a ext.
 **/
string
removeExtension(string name, string ext)
{
  string::iterator pos = 
    find_end(name.begin(),name.end(),ext.begin(),ext.end(),ciIsEqual) ;

  // make sure the occurence is at the end
  if(pos+ext.size() == name.end()) {
    return name.substr(0, pos-name.begin()) ;
  } else {
    return name ;
  }
}

// keypoint list
typedef vector<pair<VL::Sift::Keypoint,VL::float_t> > Keypoints ;

// predicate used to order keypoints by increasing scale
bool cmpKeypoints (Keypoints::value_type const&a,
		   Keypoints::value_type const&b) {
  return a.first.sigma < b.first.sigma ;
}

// -------------------------------------------------------------------
//                                                                main
// -------------------------------------------------------------------
int
main(int argc, char** argv)
{
  int    first          = -1 ;
  int    octaves        = -1 ;
  int    levels         = 3 ;
  float  threshold      = 0.04f / levels / 2.0f ;
  float  edgeThreshold  = 10.0f;
  int    nodescr        = 0 ;
  int    noorient       = 0 ;
  int    stableorder    = 0 ;
  int    savegss        = 0 ;
  int    verbose        = 0 ;
  int    binary         = 0 ;
  int    haveKeypoints  = 0 ;
  string outputFilenamePrefix ;
  string outputFilename ;
  string descriptorsFilename ;
  string keypointsFilename ;

  static struct option longopts[] = {
    { "verbose",         no_argument,            NULL,              'v' },
    { "help",            no_argument,            NULL,              'h' },
    { "output",          required_argument,      NULL,              'o' },
    { "prefix",          required_argument,      NULL,              'p' },
    { "first-octave",    required_argument,      NULL,              'f' },
    { "keypoints",       required_argument,      NULL,              'k' },
    { "octaves",         required_argument,      NULL,              'O' },
    { "levels",          required_argument,      NULL,              'S' },
    { "threshold",       required_argument,      NULL,              't' },
    { "edge-threshold",  required_argument,      NULL,              'e' },
    { "binary",          no_argument,            NULL,              'b' }, 
    { "no-descriptors",  no_argument,            &nodescr,          1   },
    { "no-orientations", no_argument,            &noorient,         1   },
    { "stable-order",    no_argument,            &stableorder,      1   },
    { "save-gss",        no_argument,            &savegss,          1   },
    { NULL,              0,                      NULL,              0   }
  };
  
  int ch ;
  try {
    while ( (ch = getopt_long(argc, argv, "hvftebk:p:o:O:S:", longopts, NULL)) != -1) {
      switch (ch) {

      case '?' :
        VL_THROW("Invalid option '"<<argv[optind-1]<<"'.") ;
        break;
        
      case ':' :
        VL_THROW("Missing option argument for '"<<argv[optind-1]<<"'.") ;
        break;
        
      case 'h' :
        std::cout
          << argv[0] << " [--verbose|=v] [--help|-h]" << endl
	  << "     [--output|-o NAME] [--prefix|-p PREFIX] [--binary|-b] [--save-gss] " << endl
          << "     [--no-descriptors] [--no-orientations] " << endl
          << "     [--levels|-S NUMBER] [--octaves|-O NUMBER] [--first-octave|-f NUMBER] " << endl
          << "     [--threshold|-t NUMBER] [--edge-threshold|-e NUMBER] " << endl
          << "     IMAGE [IMAGE2 ...]" << endl
          << endl
	  << "* Options *" <<endl
          << " --verbose             Be verbose"<<endl
          << " --help                Print this message"<<endl
          << " --output=NAME         Write to this file"<<endl
	  << " --prefix=PREFIX       Derive output filename prefixing this string to the input file"<<endl
          << " --binary              Write descriptors to a separate file in binary format"<<endl
	  << " --keypoints=FILE      Reads keypoint frames from here; do not run SIFT detector" << endl
          << " --save-gss            Save Gaussian scale space on disk" << endl
          << " --octaves=O           Number of octaves" << endl
          << " --levels=S            Number of levels per octave" << endl
          << " --first-octave=MINO   Index of the first octave" << endl
          << " --threshold=THR       Keypoint strength threhsold" << endl
          << " --edge-threshold=THR  On-edge threshold" << endl 
          << " --no-descriptors      Do not compute descriptors" << endl
          << " --no-orientations     Do not compute orientations" << endl
	  << " --stable-order        Do not reorder keypoints" << endl
	  << endl
	  << " * Examples *" << endl
	  << argv[0] << " [OPTS...] image.pgm" << endl
	  << argv[0] << " [OPTS...] image.pgm --output=file.key" << endl
	  << argv[0] << " [OPTS...] image.pgm --keypoints=frames.key" << endl
	  << argv[0] << " [OPTS...] *.pgm --prefix=/tmp/" << endl
	  << argv[0] << " [OPTS...] *.pgm --prefix=/tmp/ --binary" << endl
	  << endl
	  << " * This build: " ;
#if defined VL_USEFASTMATH
	std::cout << "has fast approximate math" ;
#else
	std::cout << "has slow accurate math" ;
#endif
	std::cout << " (fp datatype is '"
		  << VL_EXPAND_AND_STRINGIFY(VL_FASTFLOAT)
		  << "') *"<<endl ;
        return 0 ;
	
      case 'v' : // verbose
        verbose = 1 ;
        break ;
        
      case 'f': // first octave
        {
          std::istringstream iss(optarg) ;
          iss >> first ;
          if( iss.fail() )
            VL_THROW("Invalid argument '" << optarg << "'.") ;
        }
        break ;
        
      case 'O' : // octaves
        {
          std::istringstream iss(optarg) ;
          iss >> octaves ;
          if( iss.fail() )
            VL_THROW("Invalid argument '" << optarg << "'.") ;
          if( octaves < 1 ) {
            VL_THROW("Number of octaves cannot be smaller than one."); 
          }
        }
        break ;
        
      case 'S' : // levels
        {
          std::istringstream iss(optarg) ;
          iss >> levels ;
          if( iss.fail() )
            VL_THROW("Invalid argument '" << optarg << "'.") ;
          if( levels < 1 ) {
            VL_THROW("Number of levels cannot be smaller than one.") ;
          }
        }      
        break ;

      case 't' : // threshold
        {
          std::istringstream iss(optarg) ;
          iss >> threshold ;
          if( iss.fail() )
            VL_THROW("Invalid argument '" << optarg << "'.") ;
        }
        break ;

      case 'e' : // edge-threshold
        {
          std::istringstream iss(optarg) ;
          iss >> edgeThreshold ;
          if( iss.fail() )
            VL_THROW("Invalid argument '" << optarg << "'.") ;
        }
        break ;

      case 'o' : // output filename
        {
          outputFilename = std::string(optarg) ;
          break ;
        }

      case 'p' : // output prefix
        {
          outputFilenamePrefix = std::string(optarg) ;
          break ;
        }

      case 'k' : // keypoint file
	{
	  keypointsFilename = std::string(optarg) ;
	  haveKeypoints = 1 ;
	  break ;
	}

      case 'b' : // write descriptors to a binary file
        {
          binary = 1 ;
          break ;
        }

      case 0 : // all other options
        break ;
        
      default:
        assert(false) ;
      }
    }
    
    argc -= optind;
    argv += optind;

    // check for argument consistency
    if(argc == 0) VL_THROW("No input image specfied.") ;
    if(outputFilename.size() != 0 && (argc > 1 | binary)) {
      VL_THROW("--output cannot be used with multiple images or --binary.") ;
    }

    if(outputFilename.size() !=0 && 
       outputFilenamePrefix.size() !=0) {
      VL_THROW("--output cannot be used in combination with --prefix.") ;
    }
  }
  catch( VL::Exception const & e ) {
    std::cerr<<"error: "<<e.msg<<std::endl ;
    exit(1) ;
  } 

  // -----------------------------------------------------------------
  //                                            Loop over input images
  // -----------------------------------------------------------------      
  while( argc > 0 ) {
    string name(argv[0]) ;

    try {
      VL::PgmBuffer buffer ;
      
      // compute the output filenames:
      //
      // 1) if --output is specified, then we just use the one provided
      //    by the user
      //
      // 2) if --output is not specified, we derive the output filename
      //    from the input filename by
      //    - removing the extension part from the output filename
      //    - and if outputFilenamePrefix is non void, removing 
      //      the directory part and prefixing outputFilenamePrefix.
      //
      // 3) in any case we derive the binary descriptor filename by
      //    removing from the output filename the .key extension (if any)
      //    and adding a .desc extension.
      
      if(outputFilename.size() == 0) {
	// case 2) above
	outputFilename = name ;
	
	// if we specify an output directory, then extract
	// the basename
	if(outputFilenamePrefix.size() != 0) {
	  outputFilename = outputFilenamePrefix + 
	    std::string(basename(outputFilename.c_str())) ;
	}
	
      // remove .pgm extension, add .key
	outputFilename = removeExtension(outputFilename, ".pgm") ;
	outputFilename += ".key" ;
      }
      
      // remove .key extension, add .desc
      descriptorsFilename = removeExtension(outputFilename, ".key") ;
      descriptorsFilename += ".desc" ;
      
      // ---------------------------------------------------------------
      //                                                  Load PGM image
      // ---------------------------------------------------------------    
      verbose && cout<<"Lodaing PGM image '"<<name<<"' ..."<<flush;
      
      try {          
	ifstream in(name.c_str(), ios::binary) ; 
	if(! in.good()) VL_THROW("Could not open '"<<name<<"'.") ;      
	extractPgm(in, buffer) ;
      }    
      catch(VL::Exception const& e) {
	throw VL::Exception("PGM read error: "+e.msg) ;
      }
      
      verbose && cout << " done ("
		      << buffer.width<<" x "<<buffer.height<<")."<<endl ;
      
      // ---------------------------------------------------------------
      //                                            Gaussian scale space
      // ---------------------------------------------------------------    
      if( verbose ) 
	cout << "Computing Gaussian scale space ... " << flush ;
      
      int         O      = octaves ;    
      int const   S      = levels ;
      int const   omin   = first ;
      float const sigman = .5 ;
      float const sigma0 = 1.6 * powf(2.0f, 1.0f / S) ;
      
      // optionally autoselect the number number of octaves
      // we downsample up to 8x8 patches
      if(O < 1) {
	O = std::max
	  (int
	   (std::floor
	    (log2
	     (std::min(buffer.width,buffer.height))) - omin -3), 1) ;
      }
      
      // initialize scalespace
      VL::Sift sift(buffer.data, buffer.width, buffer.height, 
		    sigman, sigma0,
		    O, S,
		    omin, -1, S+1) ;
      
      verbose && cout << "("<<O<<" octaves from "<<omin<<", "
		      <<S<< " levels per octave) done."<<endl ;
      
      // ---------------------------------------------------------------
      //                                       Save Gaussian scale space
      // ---------------------------------------------------------------    
      
      if(savegss) {
	verbose && cout<<"Saving Gaussian scale space."<<endl ;
	
	string imageBasename = removeExtension(outputFilename, ".key") ;
	
	for(int o = omin ; o < omin + O ; ++o) {
	  for(int s = 0 ; s < S ; ++s) {
	    
	    ostringstream suffix ;
	    suffix<<'.'<<o<<'.'<<s<<".pgm" ;
	    string imageFilename = imageBasename + suffix.str() ;
	    
	    verbose && cout << "Saving octave " <<o
			    << " level "<<s
			    << " to '"<<imageFilename<<"' ..."<<flush ;
	    
	    ofstream fout(imageFilename.c_str(), ios::binary) ;
	    if(!fout.good()) 
	      VL_THROW("Could not open '"<<imageFilename<<'\'') ;
	    
	    VL::insertPgm(fout,
			  sift.getLevel(o,s),
			  sift.getOctaveWidth(o),
			  sift.getOctaveHeight(o)) ;
	    fout.close() ;
	    
	    verbose && cout<<" done."<<endl ;
	  }
	}
      }
      
      // ---------------------------------------------------------------
      //                                                   SIFT detector
      // ---------------------------------------------------------------    
      if( ! haveKeypoints ) {
	verbose && cout << "Running SIFT detector (threshold:"<<threshold
			<< ", edge-threshold: "<<edgeThreshold
			<< ") ..."<<flush ;
	
	sift.detectKeypoints(threshold, edgeThreshold) ;
	
	verbose && cout << " done (" 
			<< sift.keypointsEnd() - sift.keypointsBegin() 
			<< " keypoints)." << endl ;
      }
      
      // ---------------------------------------------------------------
      //                               SIFT orientations and descriptors
      // ---------------------------------------------------------------    
      
      if( verbose ) {
	if( ! noorient &&   nodescr) cout << "Computing keypoint orientations" ;
	if(   noorient && ! nodescr) cout << "Computing keypoint descriptors" ;
	if( ! noorient && ! nodescr) cout << "Computing orientations and descriptors" ;
	if(   noorient &&   nodescr) cout << "Finalizing" ; 
	cout<<flush ;
      }
      
      {      
	auto_ptr<ifstream> keypointsIn_pt ;
	auto_ptr<ofstream> descriptorsOut_pt ;
	
	// open output file
	ofstream out(outputFilename.c_str(), ios::binary) ;
	if(!out.good()) VL_THROW("Could not open output file '"<<outputFilename<<"'.") ;
	verbose && cout<<" (writing to '"<<outputFilename<<"')"<<flush ;
	out.flags(ios::fixed) ;
	
	// optionally open keypoints file
	if( haveKeypoints ) {
	  keypointsIn_pt = 
	    auto_ptr<ifstream>(new ifstream(keypointsFilename.c_str(), ios::binary)) ;
	  if(!keypointsIn_pt->good()) 
	    VL_THROW("Could not open keypoints file '"<<keypointsFilename<<"'.") ;
	  if(verbose)
	    cout<<" (reading from '"<<keypointsFilename<<"')"<<flush ;
	}
	
	// optionally open descriptors file
	if( binary ) {
	  descriptorsOut_pt = 
	    auto_ptr<ofstream>
	    (new ofstream(descriptorsFilename.c_str(), ios::binary)) ;
	  if(!descriptorsOut_pt->good())
	    VL_THROW("Could not open descriptors file '"<<descriptorsFilename<<"'.") ;
	  if(verbose)
	    cout << " (writing '"<<descriptorsFilename<<"')"<<flush;
	}
	
	verbose && cout<<" ..."<<flush ;
	
	if( haveKeypoints ) {
	  // -------------------------------------------------------------
	  //                 Reads keypoint from file, compute descriptors
	  // -------------------------------------------------------------
	  Keypoints keypoints ;

	  // readoff keypoints
	  while( !keypointsIn_pt->eof() ) {
	    VL::float_t x,y,sigma,th ;
	   
	    // extract next keypoint from file
	    (*keypointsIn_pt) >> x>>y>>sigma>>th ;
	    (*keypointsIn_pt).ignore(numeric_limits<streamsize>::max(),'\n') ;
	    if(keypointsIn_pt->eof())
	      break ;
	    if(!keypointsIn_pt->good())
	      VL_THROW("Error reading keypoints file") ;

	    // compute integer components
	    VL::Sift::Keypoint key = sift.getKeypoint(x,y,sigma) ;

	    Keypoints::value_type entry ;
	    entry.first  = key ;
	    entry.second = th ;
	    keypoints.push_back(entry) ;
	  }
	  
	  // sort keypoints
	  if(! stableorder)
	    sort(keypoints.begin(), keypoints.end(), cmpKeypoints) ;

	  // process in batch
	  for(Keypoints::const_iterator iter = keypoints.begin() ;
	      iter != keypoints.end() ;
	      ++iter) {
	    VL::Sift::Keypoint const& key = iter->first ;
	    VL::float_t th = iter->second ;

	    // write keypoint
	    out << setprecision(2) << key.x << " "
		<< setprecision(2) << key.y << " "
		<< setprecision(2) << key.sigma << " "
		<< setprecision(3) << th ;
	    
	    // compute descriptor
	    VL::float_t descr [128] ;		
	    sift.computeKeypointDescriptor(descr, key, th) ;
	    VL::uint8_t intDescr [128] ;
	    for(int i = 0 ; i < 128 ; ++i) {
	      intDescr[i] = VL::uint8_t(VL::float_t(512) * descr[i]) ;
	    }
	    if(descriptorsOut_pt.get()) {
	      descriptorsOut_pt->write(reinterpret_cast<char*>(intDescr), 128) ;
	    } else {
	      for(int i = 0 ; i < 128 ; ++i) {
		out << ' '<< uint32_t(intDescr[i]) ;
	      }
	    }
	    out << endl ;    
	  } // next keypoint
	} else {
	  // -------------------------------------------------------------
	  //            Run detector, compute orientations and descriptors
	  // -------------------------------------------------------------
	  for( VL::Sift::KeypointsConstIter iter = sift.keypointsBegin() ;
	       iter != sift.keypointsEnd() ; ++iter ) {
	    
	    // detect orientations
	    VL::float_t angles [4] ;
	    int nangles ;
	    if( ! noorient ) {
	      nangles = sift.computeKeypointOrientations(angles, *iter) ;
	    } else {
	    nangles = 1;
	    angles[0] = VL::float_t(0) ;
	  }
	    
	    // compute descriptors
	    for(int a = 0 ; a < nangles ; ++a) {

	      out << setprecision(2) << iter->x << ' '
		  << setprecision(2) << iter->y << ' '
		  << setprecision(2) << iter->sigma << ' ' 
		  << setprecision(3) << angles[a] ;
	      
	      if( ! nodescr ) {
		VL::float_t descr [128] ;
		sift.computeKeypointDescriptor(descr, *iter, angles[a]) ;
		VL::uint8_t intDescr [128] ;
		for(int i = 0 ; i < 128 ; ++i) {
		  intDescr[i] = VL::uint8_t(VL::float_t(512) * descr[i]) ;
		}
		if(descriptorsOut_pt.get()) {
		  descriptorsOut_pt->write( reinterpret_cast<char*>(intDescr), 128 ) ;
		} else {
		  for(int i = 0 ; i < 128 ; ++i) {
		    out << ' '<< uint32_t(intDescr[i]) ;
		  }
		}
	      }
	      out << endl ;
	    } // next angle
	  } // next keypoint
	}
	
	out.close() ;
	if(descriptorsOut_pt.get()) descriptorsOut_pt->close(); 
	if(keypointsIn_pt.get())    keypointsIn_pt->close(); 
	verbose && cout << " done."<<endl ;
      }
      
      argc-- ;
      argv++ ;
      outputFilename = string("") ;
    }
    catch(VL::Exception &e) {
      cerr<<endl<<"Error processing '"<<name<<"': "<<e.msg<<endl ;
      return 1 ;
    }    
  } // next image
  
  return 0 ;
}