/**
    bambam
    Copyright (C) 2019 German Tischler

    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 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
**/
#include <config.h>

#include <libmaus2/util/ArgInfo.hpp>
#include <libmaus2/lz/PlainOrGzipStream.hpp>
#include <libmaus2/util/LineBuffer.hpp>
#include <libmaus2/util/TabEntry.hpp>
#include <libmaus2/trie/TrieState.hpp>
#include <libmaus2/geometry/RangeSet.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>
#include <libmaus2/bambam/BamNumericalIndexDecoder.hpp>
#include <libmaus2/bambam/BamNumericalIndexGenerator.hpp>
#include <libmaus2/sorting/SortingBufferedOutputFile.hpp>
#include <libmaus2/aio/GenericPeeker.hpp>

#include <libmaus2/parallel/NumCpus.hpp>

#include <biobambam2/Licensing.hpp>

std::string printHelpMessage(libmaus2::util::ArgInfo const & arginfo)
{
	std::ostringstream ostr;

	ostr << ::biobambam2::Licensing::license();

	ostr << std::endl;

	ostr << "synopsis: " << arginfo.progname << " < in.gtf > out.gtf" << std::endl;

	ostr << std::endl;
	ostr << "Key=Value pairs:" << std::endl;
	ostr << std::endl;

	std::vector< std::pair<std::string,std::string> > V;

	V.push_back ( std::pair<std::string,std::string> ( "T=<filename>", "prefix for temporary files, default: create files in current directory" ) );

	#if 0
	V.push_back ( std::pair<std::string,std::string> ( "verbose=<["+::biobambam2::Licensing::formatNumber(getDefaultVerbose())+"]>", "print progress report (default: 1)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "mapqmin=<["+::biobambam2::Licensing::formatNumber(getDefaultMapQMin())+"]>", "minimum mapping quality processed (default: 255)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "mapqmax=<["+::biobambam2::Licensing::formatNumber(getDefaultMapQMax())+"]>", "maximum mapping quality processed (default: 255)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "threads=<["+::biobambam2::Licensing::formatNumber(1)+"]>", "number of threads used (default: 1, use 0 for number cores on machine)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "uncoveredthres=<["+::formatDouble(getDefaultUncoveredThres())+"]>", "threshold for fraction of uncovered exon bases per transcript (default: 0.1)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "uniqueuncoveredthres=<["+::formatDouble(getDefaultUniqueUncoveredThres())+"]>", "threshold for fraction of uncovered unique exon bases per transcript (default: 0.1)" ) );
	#endif

	::biobambam2::Licensing::printMap(ostr,V);

	return ostr.str();
}

struct Gene
{
	std::string chr;
	std::string gene_id;
	uint64_t start;

	Gene()
	{
	}
	Gene(
		std::string rchr,
		std::string rgene_id,
		uint64_t const rstart
	) : chr(rchr), gene_id(rgene_id), start(rstart)
	{
	}

	bool operator<(Gene const & G) const
	{
		if ( chr != G.chr )
			return chr < G.chr;
		else if ( start != G.start )
			return start < G.start;
		else
			return gene_id < G.gene_id;
	}

	std::ostream & serialise(std::ostream & out) const
	{
		libmaus2::util::StringSerialisation::serialiseString(out,chr);
		libmaus2::util::StringSerialisation::serialiseString(out,gene_id);
		libmaus2::util::NumberSerialisation::serialiseNumber(out,start);
		return out;
	}

	std::istream & deserialise(std::istream & in)
	{
		chr = libmaus2::util::StringSerialisation::deserialiseString(in);
		gene_id = libmaus2::util::StringSerialisation::deserialiseString(in);
		start = libmaus2::util::NumberSerialisation::deserialiseNumber(in);
		return in;
	}
};

struct Entry
{
	std::string chr;
	std::string linetype;
	std::string gene_id;
	std::string transcript_id;
	uint64_t start;
	uint64_t exon_number;
	std::string full;
	uint64_t rank;

	static bool startsWith(std::string const & s, std::string const & prefix)
	{
		return
			s.size() >= prefix.size()
			&&
			s.substr(0,prefix.size()) == prefix;
	}

	static int linetypeToNum(std::string const & s)
	{
		if ( s == "gene" )
			return 0;
		else if ( s == "transcript" )
			return 1;
		else if ( s == "exon" )
			return 2;
		else
		{
			libmaus2::exception::LibMausException lme;
			lme.getStream() << "[E] unknown type " << s << " in Entry::linetypeToNum" << std::endl;
			lme.finish();
			throw lme;
		}
	}

	bool operator<(Entry const & E) const
	{
		if ( rank != E.rank )
			return rank < E.rank;
		if ( chr != E.chr )
			return chr < E.chr;

		if ( gene_id == E.gene_id )
		{
			if ( transcript_id != E.transcript_id )
				return transcript_id < E.transcript_id;

			int const lt_lhs = linetypeToNum(linetype);
			int const lt_rhs = linetypeToNum(E.linetype);

			if ( lt_lhs != lt_rhs )
				return lt_lhs < lt_rhs;

			return exon_number < E.exon_number;
		}
		else
		{
			if ( start != E.start )
				return start < E.start;
			else
				return gene_id < E.gene_id;
		}
	}

	static uint64_t parse(std::string const & s)
	{
		std::istringstream istr(s);
		uint64_t u;
		istr >> u;
		if ( istr && istr.peek() == std::istream::traits_type::eof() )
			return u;
		else
		{
			libmaus2::exception::LibMausException lme;
			lme.getStream() << "[E] invalid number " << s << std::endl;
			lme.finish();
			throw lme;
		}
	}

	Entry()
	{

	}

	Entry(std::string const & line)
	: start(0), exon_number(0), rank(0)
	{
		std::deque<std::string> tokens = libmaus2::util::stringFunctions::tokenize<std::string>(line,std::string("\t"));

		if ( tokens.size() < 9 )
		{
			libmaus2::exception::LibMausException lme;
			lme.getStream() << "[E] invalid line " << line << std::endl;
			lme.finish();
			throw lme;
		}

		chr = tokens[0];
		linetype = tokens[2];
		start = parse(tokens[3]);

		std::deque<std::string> ctokens = libmaus2::util::stringFunctions::tokenize<std::string>(tokens[8],std::string(";"));

		for ( uint64_t i = 0; i < ctokens.size(); ++i )
		{
			std::string & s = ctokens[i];
			uint64_t j = 0;
			while ( j < s.size() && isspace(s[j]) )
				++j;
			s = s.substr(j);

			if ( startsWith(s,"gene_id ") )
			{
				gene_id = s.substr(strlen("gene_id "));
			}
			if ( startsWith(s,"transcript_id ") )
			{
				transcript_id = s.substr(strlen("transcript_id "));
			}
			if ( startsWith(s,"exon_number ") )
			{
				exon_number = parse(s.substr(strlen("exon_number ")));
			}
		}

		full = line;

		// std::cerr << chr << "\t" << linetype << "\t" << gene_id << "\t" << transcript_id << std::endl;
	}

	std::ostream & serialise(std::ostream & out) const
	{
		libmaus2::util::StringSerialisation::serialiseString(out,chr);
		libmaus2::util::StringSerialisation::serialiseString(out,linetype);
		libmaus2::util::StringSerialisation::serialiseString(out,gene_id);
		libmaus2::util::StringSerialisation::serialiseString(out,transcript_id);
		libmaus2::util::StringSerialisation::serialiseString(out,full);
		libmaus2::util::NumberSerialisation::serialiseNumber(out,exon_number);
		libmaus2::util::NumberSerialisation::serialiseNumber(out,start);
		libmaus2::util::NumberSerialisation::serialiseNumber(out,rank);
		return out;
	}

	std::istream & deserialise(std::istream & in)
	{
		chr = libmaus2::util::StringSerialisation::deserialiseString(in);
		linetype = libmaus2::util::StringSerialisation::deserialiseString(in);
		gene_id = libmaus2::util::StringSerialisation::deserialiseString(in);
		transcript_id = libmaus2::util::StringSerialisation::deserialiseString(in);
		full = libmaus2::util::StringSerialisation::deserialiseString(in);
		exon_number = libmaus2::util::NumberSerialisation::deserialiseNumber(in);
		start = libmaus2::util::NumberSerialisation::deserialiseNumber(in);
		rank = libmaus2::util::NumberSerialisation::deserialiseNumber(in);
		return in;
	}
};

static uint64_t idPush(uint64_t o, libmaus2::autoarray::AutoArray<char> & A, std::string const & s)
{
	char const * c = s.c_str();
	A.push(
		o,
		c,
		c+s.size()
	);
	A.push(o,0);

	return o;
}

struct GeneObject
{
	uint64_t offset;
	uint64_t id;

	GeneObject(uint64_t const roffset = 0, uint64_t const rid = 0)
	: offset(roffset), id(rid)
	{
	}
};

struct GeneObjectComparator
{
	char const * c;

	GeneObjectComparator(char const * rc) : c(rc) {}

	bool operator()(
		GeneObject const & A,
		GeneObject const & B
	)
	{
		char const * ca = c + A.offset;
		char const * cb = c + B.offset;

		int const r = strcmp(ca,cb);

		return r < 0;
	}
};

struct GeneObjectNameGet
{
	char const * c;
	GeneObject const * G;

	GeneObjectNameGet()
	{}
	GeneObjectNameGet(
		char const * rc,
		GeneObject const * rG
	) : c(rc), G(rG)
	{}

	std::string get(uint64_t const i) const
	{
		return c + G[i].offset;
	}
};

int filtergtf(libmaus2::util::ArgInfo const & arginfo)
{
	std::string line;
	std::string const tmpfilenamebase = arginfo.getValue<std::string>("T",arginfo.getDefaultTmpFileName());
	std::string const tmpfilenamelist = tmpfilenamebase + "_entries_list";
	std::string const tmpfilenamegenelist = tmpfilenamebase + "_gene_list";
	std::string const tmpfilename = tmpfilenamebase + "_entries";
	libmaus2::util::TempFileRemovalContainer::addTempFile(tmpfilenamelist);
	libmaus2::util::TempFileRemovalContainer::addTempFile(tmpfilenamegenelist);
	libmaus2::util::TempFileRemovalContainer::addTempFile(tmpfilename);

	libmaus2::aio::OutputStreamInstance::unique_ptr_type pOSI(new libmaus2::aio::OutputStreamInstance(tmpfilenamelist));

	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Gene>::unique_ptr_type Pgenesorter(
		new libmaus2::sorting::SerialisingSortingBufferedOutputFile<Gene>(tmpfilenamegenelist,64*1024)
	);

	uint64_t c = 0;
	while ( std::getline(std::cin,line) )
		if ( line.size() )
		{
			if ( line[0] == '#' )
				std::cout << line << "\n";
			else
			{
				Entry const E(line);

				if (
					E.linetype == "gene"
					||
					E.linetype == "transcript"
					||
					E.linetype == "exon"
				)
				{
					if ( E.linetype == "gene" )
					{
						Pgenesorter->put(Gene(E.chr,E.gene_id,E.start));
					}
					E.serialise(*pOSI);

					if ( ++c % (64*1024) == 0 )
						std::cerr << "[V] "<< c << std::endl;
				}
			}
		}

	pOSI->flush();
	pOSI.reset();

	libmaus2::autoarray::AutoArray<char> Aid;
	libmaus2::autoarray::AutoArray<GeneObject> GO;

	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Gene>::merger_ptr_type Pgenemerger(Pgenesorter->getMerger());
	Gene G;
	uint64_t Aid_o = 0;
	uint64_t GO_o = 0;
	for ( uint64_t id = 0; Pgenemerger->getNext(G); ++id )
	{
		uint64_t const name_offset = Aid_o;
		Aid_o = idPush(Aid_o,Aid,G.gene_id);
		GO.push(GO_o,GeneObject(name_offset,id));
	}

	std::sort(GO.begin(),GO.begin()+GO_o,GeneObjectComparator(Aid.begin()));

	for ( uint64_t i = 1; i < GO_o; ++i )
	{
		char const * c_a = Aid.begin() + GO[i-1].offset;
		char const * c_b = Aid.begin() + GO[i-0].offset;
		assert ( strcmp(c_a,c_b) < 0 );
	}

	typedef libmaus2::util::ConstIterator<GeneObjectNameGet,std::string> name_iterator;
	GeneObjectNameGet const GONG(Aid.begin(),GO.begin());
	name_iterator it_a(&GONG,0);
	name_iterator it_e(&GONG,GO_o);

	libmaus2::aio::InputStreamInstance::unique_ptr_type ISI(new libmaus2::aio::InputStreamInstance(tmpfilenamelist));

	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Entry>::unique_ptr_type Psorter(
		new libmaus2::sorting::SerialisingSortingBufferedOutputFile<Entry>(tmpfilename,64*1024)
	);

	Entry E;
	while ( *ISI && ISI->peek() != std::istream::traits_type::eof() )
	{
		E.deserialise(*ISI);

		name_iterator it_c = ::std::lower_bound(it_a,it_e,E.gene_id);

		assert ( it_c != it_e && *it_c == E.gene_id );

		uint64_t const gene_rank = GO[it_c - it_a].id;

		// std::cerr << "name " << E.gene_id << " rank " << gene_rank << std::endl;

		E.rank = gene_rank;

		Psorter->put(E);
	}

	ISI.reset();
	libmaus2::aio::FileRemoval::removeFile(tmpfilenamelist);

	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Entry>::merger_ptr_type Pmerger(Psorter->getMerger());

	while ( Pmerger->getNext(E) )
	{
		std::cout << E.full << "\n";
	}

	Pmerger.reset();
	Psorter.reset();

	libmaus2::aio::FileRemoval::removeFile(tmpfilename);

	if ( ! std::cout )
	{
		libmaus2::exception::LibMausException lme;
		lme.getStream() << "[E] failed to write output" << std::endl;
		lme.finish();
		throw lme;
	}

	return EXIT_SUCCESS;
}

int main(int argc, char * argv[])
{
	try
	{
		::libmaus2::util::ArgInfo const arginfo(argc,argv);

		for ( uint64_t i = 0; i < arginfo.restargs.size(); ++i )
			if (
				arginfo.restargs[i] == "-v"
				||
				arginfo.restargs[i] == "--version"
			)
			{
				std::cerr << ::biobambam2::Licensing::license();
				return EXIT_SUCCESS;
			}
			else if (
				arginfo.restargs[i] == "-h"
				||
				arginfo.restargs[i] == "--help"
			)
			{
				std::cerr << printHelpMessage(arginfo);
				return EXIT_SUCCESS;
			}

		return filtergtf(arginfo);
	}
	catch(std::exception const & ex)
	{
		std::cerr << ex.what() << std::endl;
		return EXIT_FAILURE;
	}
}
