// PatMaN DNA pattern matcher
// (C) 2007 Kay Pruefer, Udo Stenzel
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or (at
// your option) any later version.  See the LICENSE file for details.


#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <cstdlib>
#include <cstring>

#include <popt.h>

#include "prefix_tree.h"
#include "fasta.h"

using namespace std ;

// global variables
size_t cutoff = 2 ;
size_t allow_gaps = 0 ;
size_t minlength = 0 ;
long long total_nodes = 0 ;
int only_plus_strand = 0 ;
int discount_adenine = 0 ;
int ambi_codes = 0 ; 
int chop5 = 0, chop3 = 0 ;
int debug_flags = debug_notquiet ;
int do_prefetch = 3 ;
ostream *output = 0 ;

long position_genome ;
const char* genome_name ;

enum option_flag { opt_none, opt_db, opt_pat, opt_output, opt_version } ;

struct poptOption options[] = {
	{ "version",      'V', POPT_ARG_NONE,                             0,       opt_version, "Print version number and exit", 0 },
	{ "edits",        'e', POPT_ARG_INT  | POPT_ARGFLAG_SHOW_DEFAULT, &cutoff,           0, "Set maximum edit distance to N", "N" },
	{ "gaps",         'g', POPT_ARG_INT  | POPT_ARGFLAG_SHOW_DEFAULT, &allow_gaps,       0, "Set maximum number of gaps to N", "N" },
	{ "databases",    'D', POPT_ARG_NONE,                             0,            opt_db, "Following files are databases", 0 },
	{ "patterns",     'P', POPT_ARG_NONE,                             0,           opt_pat, "Following files contain patterns", 0 },
	{ "output",       'o', POPT_ARG_STRING,                           0,        opt_output, "Write output to FILE", "FILE" },
	{ "ambicodes",    'a', POPT_ARG_NONE,                             &ambi_codes,       0, "Interpret ambiguity codes in patterns", 0 },
	{ "singlestrand", 's', POPT_ARG_NONE,                             &only_plus_strand, 0, "Do not match reverse-complements", 0 },
	{ "prefetch",     'p', POPT_ARG_INT  | POPT_ARGFLAG_SHOW_DEFAULT, &do_prefetch,      0, "Prefetch N nodes", "N" },
	{ "min-length",   'l', POPT_ARG_INT  | POPT_ARGFLAG_DOC_HIDDEN,   &minlength,        0, "Only consider patterns longer than N", "N" },
	{ "chop-3",       'x', POPT_ARG_INT  | POPT_ARGFLAG_DOC_HIDDEN,   &chop3,            0, "Chop N bp off the 3' end of patterns", "N" },
	{ "chop-5",       'X', POPT_ARG_INT  | POPT_ARGFLAG_DOC_HIDDEN,   &chop5,            0, "Chop N bp off the 5' end of patterns", "N" },
	{ "adenine-hack", 'A', POPT_ARG_NONE | POPT_ARGFLAG_DOC_HIDDEN,   &discount_adenine, 0, "Ignore inserted adenines in db", 0 },
	{ "quiet",        'q', POPT_ARG_NONE | POPT_BIT_CLR, &debug_flags,      debug_notquiet, "Turn off warnings", 0 },
	{ "verbose",      'v', POPT_ARG_NONE | POPT_BIT_SET, &debug_flags,       debug_verbose, "Print additional progress reports", 0 },
	{ "debug",        'd', POPT_ARG_INT  | POPT_ARGFLAG_DOC_HIDDEN,   &debug_flags,      0, "Set debug flags to N (see man page)", "N" },
	POPT_AUTOHELP POPT_TABLEEND
} ;


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

	std::vector<std::string> dbfiles, patfiles ;

	poptContext pc = poptGetContext( "patman", argc, argv, options, 0 ) ;
	poptSetOtherOptionHelp( pc, "[OPTION...] [input-file...]" ) ;
	poptReadDefaultConfig( pc, 0 ) ;
	bool isdbfiles = false ;

	if( argc <= 1 ) { poptPrintHelp( pc, stderr, 0 ) ; return 1 ; }
	int rc = poptGetNextOpt(pc) ;
	for(;;) {
		if( rc == opt_version ) {
			std::cout << poptGetInvocationName(pc) << ", revision " << VERSION << std::endl ;
			return 0 ;
		}
		else if( rc == opt_db || rc == opt_pat || rc == -1 ) {
			while( poptPeekArg(pc) )
				( isdbfiles ? dbfiles : patfiles ).push_back( poptGetArg(pc) ) ;
			isdbfiles = (rc == opt_db) ;
		}
		else if( rc == opt_output )
		{
			delete output ;
			const char *arg = poptGetOptArg(pc) ;
			if( 0 == strcmp( arg, "-" ) ) output = 0 ;
			else output = new std::ofstream( arg ) ;
		}
		else {
			std::cerr << argv[0] << ": " << poptStrerror( rc ) 
				<< ' ' << poptBadOption( pc, 0 ) << std::endl ;
			return 1 ; 
		}
		if( rc == -1 ) break ; else rc = poptGetNextOpt(pc) ;
	}

	if( dbfiles.empty() )
	{ 
		dbfiles.push_back( "-" ) ;
		if( debug_flags & debug_notquiet ) 
			clog << "No databases given, using stdin." << endl ;
	}
	if( patfiles.empty() )
	{
		clog << "No patterns given, aborting." << endl ;
		exit(1) ;
	}

	prefix_tree trie ;
	for( size_t i = 0 ; i != patfiles.size() ; ++i )
	{
		ifstream *infile = 0 ;
		if( patfiles[i] != "-" ) {
			infile = new ifstream( patfiles[i].c_str() ) ;
			if( debug_flags & debug_verbose ) clog << "reading patterns from " << patfiles[i] << endl ;
			if( !infile || !*infile ) { cerr << "cannot read " << patfiles[i] << endl ; return 1 ; }
		}
		else if( debug_flags & debug_verbose ) clog << "reading patterns from STDIN" << endl ;

		// add pattern to prefix tree
		fasta_fac f( infile ? *infile : cin ) ;
		fasta* ff ;
		while ( (ff = f.get()) && !ff->null ) {
			if( minlength <= ff->seq.size() )
				trie.add_seq(
						ff->seq.substr( chop5, ff->seq.size() - chop5 - chop3 ),
						ff->headerline, !!ambi_codes ) ;
			delete ff ;
		}
		delete infile ;
	}

	if( debug_flags & debug_verbose ) clog << "post-processing trie" << endl ;
	trie.add_suffix_links() ;

#ifndef NDEBUG
	trie.debug( cerr ) ;
#endif

	mismatch_ptr *ptrs = 0 ;

	// read database
	string headerline ;
	for( size_t i = 0 ; i != dbfiles.size() ; ++i ) {
		ifstream *dbfile = 0 ;
		if( dbfiles[i] != "-" ) {
			dbfile = new ifstream( dbfiles[i].c_str() ) ;
			if( debug_flags & debug_verbose ) clog << "processing database from " << dbfiles[i] << endl ;
			if( !dbfile || !*dbfile ) { cerr << "cannot read " << dbfiles[i] << endl ; return 1 ; }
		}
		else if( debug_flags & debug_verbose ) clog << "processing database from STDIN" << endl ;

		getline( dbfile ? *dbfile : cin, headerline ) ;
		while( dbfile ? *dbfile : cin ) {
			if ( headerline[0] != '>' ) {
				cerr << "database is not in fasta format" << endl ;
				exit( 2 ) ;
			}
			headerline = headerline.substr( 1, headerline.size()-1 ) ; 
			genome_name = headerline.c_str() ; // global variable

			string seq ;
			bool eat_n = 0 ;
			size_t last_n = 0 ;
			position_genome = 1 ; // global variable

			ptrs = trie.init( ptrs ) ;

			while ( getline( dbfile ? *dbfile : cin, seq )
					&& ( seq.empty() || seq[0] != '>' ) ) {
				for ( size_t i = 0 ; i < seq.size() ; i++,position_genome++ ) 
				{
					if ( toupper( seq[i] ) == 'N' ) 
						if ( eat_n ) continue ;
						else { 
							last_n ++ ;
							if ( last_n == cutoff + 1 ) eat_n = 1 ;
						}
					else if ( eat_n ) { eat_n = 0 ; last_n = 0 ; }

					if( debug_flags & debug_sequence ) std::clog << seq[i] << std::flush ;
					ptrs = trie.compare( seq[i], ptrs ) ;
				}
			}
			trie.compare( 0, ptrs ) ;
			headerline = seq ;
		}

		delete dbfile ;
	}
	( output ? *output : cout ) << std::flush ;
	delete output ;

	if( debug_flags & debug_countnodes )
		clog << "Fetched " << total_nodes << " nodes in total" << endl ;

	poptFreeContext( pc ) ;
	return 0 ;
}
