/**
    bambam
    Copyright (C) 2009-2013 German Tischler
    Copyright (C) 2011-2013 Genome Research Limited

    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 <iostream>
#include <queue>

#include <libmaus2/bambam/BamAlignment.hpp>
#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/bambam/BamWriter.hpp>
#include <libmaus2/bambam/ProgramHeaderLineSet.hpp>
#include <libmaus2/sorting/SortingBufferedOutputFile.hpp>
#include <libmaus2/aio/GenericPeeker.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>

#include <libmaus2/util/ArgInfo.hpp>

#include <biobambam2/Licensing.hpp>

static int getDefaultVerbose() { return 1; }

struct Element
{
	uint32_t p;
	char c;

	Element() {}
	Element(uint32_t const rp, char const rc) : p(rp), c(rc) {}

	void serialise(std::ostream & out) const
	{
		out.write(reinterpret_cast<char const *>(&p),4);
		out.write(reinterpret_cast<char const *>(&c),1);
	}

	void deserialise(std::istream & in)
	{
		in.read(reinterpret_cast<char *>(&p),4);
		in.read(reinterpret_cast<char *>(&c),1);
	}

	bool operator<(Element const & O) const
	{
		if ( p != O.p )
			return p < O.p;
		else
			return c < O.c;
	}
};

struct Key
{
	uint64_t p;
	char c;

	Key()
	{}

	Key(uint64_t const rp, char const rc) : p(rp), c(rc) {}

	bool operator<(Key const & O) const
	{
		if ( p != O.p )
			return p < O.p;
		else
			return c < O.c;
	}
};

void handle(
	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>::unique_ptr_type & pSB,
	int64_t const prevrefid,
	::libmaus2::bambam::BamHeader const & header
)
{
	if ( prevrefid >= 0 )
	{
		std::cout << ">" << header.getRefIDName(prevrefid) << "\n";

		libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>::merger_ptr_type pmerger(pSB->getMerger(1024 /* back block size */,128));
		libmaus2::aio::GenericPeeker<libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>::merger_type> GP(*pmerger);
		Element EL;
		libmaus2::autoarray::AutoArray<char> A;

		while ( GP.peekNext(EL) )
		{
			uint64_t const p = EL.p;

			uint64_t o = 0;
			while ( GP.peekNext(EL) && EL.p == p )
			{
				A.push(o,EL.c);
				GP.getNext(EL);
			}

			std::sort(A.begin(),A.begin()+o);

			std::cerr << prevrefid << "," << p << "\t" << std::string(A.begin(),A.begin()+o) << std::endl;
		}
	}

	pSB.reset();
}

struct EElement
{
	char c;
	uint64_t f;

	EElement() {}
	EElement(char const rc, uint64_t const rf) : c(rc), f(rf) {}

	bool operator<(EElement const & E) const
	{
		if ( f != E.f )
			return f > E.f;
		else
			return c < E.c;
	}
};

void handle(std::map<Key,uint64_t> & M, std::vector<char> & VC, uint64_t const v = std::numeric_limits<uint64_t>::max(), uint64_t const drift = 1000)
{
	while ( M.size() && M.begin()->first.p + drift < v )
	{
		uint64_t const p = M.begin()->first.p;

		std::vector<EElement> V;
		while ( M.size() && M.begin()->first.p == p )
		{
			V.push_back(EElement(M.begin()->first.c,M.begin()->second));
			M.erase(M.begin());
		}

		std::sort(V.begin(),V.end());

		if ( V.size() > 1 )
		{
			std::cerr << "p=" << p << " ";

			for ( uint64_t i = 0; i < V.size(); ++i )
				std::cerr << "(" << V[i].c << "," << V[i].f << ")";

			std::cerr << std::endl;
		}

		VC[p] = V[0].c;
	}
}

