/**
    biobambam
    Copyright (C) 2020 German Tischler-Höhle

    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 <biobambam2/BamBamConfig.hpp>
#include <biobambam2/Licensing.hpp>
#include <biobambam2/DepthInterval.hpp>

#include <iomanip>

#include <config.h>

#include <libmaus2/util/TempFileRemovalContainer.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>
#include <libmaus2/util/ArgParser.hpp>
#include <libmaus2/util/ArgInfo.hpp>
#include <libmaus2/bambam/CircularHashCollatingBamDecoder.hpp>

struct Callback : public libmaus2::bambam::CollatingBamDecoderAlignmentInputCallback
{
	uint64_t sec;
	uint64_t sup;

	Callback() : sec(0), sup(0)
	{
	}

	virtual void operator()(::libmaus2::bambam::BamAlignment const & A)
	{
		if ( A.isSecondary() )
		{
			sec++;
		}
		if ( A.isSupplementary() )
		{
			sup++;
		}
	}
};

struct Counter
{
	uint64_t pairs;
	uint64_t pairs_both;
	uint64_t pairs_none;
	uint64_t pairs_1;
	uint64_t pairs_2;
	uint64_t pairs_proper;

	uint64_t o1;
	uint64_t o1_mapped;
	uint64_t o1_unmapped;

	uint64_t o2;
	uint64_t o2_mapped;
	uint64_t o2_unmapped;

	Counter()
	: pairs(0), pairs_both(0), pairs_none(0), pairs_1(0), pairs_2(0), pairs_proper(0),
	  o1(0), o1_mapped(0), o1_unmapped(0), o2(0), o2_mapped(0), o2_unmapped(0)
	{

	}

	void update(std::pair <libmaus2::bambam::BamAlignment const *, libmaus2::bambam::BamAlignment const *> const & P)
	{
		if ( P.first && P.second )
		{
			pairs += 1;

			if ( P.first->isMapped() && P.second->isMapped() )
			{
				pairs_both += 1;

				if ( P.first->isProper() && P.second->isProper() )
				{
					pairs_proper += 1;
				}
			}
			else if ( P.first->isMapped() )
			{
				pairs_1 += 1;
			}
			else if ( P.second->isMapped() )
			{
				pairs_2 += 1;
			}
			else
			{
				pairs_none += 1;
			}
		}
		else if ( P.first )
		{
			o1 += 1;

			if ( P.first->isMapped() )
				o1_mapped++;
			else
				o1_unmapped++;
		}
		else if ( P.second )
		{
			o2 += 1;

			if ( P.second->isMapped() )
				o2_mapped++;
			else
				o2_unmapped++;
		}
	}
};

std::ostream & operator<<(std::ostream & out, Counter const & C)
{
	out << C.pairs << "/" << C.o1 << "/" << C.o2 << "\t"
		<< C.pairs_both << "/" << C.pairs_proper << "/" << C.pairs_1 << "/" << C.pairs_2 << "/" << C.pairs_none << "\t"
		<< C.o1_mapped << "/" << C.o1_unmapped << "\t"
		<< C.o2_mapped << "/" << C.o2_unmapped
		;
	return out;
}

struct CounterContainer
{
	Counter global;
	Counter split;
	std::map<uint64_t,Counter> M;

	void update(std::pair <libmaus2::bambam::BamAlignment const *, libmaus2::bambam::BamAlignment const *> const & P)
	{
		global.update(P);

		if ( P.first && P.second )
		{
			if ( P.first->isMapped() && P.second->isMapped() )
			{
				if ( P.first->getRefID() == P.second->getRefID() )
					M[P.first->getRefID()].update(P);
				else
					split.update(P);
			}
			else if ( P.first->isMapped() )
				M[P.first->getRefID()].update(P);
			else if ( P.second->isMapped() )
				M[P.second->getRefID()].update(P);
		}
		else if ( P.first )
		{
			if ( P.first->isMapped() )
				M[P.first->getRefID()].update(P);
		}
		else if ( P.second )
		{
			if ( P.second->isMapped() )
				M[P.second->getRefID()].update(P);
		}
	}
};

void bamcountflags(libmaus2::util::ArgParser const & arg)
{
	libmaus2::bambam::BamAlignmentDecoderWrapper::unique_ptr_type decwrapper(
		libmaus2::bambam::BamMultiAlignmentDecoderFactory::construct(arg,true /* put rank */));
	::libmaus2::bambam::BamAlignmentDecoder * ppdec = &(decwrapper->getDecoder());
	::libmaus2::bambam::BamAlignmentDecoder & dec = *ppdec;
	::libmaus2::bambam::BamHeader const & header = dec.getHeader();
	std::string const tmpfilebase = arg.uniqueArgPresent("T") ? arg["T"] : libmaus2::util::ArgInfo::getDefaultTmpFileName(arg.progname);

	libmaus2::bambam::CircularHashCollatingBamDecoder CHC(
		dec,tmpfilebase+".coltmp",
		libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_FSECONDARY|libmaus2::bambam::BamFlagBase::LIBMAUS2_BAMBAM_FSUPPLEMENTARY);

	Callback CB;

	CHC.addInputCallback(&CB);

	std::pair <libmaus2::bambam::BamAlignment const *, libmaus2::bambam::BamAlignment const *> P;

	uint64_t c = 0;
	uint64_t p = 0;
	uint64_t const mod = 1024*1024;

	CounterContainer C;

	uint64_t mapped = 0;
	uint64_t unmapped = 0;

	while ( CHC.tryPair(P) )
	{
		C.update(P);

		if ( P.first )
		{
			if ( P.first->isMapped() )
				mapped += 1;
			else
				unmapped += 1;

			++c;
		}
		if ( P.second )
		{
			if ( P.second->isMapped() )
				mapped += 1;
			else
				unmapped += 1;

			++c;
		}

		if ( c / mod != p / mod )
		{
			std::cerr << "[V]\t" << ((c/mod)*mod) << "\t" << C.global << "\t" << mapped << "/" << unmapped << std::endl;
			p = c;
		}
	}

	std::cout << "global\t" << C.global << "\t" << mapped << "/" << unmapped << std::endl;
	for ( std::map<uint64_t,Counter>::const_iterator it = C.M.begin(); it != C.M.end(); ++it )
	{
		std::cout << header.getRefIDName(it->first) << "\t" << it->second << std::endl;
	}

	#if 0
	libmaus2::bambam::BamAlignment const & algn = dec.getAlignment();
	uint64_t const mindepth = arg.getParsedArgOrDefault<uint64_t>("mindepth",0);
	uint64_t const maxdepth = arg.getParsedArgOrDefault<uint64_t>("maxdepth",std::numeric_limits<int64_t>::max());
	bool const binary = arg.getParsedArgOrDefault<int>("binary",0);
	int const verbose = arg.getParsedArgOrDefault<int>("verbose",1);

	libmaus2::util::FiniteSizeHeap<uint64_t> FSH(0);

	uint64_t prevstart = 0;
	uint64_t depth = 0;
	int64_t prevrefid = -1;

	IntervalHandler IH(mindepth,maxdepth,binary,header);

	while ( dec.readAlignment() )
	{
		if ( algn.isMapped() )
		{
			if ( verbose && (algn.getRefID() != prevrefid) )
				std::cerr << "[V] " << dec.getHeader().getRefIDName(algn.getRefID()) << std::endl;

			libmaus2::math::IntegerInterval<int64_t> const I = algn.getReferenceInterval();

			int64_t const i_from = I.from;
			int64_t const i_to = i_from + I.diameter();

			// std::cerr << "[" << i_from << "," << i_to << ")" << std::endl;

			// stack is not empty and top element is finished
			while (
				(!FSH.empty())
				&&
				(
					algn.getRefID() != prevrefid
					||
					static_cast<int64_t>(FSH.top()) <= i_from
				)
			)
			{
				uint64_t const end = FSH.pop();

				if ( end > prevstart )
				{
					biobambam2::DepthInterval const D(prevrefid,prevstart,end,depth);
					IH.handle(D);
				}

				depth -= 1;
				prevstart = end;
			}

			// we have reached the end of a refid
			if ( algn.getRefID() != prevrefid && prevrefid >= 0 )
			{
				assert ( depth == 0 );

				// length of ref id
				int64_t const len = dec.getHeader().getRefIDLength(prevrefid);

				// if there is a depth 0 stretch in the end
				if ( len > static_cast<int64_t>(prevstart) )
				{
					biobambam2::DepthInterval const D(prevrefid,prevstart,len,depth);
					IH.handle(D);
				}

				prevstart = 0;
			}

			if ( i_from > static_cast<int64_t>(prevstart) )
			{
				biobambam2::DepthInterval const D(algn.getRefID(),prevstart,i_from,depth);
				IH.handle(D);
			}

			depth += 1;
			prevstart = i_from;

			FSH.pushBump(i_to);

			prevrefid = algn.getRefID();
		}
	}

	while (
		(!FSH.empty())
	)
	{
		uint64_t const end = FSH.pop();

		if ( end > prevstart )
		{
			biobambam2::DepthInterval const D(prevrefid,prevstart,end,depth);
			IH.handle(D);
		}

		depth -= 1;
		prevstart = end;
	}

	if ( prevrefid >= 0 )
	{
		assert ( depth == 0 );

		int64_t const len = dec.getHeader().getRefIDLength(prevrefid);

		if ( len > static_cast<int64_t>(prevstart) )
		{
			biobambam2::DepthInterval const D(prevrefid,prevstart,len,depth);
			IH.handle(D);
		}

		prevstart = 0;
	}

	IH.finish();
	#endif
}

int main(int argc, char * argv[])
{
	try
	{
		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> Vformat = libmaus2::bambam::BamAlignmentDecoderInfo::getArgumentDefinitions();
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("v","verbose",true));
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("T","",true));

		::libmaus2::util::ArgParser arg(argc,argv,Vformat);

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