/**
    libmaus2
    Copyright (C) 2009-2021 German Tischler-Höhle
    Copyright (C) 2011-2014 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 <libmaus2/suffixsort/bwtb3m/BwtMergeParallelGapFragMerge.hpp>

uint64_t libmaus2::suffixsort::bwtb3m::BwtMergeParallelGapFragMerge::getRlLength(std::vector < std::vector < std::string > > const & bwtfilenames, uint64_t const numthreads)
{
	std::atomic<uint64_t> fs(0);
	for ( uint64_t i = 0; i < bwtfilenames.size(); ++i )
	{
		uint64_t const lfs = rl_decoder::getLength(bwtfilenames[i],numthreads);
		fs += lfs;
	}
	return fs.load();
}

std::vector<std::string> libmaus2::suffixsort::bwtb3m::BwtMergeParallelGapFragMerge::parallelGapFragMerge(
	libmaus2::util::TempFileNameGenerator & gtmpgen,
	std::vector < std::vector < std::string > > const & bwtfilenames,
	std::vector < std::vector < std::string > > const & gapfilenames,
	// std::string const & outputfilenameprefix,
	// std::string const & tempfilenameprefix,
	uint64_t const numthreads,
	uint64_t const lfblockmult,
	uint64_t const rlencoderblocksize,
	std::ostream * logstr,
	int const verbose
)
{
	if ( verbose >= 5 && logstr )
	{
		(*logstr) << "[V] entering parallelGapFragMerge for " << bwtfilenames.size() << " inputs" << std::endl;
	}

	std::vector<std::string> goutputfilenames;

	// no bwt input files, create empty bwt file
	if ( ! bwtfilenames.size() )
	{
		std::string const outputfilename = gtmpgen.getFileName() + ".bwt";

		rl_encoder rlenc(outputfilename,0 /* alphabet */,0,rlencoderblocksize);

		rlenc.flush();

		goutputfilenames = std::vector<std::string>(1,outputfilename);
	}
	// no gap files, rename input
	else if ( ! gapfilenames.size() )
	{
		assert ( bwtfilenames.size() == 1 );

		std::vector<std::string> outputfilenames;

		for ( uint64_t i = 0; i < bwtfilenames[0].size(); ++i )
		{
			std::ostringstream outputfilenamestr;
			outputfilenamestr << gtmpgen.getFileName() << '_'
				<< std::setw(4) << std::setfill('0') << i << std::setw(0)
				<< ".bwt";
			std::string const outputfilename = outputfilenamestr.str();

			// ::libmaus2::util::GetFileSize::copy(bwtfilenames[0][i],outputfilename);
			libmaus2::aio::OutputStreamFactoryContainer::rename ( bwtfilenames[0][i].c_str(), outputfilename.c_str() );

			outputfilenames.push_back(outputfilename);
		}

		goutputfilenames = outputfilenames;
	}
	// at least one gap file, merge
	else
	{
		#if defined(LIBMAUS2_SUFFIXSORT_BWTB3M_HUFGAP)
		typedef ::libmaus2::huffman::GapDecoder gapfile_decoder_type;
		#else
		typedef ::libmaus2::gamma::GammaGapDecoder gapfile_decoder_type;
		#endif

		unsigned int const albits = rl_decoder::haveAlphabetBits() ? rl_decoder::getAlBits(bwtfilenames.front()) : 0;

		uint64_t const firstblockgapfilesize = gapfilenames.size() ? gapfile_decoder_type::getLength(gapfilenames[0]) : 0;
		assert ( firstblockgapfilesize );

		// uint64_t const fs = rl_decoder::getLength(bwtfilenames);
		uint64_t const fs = getRlLength(bwtfilenames,numthreads);

		// first gap file meta information
		::libmaus2::huffman::IndexDecoderDataArray::unique_ptr_type Pfgapidda(new ::libmaus2::huffman::IndexDecoderDataArray(gapfilenames[0],numthreads));
		// first gap file
		gapfile_decoder_type::unique_ptr_type fgap(new gapfile_decoder_type(*Pfgapidda /* gapfilenames[0] */));
		// low marker
		uint64_t hlow = 0;
		// target g parts
		uint64_t const tgparts = numthreads*lfblockmult;
		// size of parts
		uint64_t const gpartsize = (fs + tgparts - 1) / tgparts;
		// maximum parts
		uint64_t const maxgparts = (fs + gpartsize - 1) / gpartsize;
		::libmaus2::autoarray::AutoArray< ::libmaus2::suffixsort::GapMergePacket> gmergepackets(maxgparts);
		// actual merge parts
		uint64_t actgparts = 0;
		uint64_t gs = 0;

		// (*logstr) << "fs=" << fs << " tgparts=" << tgparts << " gpartsize=" << gpartsize << std::endl;

		// compute number of suffixes per gblock
		while ( hlow != firstblockgapfilesize )
		{
			uint64_t const gsrest = (fs-gs);
			uint64_t const gskip = std::min(gsrest,gpartsize);
			libmaus2::huffman::KvInitResult kvinit;
			gapfile_decoder_type lgapdec(*Pfgapidda,gs + gskip, kvinit);

			// avoid small rest
			if ( fs - kvinit.kvoffset <= 1024 )
			{
				kvinit.kvoffset = fs;
				kvinit.koffset  = firstblockgapfilesize;
				// (*logstr) << "+++ end" << std::endl;
			}
			// we did not end up on a gap array element, go to the next one
			else if ( kvinit.kvtarget )
			{
				kvinit.koffset += 1;
				kvinit.voffset += kvinit.kvtarget + lgapdec.peek();
				kvinit.kvoffset += (1 + kvinit.kvtarget + lgapdec.peek());

				// there is no new suffix after the last gap element, so deduct 1 if we reached the end of the array
				if ( kvinit.koffset == firstblockgapfilesize )
					kvinit.kvoffset -= 1;

				// (*logstr) << "intermediate partial" << std::endl;
			}

			#if 1
			uint64_t const s = kvinit.kvoffset-gs;
			uint64_t const hhigh = kvinit.koffset;
			#else
			uint64_t hhigh = hlow;
			uint64_t s = 0;

			// sum up until we have reached at least gpartsize or the end of the file
			while ( hhigh != firstblockgapfilesize && s < gpartsize )
			{
				uint64_t const d = fgap->decode();
				s += (d + 1);
				hhigh += 1;
			}

			// last gap file value has no following suffix in block merged into
			if ( hhigh == firstblockgapfilesize )
				s -= 1;

			// if there is only one suffix left, then include it
			if ( hhigh+1 == firstblockgapfilesize && fgap->peek() == 0 )
			{
				fgap->decode();
				hhigh += 1;
			}

			// (*logstr) << "*** hhigh: " << hhigh << " kvinit.koffset: " << kvinit.koffset << " s=" << s << " kvinit.kvoffset-gs=" << kvinit.kvoffset-gs << std::endl;
			// kvinit.koffset, kvinit.kvoffset, rest: kvinit.kvtarget

			if ( actgparts >= maxgparts )
			{
				::libmaus2::exception::LibMausException se;
				se.getStream() << "acgtgparts=" << actgparts << " >= maxgparts=" << maxgparts << std::endl;
				se.finish();
				throw se;
				// assert ( actgparts < maxgparts );
			}
			#endif

			// save interval on first gap array and number of suffixes on this interval
			gmergepackets[actgparts++] = ::libmaus2::suffixsort::GapMergePacket(hlow,hhigh,s);

			// (*logstr) << "got packet " << gmergepackets[actgparts-1] << " firstblockgapfilesize=" << firstblockgapfilesize << std::endl;

			// add suffixes in this block to global count
			gs += s;

			// set new low marker
			hlow = hhigh;
		}

		// we should have seen all the suffixes
		assert ( gs == fs );

		#if 0
		(*logstr) << "actgparts=" << actgparts << std::endl;

		for ( uint64_t i = 0; i < actgparts; ++i )
			(*logstr) << gmergepackets[i] << std::endl;
		#endif

		// compute prefix sums over number of suffixes per gblock
		::libmaus2::autoarray::AutoArray<uint64_t> spref(actgparts+1,false);
		for ( uint64_t i = 0; i < actgparts; ++i )
			spref[i] = gmergepackets[i].s;
		//spref.prefixSums();
		libmaus2::util::PrefixSums::prefixSums(spref.begin(),spref.end());

		// compute prefix sums over number of suffixes per block used for each gpart
		::libmaus2::autoarray::AutoArray < uint64_t > bwtusedcntsacc( (actgparts+1)*(gapfilenames.size()+1), false );

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

		#if defined(_OPENMP)
		#pragma omp parallel for schedule(dynamic,1) num_threads(numthreads)
		#endif
		for ( int64_t z = 0; z < static_cast<int64_t>(spref.size()); ++z )
		{
			try
			{
				#if 0
				(*logstr) << "proc first: " << gmergepackets[z] << ",spref=" << spref[z] << std::endl;
				#endif

				// offset in first gap array
				uint64_t lspref = spref[z];

				// array of decoders
				::libmaus2::autoarray::AutoArray < gapfile_decoder_type::unique_ptr_type > gapdecs(gapfilenames.size());

				for ( uint64_t j = 0; j < gapfilenames.size(); ++j )
				{
					::libmaus2::huffman::KvInitResult kvinitresult;
					gapfile_decoder_type::unique_ptr_type tgapdecsj(
						new gapfile_decoder_type(
							gapfilenames[j],
							lspref,kvinitresult,
							numthreads
						)
					);
					gapdecs[j] = std::move(tgapdecsj);

					// key offset for block z and file j
					bwtusedcntsacc [ j * (actgparts+1) + z ] = kvinitresult.koffset;

					#if 0
					(*logstr) << "lspref=" << lspref << "," << kvinitresult << std::endl;
					#endif

					// we should be on a key if j is the first file
					if ( j == 0 )
					{
						if ( kvinitresult.kvtarget != 0 )
						{
							if ( logstr )
								(*logstr) << "j=0 " << " z=" << z << " lspref=" << lspref <<
									" kvinitresult.koffset=" << kvinitresult.koffset <<
									" kvinitresult.voffset=" << kvinitresult.voffset <<
									" kvinitresult.kvoffset=" << kvinitresult.kvoffset <<
									" kvinitresult.kvtarget=" << kvinitresult.kvtarget << std::endl;
						}
						assert ( kvinitresult.kvtarget == 0 );
					}

					// offset for next gap file:
					// sum of values up to key lspref in this file + number of values used for next key
					lspref = kvinitresult.voffset + kvinitresult.kvtarget;

				}

				#if 0
				if ( logstr )
					(*logstr) << "lspref=" << lspref << std::endl;
				#endif

				// set end pointer
				bwtusedcntsacc [ gapfilenames.size() * (actgparts+1) + z ] = lspref;
			}
			catch(std::exception const & ex)
			{
				libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
				if ( logstr )
					(*logstr) << ex.what() << std::endl;
				parfailed = 1;
			}
		}

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

		// how many suffixes of each block do we use in each gpart?
		// (turn prefix sums in count per block)
		::libmaus2::autoarray::AutoArray < uint64_t > bwtusedcnts( (gapfilenames.size()+1) * actgparts, false );
		for ( uint64_t block = 0; block < gapfilenames.size()+1; ++block )
		{
			uint64_t const * lbwtusedcntsacc = bwtusedcntsacc.begin() + block*(actgparts+1);

			for ( uint64_t i = 0; i < actgparts; ++i )
				bwtusedcnts [ block * actgparts + i ] = lbwtusedcntsacc[i+1]-lbwtusedcntsacc[i];

			assert ( lbwtusedcntsacc [ actgparts ] == rl_decoder::getLength(bwtfilenames[block],numthreads) );
		}

		// vector of output file names
		std::vector < std::string > gpartfrags(actgparts);

		// now use the block counts computed above
		#if defined(_OPENMP)
		#pragma omp parallel for schedule(dynamic,1) num_threads(numthreads)
		#endif
		for ( int64_t z = 0; z < static_cast<int64_t>(actgparts); ++z )
		{
			try
			{
				std::ostringstream ostr;
				ostr << gtmpgen.getFileName() << "_" << std::setw(4) << std::setfill('0') << z << std::setw(0) << ".bwt";
				std::string const gpartfrag = ostr.str();
				::libmaus2::util::TempFileRemovalContainer::addTempFile(gpartfrag);
				gpartfrags[z] = gpartfrag;

				#if 0
				if ( logstr )
					(*logstr) << gmergepackets[z] << ",spref=" << spref[z] << std::endl;
				#endif
				uint64_t lspref = spref[z];
				::libmaus2::autoarray::AutoArray < gapfile_decoder_type::unique_ptr_type > gapdecoders(gapfilenames.size());
				::libmaus2::autoarray::AutoArray< uint64_t > gapcur(gapfilenames.size());

				// set up gap file decoders at the proper offsets
				for ( uint64_t j = 0; j < gapfilenames.size(); ++j )
				{
					// sum up number of suffixes in later blocks for this gpart
					uint64_t suflat = 0;
					for ( uint64_t k = j+1; k < bwtfilenames.size(); ++k )
						suflat += bwtusedcnts [ k*actgparts + z ];

					::libmaus2::huffman::KvInitResult kvinitresult;
					gapfile_decoder_type::unique_ptr_type tgapdecodersj(
						new gapfile_decoder_type(
							gapfilenames[j],
							lspref,kvinitresult,
							numthreads
						)
					);
					gapdecoders[j] = std::move(tgapdecodersj);
					if ( suflat )
						gapcur[j] = gapdecoders[j]->decode();
					else
						gapcur[j] = 0;

					lspref = kvinitresult.voffset + kvinitresult.kvtarget;

					if ( j == 0 )
						assert ( kvinitresult.kvtarget == 0 );
				}

				::libmaus2::autoarray::AutoArray < uint64_t > bwttowrite(bwtfilenames.size(),false);
				::libmaus2::autoarray::AutoArray < rl_decoder::unique_ptr_type > bwtdecoders(bwtfilenames.size());

				for ( uint64_t j = 0; j < bwtfilenames.size(); ++j )
				{
					uint64_t const bwtoffset = bwtusedcntsacc [ j * (actgparts+1) + z ];
					bwttowrite[j] = bwtusedcnts [ j * actgparts + z ];

					rl_decoder::unique_ptr_type tbwtdecodersj(
						new rl_decoder(bwtfilenames[j],bwtoffset,numthreads)
					);
					bwtdecoders[j] = std::move(tbwtdecodersj);

					#if 0
					if ( logstr )
						(*logstr) << "block=" << j << " offset=" << bwtoffset << " bwttowrite=" << bwttowrite[j] << std::endl;
					#endif
				}

				uint64_t const totalbwt = std::accumulate(bwttowrite.begin(),bwttowrite.end(),0ull);

				rl_encoder bwtenc(gpartfrag,albits /* alphabet */,totalbwt,rlencoderblocksize);

				// start writing loop
				uint64_t written = 0;
				while ( written < totalbwt )
				{
					// determine file we next read/decode from
					// this is the leftmost one with gap value 0 and still data to write
					uint64_t writeindex = bwtdecoders.size()-1;
					for (
						uint64_t i = 0;
						i < gapcur.size();
						++i
					)
						if (
							(! gapcur[i])
							&&
							(bwttowrite[i])
						)
						{
							writeindex = i;
							break;
						}

					// sanity check
					if ( ! bwttowrite[writeindex] )
					{
						assert ( bwttowrite[writeindex] );
					}

					// adjust counters
					written++;
					bwttowrite[writeindex]--;
					// get next gap value if block is not the last one
					if ( bwttowrite[writeindex] && writeindex < gapcur.size() )
						gapcur[writeindex] = gapdecoders[writeindex]->decode();

					// copy symbol
					uint64_t const sym = bwtdecoders[writeindex]->decode();
					bwtenc.encode(sym);

					// adjust gap values of blocks to the left
					for ( uint64_t i = 0; i < writeindex; ++i )
						if ( bwttowrite[i] )
						{
							assert ( gapcur[i] > 0 );
							gapcur[i]--;
						}
				}
				//if ( logstr )
					// (*logstr) << "(1)";

				// all data should have been written now
				for ( uint64_t i = 0; i < bwttowrite.size(); ++i )
					assert ( !bwttowrite[i] );

				bwtenc.flush();
			}
			catch(std::exception const & ex)
			{
				libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
				if ( logstr )
					(*logstr) << ex.what() << std::endl;
				parfailed = 1;
			}
		}

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

		goutputfilenames = gpartfrags;
	}

	if ( verbose >= 5 && logstr )
	{
		(*logstr) << "[V] checking length consistency" << std::endl;
	}

	assert (
		rl_decoder::getLength(goutputfilenames,numthreads) ==
		rl_decoder::getLength(bwtfilenames,numthreads)
	);

	if ( verbose >= 5 && logstr )
	{
		(*logstr) << "[V] leaving parallelGapFragMerge for " << bwtfilenames.size() << " inputs" << std::endl;
	}

	return goutputfilenames;
}