int bamrefextract(::libmaus2::util::ArgInfo const & arginfo)
{
	int const verbose = arginfo.getValue<int>("verbose",getDefaultVerbose());

	libmaus2::bambam::BamAlignmentDecoderWrapper::unique_ptr_type decwrapper(
		libmaus2::bambam::BamMultiAlignmentDecoderFactory::construct(arginfo));
	::libmaus2::bambam::BamAlignmentDecoder * ppdec = &(decwrapper->getDecoder());
	::libmaus2::bambam::BamAlignmentDecoder & dec = *ppdec;
	::libmaus2::bambam::BamHeader const & header = dec.getHeader();
	::libmaus2::bambam::BamAlignment const & algn = dec.getAlignment();

	std::string const tmpfilenamebase = arginfo.getValue<std::string>("tmpfile",arginfo.getDefaultTmpFileName());
	std::string const tmpfilenameout = tmpfilenamebase + "_bamsort";
	::libmaus2::util::TempFileRemovalContainer::addTempFile(tmpfilenameout);

	// construct new header
	uint64_t c = 0;

	int64_t prevrefid = -1;
	int64_t prevpos = -1;

	// std::vector< std::pair<uint64_t,char> > V;
	libmaus2::autoarray::AutoArray<libmaus2::bambam::cigar_operation> cigop;

	libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>::unique_ptr_type pSB;

	std::map<Key,uint64_t> M;
	std::vector<char> V;
	int64_t Vid = -1;
	int64_t Vnext = 0;

	while ( dec.readAlignment() )
	{
		// algn.serialise(writer->getStream());

		if ( algn.isMapped() )
		{
			int64_t const currefid = algn.getRefID();

			if ( currefid < prevrefid )
			{
				libmaus2::exception::LibMausException lme;
				lme.getStream() << "[E] file is not in sorted order" << std::endl;
				lme.finish();
				throw lme;
			}

			if ( currefid > prevrefid )
			{
				assert ( currefid >= 0 );
				assert ( currefid < static_cast<int64_t>(header.getNumRef()) );
				// uint64_t const l = header.getRefIDLength(currefid);

				// handle(pSB,prevrefid,header);
				handle(M,V);

				libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>::unique_ptr_type tptr(
					new libmaus2::sorting::SerialisingSortingBufferedOutputFile<Element>(
						tmpfilenameout,128*1024 /* buf size */
					)
				);
				pSB = std::move(tptr);

				prevpos = -1;

				if ( Vid >= 0 )
				{
					while ( Vnext < Vid )
					{
						std::vector<char> V(header.getRefIDLength(Vnext),'N');
						std::cout << ">" << header.getRefIDName(Vnext) << "\n";
						std::cout << std::string(V.begin(),V.end()) << "\n";
						Vnext += 1;
					}

					assert ( Vnext == Vid );

					std::cout << ">" << header.getRefIDName(Vid) << "\n";
					std::cout << std::string(V.begin(),V.end()) << "\n";
					Vnext += 1;
				}

				V.resize(header.getRefIDLength(currefid));
				Vid = currefid;
				std::fill(V.begin(),V.end(),'N');
			}

			int64_t const curpos = algn.getPos();

			if ( curpos < prevpos )
			{
				libmaus2::exception::LibMausException lme;
				lme.getStream() << "[E] file is not in sorted order" << std::endl;
				lme.finish();
				throw lme;

				prevpos = curpos;
			}

			// std::cerr << "(" << currefid << "," << curpos << ")" << " l=" << algn.getLseq() << " ll=" << algn.getLseqByCigar() << std::endl;

			std::string const R = algn.getRead();
			uint32_t const ncig = algn.getCigarOperations(cigop);

			uint64_t refpos = curpos;
			std::string::const_iterator itr = R.begin();

			for ( uint64_t z = 0; z < ncig; ++z )
			{
				libmaus2::bambam::cigar_operation const op = cigop[z];

				for ( int64_t i = 0; i < op.second; ++i )
				{
					//std::cerr << itr-R.begin() << "," << i << "/" << op.second << "," << op.first << std::endl;

					switch ( op.first )
					{
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CMATCH:
						{
							refpos++;
							itr++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CEQUAL:
						{
							#if 0
							if ( *it != *itr )
								return false;
							#endif


							// V.push_back(std::pair<uint64_t,char>(refpos,*itr));
							// pSB->put(Element(refpos,*itr));

							M [ Key(refpos,*itr) ] ++;

							#if 0
							if ( *it == 'N' )
								*it = *itr;
							else
							{
								if ( *it != *itr )
								{
									std::cerr << "[W] " << algn.getName() << " ref=" << *it << " != read=" << *itr << " p=" << (itr-R.begin()) << " " << algn.getCigarString() << std::endl;
								}
								// assert ( *it == *itr );
							}
							#endif

							refpos++;
							itr++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CDIFF:
						{
							#if 0
							if ( *it == *itr )
								return false;
							#endif

							refpos++;
							itr++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CINS:
						{
							itr++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CDEL:
						{
							refpos++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CSOFT_CLIP:
						{
							itr++;
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CHARD_CLIP:
						{
							assert(false);
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CPAD:
						{
							assert(false);
							break;
						}
						case libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_CREF_SKIP:
						{
							assert(false);
							refpos++;
							break;
						}
						default:
						{
							assert(false);
							break;
						}
					}
				}
			}

			handle(M,V,curpos);

			assert ( itr == R.end() );

			prevrefid = currefid;
			prevpos = curpos;
		}

		++c;
		if ( verbose && (c & (1*1024-1)) == 0 )
		{
			if ( prevrefid >= 0 )
	 			std::cerr << "[V] " << c << " p=" << prevpos << "/" << header.getRefIDLength(prevrefid) << std::endl;
			else
	 			std::cerr << "[V] " << c << " p=" << prevpos << std::endl;
		}
	}
	std::cerr << "[V] " << c << std::endl;

	// handle(pSB,prevrefid,header);
	handle(M,V);


	if ( Vid >= 0 )
	{
		while ( Vnext < Vid )
		{
			std::vector<char> V(header.getRefIDLength(Vnext),'N');
			std::cout << ">" << header.getRefIDName(Vnext) << "\n";
			std::cout << std::string(V.begin(),V.end()) << "\n";
			Vnext += 1;
		}

		assert ( Vnext == Vid );

		std::cout << ">" << header.getRefIDName(Vid) << "\n";
		std::cout << std::string(V.begin(),V.end()) << "\n";
		Vnext += 1;
	}

	while ( Vnext < static_cast<int64_t>(header.getNumRef()) )
	{
		std::vector<char> V(header.getRefIDLength(Vnext),'N');
		std::cout << ">" << header.getRefIDName(Vnext) << "\n";
		std::cout << std::string(V.begin(),V.end()) << "\n";
		Vnext += 1;
	}

	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 << ::biobambam2::Licensing::license();
				std::cerr << std::endl;
				std::cerr << "Key=Value pairs:" << std::endl;
				std::cerr << std::endl;

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

				::biobambam2::Licensing::printMap(std::cerr,V);

				std::cerr << std::endl;

				std::cerr << "The keep and remove keys are mutually exclusive. Tags are given by their two character ids. Multiple ids are separated by commas." << std::endl;

				return EXIT_SUCCESS;
			}

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