/**
    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 <libmaus2/lz/BgzfInflateBase.hpp>
#include <libmaus2/bambam/BamHeader.hpp>
#include <libmaus2/bambam/SamDecoder.hpp>
#include <libmaus2/util/ArgInfo.hpp>
#include <libmaus2/util/LineBuffer.hpp>
#include <libmaus2/util/Concat.hpp>
#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/bambam/BamNumericalIndexDecoder.hpp>
#include <libmaus2/bambam/BamNumericalIndexGenerator.hpp>
#include <libmaus2/fastx/StreamFastAReader.hpp>
#include <libmaus2/lz/BgzfDeflate.hpp>
#include <libmaus2/util/GetFileSize.hpp>
#include <libmaus2/util/PrefixSums.hpp>
#include <libmaus2/aio/TempFileArray.hpp>
#include <libmaus2/parallel/NumCpus.hpp>
#include <libmaus2/clustering/KmerBase.hpp>
#include <libmaus2/fastx/SingleWordDNABitBuffer.hpp>
#include <libmaus2/hashing/hash.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>

#include <biobambam2/BamBamConfig.hpp>
#include <biobambam2/Licensing.hpp>
#include <biobambam2/ReadHeader.hpp>
#include <biobambam2/RunEOFFilter.hpp>

int getDefaultLevel()
{
	return -1;
}

void printHelpMessage(libmaus2::util::ArgInfo const & arginfo)
{
	std::cerr << ::biobambam2::Licensing::license();
	std::cerr << std::endl;

	std::cerr << "synopsis: " << arginfo.progname << " [thresval=3] [k=24] [threads=1] [level=6] <in.bam> <in.fasta> >out.bam" << std::endl;
	std::cerr << std::endl;
	std::cerr << "Key=Value pairs:" << std::endl;
	std::cerr << std::endl;

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

	V.push_back ( std::pair<std::string,std::string> ( "level=<["+::biobambam2::Licensing::formatNumber(getDefaultLevel())+"]>", libmaus2::bambam::BamBlockWriterBaseFactory::getBamOutputLevelHelpText() ) );
	V.push_back ( std::pair<std::string,std::string> ( "threads=<["+::biobambam2::Licensing::formatNumber(1)+"]>", "number of threads used (default: 1)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "k=<["+::biobambam2::Licensing::formatNumber(24)+"]>", "k-mer size used (default: 24)" ) );
	V.push_back ( std::pair<std::string,std::string> ( "thresval=<["+::biobambam2::Licensing::formatNumber(3)+"]>", "threshold for number of k-mers required to keep read (default: 3)" ) );

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

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

std::vector<std::string> loadRef(std::string const & fn)
{
	libmaus2::aio::InputStreamInstance ISI(fn);
	libmaus2::fastx::StreamFastAReaderWrapper SFAR(ISI);
	libmaus2::fastx::StreamFastAReader::pattern_type pattern;
	std::vector<std::string> V;

	while ( SFAR.getNextPatternUnlocked(pattern) )
	{
		std::string & s = pattern.spattern;

		for ( uint64_t i = 0; i < s.size(); ++i )
			s[i] = libmaus2::fastx::remapChar(libmaus2::fastx::mapChar(s[i]));

		V.push_back(s);
	}

	return V;
}

template<unsigned int kl>
struct KMerBlockTemplate
{
	uint64_t kmer;
	uint64_t hash;
	uint64_t from;
	uint64_t to;

	static uint64_t getHashSize()
	{
		return 1ull << kl;
	}

	static uint64_t getHash(uint64_t const kmer)
	{
		return libmaus2::hashing::EvaHash::hash642(&kmer,1) & (getHashSize()-1);
	}

	KMerBlockTemplate() : kmer(0), hash(0), from(0), to(0) {}
	KMerBlockTemplate(uint64_t const rkmer)
	: kmer(rkmer), hash(getHash(kmer)), from(0), to(0) {}
	KMerBlockTemplate(uint64_t const rkmer, uint64_t const rfrom, uint64_t const rto)
	: kmer(rkmer), hash(getHash(kmer)), from(rfrom), to(rto) {}

	bool operator<(KMerBlockTemplate<kl> const & K) const
	{
		if ( hash != K.hash )
			return hash < K.hash;
		else
			return kmer < K.kmer;
	}
};

struct HashBlock
{
	uint64_t from;
	uint64_t to;

	HashBlock() : from(0), to(0) {}
	HashBlock(uint64_t const rfrom, uint64_t const rto)
	: from(rfrom), to(rto) {}
};

static unsigned int const kl = 16;

typedef KMerBlockTemplate<kl> KMerBlock;

static libmaus2::bambam::BamHeader::unique_ptr_type readBamHeader(libmaus2::util::ArgInfo const & arginfo, std::string const & bamfn)
{
	libmaus2::util::ArgInfo arginfoCopy = arginfo;
	arginfoCopy.replaceKey("I",bamfn);

	libmaus2::bambam::BamAlignmentDecoderWrapper::unique_ptr_type bamdecwrapper(
		libmaus2::bambam::BamMultiAlignmentDecoderFactory::construct(arginfoCopy)
	);
	libmaus2::bambam::BamHeader::unique_ptr_type ubamheader(bamdecwrapper->getDecoder().getHeader().uclone());

	return ubamheader;
}

struct Match
{
	uint64_t seqid;
	uint64_t strand;
	uint64_t pos;
	uint64_t readpos;

	Match()
	{}

	Match(uint64_t const rseqid, uint64_t const rstrand, uint64_t const rpos, uint64_t const rreadpos)
	: seqid(rseqid), strand(rstrand), pos(rpos), readpos(rreadpos) {}

	bool operator<(Match const & M) const
	{
		if ( seqid != M.seqid )
			return seqid < M.seqid;
		else if ( strand != M.strand )
			return strand < M.strand;
		else
			return pos < M.pos;
	}
};

struct KMer
{
	uint64_t kmer;
	uint64_t seqid;
	uint64_t pos;
	uint64_t strand;

	KMer()
	{

	}

	KMer(
		uint64_t rkmer,
		uint64_t rseqid,
		uint64_t rpos,
		uint64_t rstrand
	) : kmer(rkmer), seqid(rseqid), pos(rpos), strand(rstrand)
	{

	}

	std::string toString(uint64_t const k) const
	{
		libmaus2::fastx::SingleWordDNABitBuffer forw(k);
		forw.buffer = kmer;

		std::ostringstream ostr;
		ostr << "KMer(" << forw << "," << seqid << "," << pos << "," << strand << ")";
		return ostr.str();
	}
};


struct KMerKSorter
{
	bool operator()(KMer const & A, KMer const & B) const
	{
		return A.kmer < B.kmer;
	}
};

struct MatchCallback
{
	uint64_t AMo;
	libmaus2::autoarray::AutoArray<HashBlock> const & KH;
	libmaus2::autoarray::AutoArray< KMerBlock > const & AKB;
	libmaus2::autoarray::AutoArray<KMer> const & AK;
	libmaus2::autoarray::AutoArray<Match> & AM;

	MatchCallback(
		libmaus2::autoarray::AutoArray<HashBlock> const & rKH,
		libmaus2::autoarray::AutoArray< KMerBlock > const & rAKB,
		libmaus2::autoarray::AutoArray<KMer> const & rAK,
		libmaus2::autoarray::AutoArray<Match> & rAM
	) : AMo(0), KH(rKH), AKB(rAKB), AK(rAK), AM(rAM)
	{

	}

	void operator()(uint64_t const kmer, uint64_t const rpos)
	{
		uint64_t const hash = KMerBlock::getHash(kmer);
		assert ( hash < KH.size() );
		HashBlock const & HB = KH[hash];
		if ( HB.from < HB.to )
		{
			for ( uint64_t h = HB.from; h < HB.to; ++h )
			{
				KMerBlock const & KB = AKB[h];
				assert ( KB.hash == hash );

				if ( KB.kmer == kmer )
				{
					uint64_t const kfrom = KB.from;
					uint64_t const kto = KB.to;

					for ( uint64_t z = kfrom; z < kto; ++z )
					{
						assert ( AK[z].kmer == kmer );

						AM.push(AMo,Match(AK[z].seqid,AK[z].strand,AK[z].pos,rpos));
					}
				}
			}
		}
	}

	bool thres(uint64_t const rthres)
	{
		std::sort(AM.begin(),AM.begin()+AMo);

		uint64_t ilow = 0;
		bool found = false;

		while ( ilow < AMo )
		{
			uint64_t ihigh = ilow+1;

			// look for same seqid and strand
			while (
				ihigh < AMo
				&&
				AM[ihigh].seqid  == AM[ilow].seqid
				&&
				AM[ihigh].strand == AM[ilow].strand
			)
				++ihigh;

			// std::cerr << AM[ilow].seqid << "," << AM[ilow].strand << "=" << ihigh-ilow << std::endl;

			uint64_t c = 1;
			for ( uint64_t z = ilow+1; z < ihigh; ++z )
				if ( AM[z-1].pos != AM[z].pos )
					++c;

			if ( c >= rthres )
				found = true;

			ilow = ihigh;
		}

		return found;
	}
};


struct Filter
{
	libmaus2::autoarray::AutoArray<HashBlock> const & KH;
	libmaus2::autoarray::AutoArray< KMerBlock > const & AKB;
	libmaus2::autoarray::AutoArray<KMer> const & AK;
	libmaus2::clustering::KmerBase const & KB;
	unsigned int const k;
	uint64_t const thresval;

	Filter(
		libmaus2::autoarray::AutoArray<HashBlock> const & rKH,
		libmaus2::autoarray::AutoArray< KMerBlock > const & rAKB,
		libmaus2::autoarray::AutoArray<KMer> const & rAK,
		libmaus2::clustering::KmerBase const & rKB,
		unsigned int const rk,
		uint64_t const rthresval
	) : KH(rKH), AKB(rAKB), AK(rAK), KB(rKB), k(rk), thresval(rthresval)
	{

	}

	bool checkFilter(
		std::pair<
			std::pair<uint8_t const *,uint64_t>,
			libmaus2::bambam::BamRawDecoderBase::RawInterval
		> const & P,
		libmaus2::autoarray::AutoArray<char> & A,
		libmaus2::fastx::SingleWordDNABitBuffer & forw,
		libmaus2::autoarray::AutoArray<Match> & AM
	) const
	{
		uint64_t const lseq = libmaus2::bambam::BamAlignmentDecoderBase::decodeRead(P.first.first,A);

		#if 0
		for ( uint64_t z = 0; z < lseq; ++z )
			A[z] = libmaus2::fastx::mapChar(A[z]);
		#endif

		// std::cerr << "j=" << j << " lseq=" << lseq << " [" << low << "," << high << ")" << std::endl;


		MatchCallback CB(KH,AKB,AK,AM);

		KB.kmerCallbackPosForwardOnly(
			A.begin(),
			lseq,
			CB,
			forw,
			k
		);

		bool const thres = CB.thres(thresval);

		return thres;
	}
};

static void putAlignment(
	std::pair<
		std::pair<uint8_t const *,uint64_t>,
		libmaus2::bambam::BamRawDecoderBase::RawInterval
	> const & P,
	libmaus2::lz::BgzfDeflate<std::ostream> & BOS
)
{
	// write block size
	::libmaus2::bambam::EncoderBase::putLE<libmaus2::lz::BgzfDeflate<std::ostream>,uint32_t>(BOS,P.first.second);

	BOS.write(
		reinterpret_cast<char const *>(P.first.first),
		P.first.second
	);
}

int bamfilterk(libmaus2::util::ArgInfo const & arginfo)
{
	uint64_t const numrest = arginfo.getNumRestArgs();
	int const level = libmaus2::bambam::BamBlockWriterBaseFactory::checkCompressionLevel(arginfo.getValue("level",getDefaultLevel()));
	unsigned int numthreads = arginfo.getValue<unsigned int>("threads",1);
	unsigned int const k = arginfo.getValue<unsigned int>("k",24);
	uint64_t const thresval = arginfo.getValue<unsigned int>("thresval",3);
	assert ( k );
	if ( ! numthreads )
		numthreads = libmaus2::parallel::NumCpus::getNumLogicalProcessors();

	std::string const tmpfilenamebase = arginfo.getValue<std::string>("tmpfile",arginfo.getDefaultTmpFileName());
	std::string const tmpfilenameout = tmpfilenamebase + "_bamfilterk";

	if ( numrest < 2 )
	{
		printHelpMessage(arginfo);

		libmaus2::exception::LibMausException lme;
		lme.getStream() << "[E] insufficient number of arguments provided" << std::endl;
		lme.finish();
		throw lme;
	}
	std::string const bamfn = arginfo.getUnparsedRestArg(0);
	std::string const kfn = arginfo.getUnparsedRestArg(1);

	libmaus2::bambam::BamHeader::unique_ptr_type ubamheader(readBamHeader(arginfo,bamfn));

	std::string bgzfhead;
	{
		std::ostringstream headerstr;
		ubamheader->serialise(headerstr);
		std::string const headerser = headerstr.str();

		std::ostringstream defstr;
		libmaus2::lz::BgzfDeflate<std::ostream> def(defstr,level);
		def.write(headerser.c_str(),headerser.size());
		def.flush();
		bgzfhead = biobambam2::runEOFFilter(defstr.str());
	}

	std::string const indexfn = libmaus2::bambam::BamNumericalIndexBase::getIndexName(bamfn);

	libmaus2::bambam::BamNumericalIndexGenerator::indexFileCheck(bamfn,1024,numthreads);

	libmaus2::bambam::BamNumericalIndexDecoder indexdec(indexfn);
	uint64_t const n = indexdec.getAlignmentCount();
	uint64_t const packsperthread = 8;
	uint64_t const tnumpacks = packsperthread*numthreads;
	assert ( tnumpacks );
	uint64_t packsize = (n + tnumpacks - 1)/tnumpacks;
	if ( packsize % 2 )
		packsize += 1;
	assert ( packsize % 2 == 0 );
	uint64_t const numpacks = packsize ? (n + packsize - 1)/packsize : 0;

	std::vector<std::string> const Vref = loadRef(kfn);
	libmaus2::clustering::KmerBase const KB;

	libmaus2::autoarray::AutoArray<KMer> AK;
	uint64_t o_AK = 0;

	for ( uint64_t i = 0; i < Vref.size(); ++i )
	{
		libmaus2::fastx::SingleWordDNABitBuffer forw(k);
		libmaus2::fastx::SingleWordDNABitBuffer reve(k);

		struct Callback
		{
			uint64_t const seqid;
			libmaus2::autoarray::AutoArray<KMer> & AK;
			uint64_t o_AK;

			Callback(uint64_t const rseqid, libmaus2::autoarray::AutoArray<KMer> & rAK, uint64_t ro_AK) : seqid(rseqid), AK(rAK), o_AK(ro_AK) {}

			void operator()(uint64_t const kmer, uint64_t const pos, bool const z)
			{
				//std::cerr << "kmer " << kmer << " seqid " << seqid << " pos " << pos << " z " << z << std::endl;
				AK.push(o_AK,KMer(kmer,seqid,pos,z));
			}
		};

		Callback CB(i,AK,o_AK);

		KB.kmerCallbackPosFR(
			Vref[i].c_str(),
			Vref[i].size(),
			CB,
			forw,
			reve,
			k
		);

		o_AK = CB.o_AK;
	}

	std::sort(AK.begin(),AK.begin()+o_AK,KMerKSorter());

	libmaus2::autoarray::AutoArray< KMerBlock > AKB;
	uint64_t o_AKB = 0;

	uint64_t ilow = 0;
	while ( ilow < o_AK )
	{
		uint64_t const kmer = AK[ilow].kmer;
		uint64_t ihigh = ilow+1;

		while ( ihigh < o_AK && AK[ihigh].kmer == kmer )
			++ihigh;

		AKB.push(o_AKB,KMerBlock(kmer,ilow,ihigh));

		ilow = ihigh;
	}

	std::sort(AKB.begin(),AKB.begin()+o_AKB);

	libmaus2::autoarray::AutoArray<HashBlock> KH(KMerBlock::getHashSize());

	ilow = 0;
	while ( ilow < o_AKB )
	{
		uint64_t const hash = AKB[ilow].hash;
		uint64_t ihigh = ilow+1;

		while ( ihigh < o_AKB && AKB[ihigh].hash == hash )
			++ihigh;

		assert ( hash < KH.size() );

		KH[hash] = HashBlock(ilow,ihigh);

		// std::cerr << "H\t" << hash << "\t" << ihigh-ilow << std::endl;

		ilow = ihigh;
	}

	std::vector<std::string> Vtmp(numpacks);
	#if defined(_OPENMP)
	#pragma omp parallel for num_threads(numthreads)
	#endif
	for ( uint64_t i = 0; i < numpacks; ++i )
	{
		std::ostringstream fnostr;
		fnostr << tmpfilenameout << "_" << std::setw(6) << std::setfill('0') << i << std::setw(0);
		Vtmp[i] = fnostr.str();
	}
	for ( uint64_t i = 0; i < numpacks; ++i )
		libmaus2::util::TempFileRemovalContainer::addTempFile(Vtmp[i]);
	std::atomic<uint64_t> packfinished(0);
	std::atomic<uint64_t> written(0);

	Filter const filter(KH,AKB,AK,KB,k,thresval);

	std::atomic<uint64_t> parfailed(0);

	#if defined(_OPENMP)
	#pragma omp parallel for num_threads(numthreads) schedule(dynamic,1)
	#endif
	for ( uint64_t i = 0; i < numpacks; ++i )
	{
		try
		{
			std::string const & tmpfn = Vtmp[i];
			uint64_t const low = i * packsize;
			uint64_t const high = std::min(low+packsize,n);
			assert ( high > low );

			assert ( low % 2 == 0 );
			assert ( high % 2 == 0 );

			uint64_t const nump = (high-low)/2;

			// uint64_t const size = high-low;

			libmaus2::aio::OutputStreamInstance::unique_ptr_type pOS(new libmaus2::aio::OutputStreamInstance(tmpfn));
			libmaus2::lz::BgzfDeflate<std::ostream>::unique_ptr_type pBOS(new libmaus2::lz::BgzfDeflate<std::ostream>(*pOS,level));

			std::pair<
				std::pair<uint8_t const *,uint64_t>,
				libmaus2::bambam::BamRawDecoderBase::RawInterval
				> P[2];

			libmaus2::bambam::BamRawDecoder::unique_ptr_type pdec(indexdec.getRawDecoderAt(bamfn,low));

			libmaus2::autoarray::AutoArray<uint8_t> U0;
			libmaus2::autoarray::AutoArray<uint8_t> U1;

			libmaus2::autoarray::AutoArray<char> A;
			libmaus2::fastx::SingleWordDNABitBuffer forw(k);
			libmaus2::autoarray::AutoArray<Match> AM;

			for ( uint64_t pindex = 0; pindex < nump; ++pindex )
			{
				P[0] = pdec->getPos();
				assert ( P[0].first.first );
				U0.ensureSize(P[0].first.second);
				std::copy(P[0].first.first,P[0].first.first + P[0].first.second,U0.begin());
				P[0].first.first = U0.begin();

				P[1] = pdec->getPos();
				assert ( P[1].first.first );

				#if 0
				U1.ensureSize(P[1].first.second);
				std::copy(P[1].first.first,P[1].first.first + P[1].first.second,U1.begin());
				P[1].first.first = U1.begin();
				#endif

				char const * c0 = libmaus2::bambam::BamAlignmentDecoderBase::getReadName(P[0].first.first);
				char const * c1 = libmaus2::bambam::BamAlignmentDecoderBase::getReadName(P[1].first.first);

				bool const samename = (strcmp(c0,c1) == 0);

				if ( ! samename )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] samename check failed, mismatching are" << std::endl;
					lme.getStream() << "[E] " << c0 << std::endl;
					lme.getStream() << "[E] " << c1 << std::endl;
					lme.finish();
					throw lme;
				}
				assert ( samename );

				bool const thres0 = filter.checkFilter(P[0],A,forw,AM);
				bool const thres1 = filter.checkFilter(P[1],A,forw,AM);

				if ( thres0 || thres1 )
				{
					putAlignment(P[0],*pBOS);
					putAlignment(P[1],*pBOS);
					written += 1;
				}
			}

			pBOS->flush();
			pBOS.reset();
			pOS->flush();
			pOS.reset();

			uint64_t const lpackfinished = ++packfinished;

			{
				libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
				std::cerr << lpackfinished << "/" << numpacks << "\t" << written << std::endl;
			}
		}
		catch(std::exception const & ex)
		{
			parfailed++;
			{
				libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
				std::cerr << ex.what() << std::endl;
			}
		}
	}

	if ( parfailed )
	{
		libmaus2::exception::LibMausException lme;
		lme.getStream() << "[E] parallel loop failed" << std::endl;
		lme.finish();
		throw lme;
	}

	std::cerr << "[V] wrote " << written << " pairs" << std::endl;

	std::cout.write(bgzfhead.c_str(),bgzfhead.size());
	for ( uint64_t i = 0; i < numpacks; ++i )
	{
		std::string const & tmpfn = Vtmp[i];

		// std::cerr << tmpfn << "\t" << libmaus2::util::GetFileSize::getFileSize(tmpfn) << std::endl;

		{
		libmaus2::aio::InputStreamInstance ISI(tmpfn);
		libmaus2::util::Concat::concat(ISI,std::cout);
		}
		libmaus2::aio::FileRemoval::removeFile(tmpfn);
	}

	std::string const eofblock = libmaus2::lz::BgzfDeflate<std::ostream>::getEOFBlock();
	std::cout.write(eofblock.c_str(),eofblock.size());
	std::cout.flush();

	if ( ! std::cout )
	{
		libmaus2::exception::LibMausException lme;
		lme.getStream() << "[E] failed to write output to standard output channel" << 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"
			)
			{
				printHelpMessage(arginfo);
				return EXIT_SUCCESS;
			}

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