/*
    libmaus2
    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 <libmaus2/util/ArgParser.hpp>
#include <libmaus2/huffman/RLDecoder.hpp>
#include <libmaus2/rank/RunLengthBitVectorGenerator.hpp>
#include <libmaus2/rank/RunLengthBitVector.hpp>
#include <libmaus2/parallel/NumCpus.hpp>

std::string getTmpFileBase(libmaus2::util::ArgParser const & arg)
{
	std::string const tmpfilebase = arg.uniqueArgPresent("T") ? arg["T"] : libmaus2::util::ArgInfo::getDefaultTmpFileName(arg.progname);
	return tmpfilebase;
}

static uint64_t getDefaultNumThreads()
{
	return libmaus2::parallel::NumCpus::getNumLogicalProcessors();
}

void testRunLengthBitVector(libmaus2::util::ArgParser const & arg)
{
	std::string const tmpprefix = getTmpFileBase(arg);
	std::vector<std::string> Vin(1,arg[0]);
	uint64_t const bs = 64*1024;
	typedef libmaus2::huffman::RLDecoder decoder_type;
	decoder_type::unique_ptr_type RL(new decoder_type(Vin,0,1));
	uint64_t const n = RL->getN();
	decoder_type::run_type R;
	uint64_t const numthreads = arg.uniqueArgPresent("threads") ? arg.getUnsignedNumericArg<uint64_t>("threads") : getDefaultNumThreads();

	std::cerr << "[V] numthreads=" << numthreads << std::endl;

	std::string const rlfn = tmpprefix + ".rl";
	std::string const rlindexfn = tmpprefix + ".rl.index";

	libmaus2::aio::OutputStreamInstance::unique_ptr_type optr(
		new libmaus2::aio::OutputStreamInstance(rlfn)
	);

	libmaus2::rank::RunLengthBitVectorGenerator::unique_ptr_type RLgen(
		new libmaus2::rank::RunLengthBitVectorGenerator(*optr,rlindexfn,n,bs,true /* putheader */));

	uint64_t const targetpacks = 256*numthreads;
	uint64_t const packsize = (n + targetpacks - 1)/targetpacks;
	uint64_t const numpacks = packsize ? (n + packsize - 1)/packsize : 0;

	uint64_t symsum = 0;
	uint64_t runsum = 0;
	std::vector < std::pair<uint64_t,uint64_t > > VRU(numpacks);
	uint64_t TRU[2] = {0,0};
	while ( RL->decodeRun(R) )
	{
		int64_t const sym = R.sym;
		uint64_t const rlen = R.rlen;
		bool const b = (static_cast<uint64_t>(sym) >> 1) & 1;

		// std::cerr << "sym=" << sym << " rlen=" << rlen << std::endl;

		for ( uint64_t i = 0; i < rlen; ++i )
		{
			RLgen->putbit(b);

			++TRU[b];

			if ( ++symsum % packsize == 0 )
			{
				uint64_t const pack = symsum / packsize;
				VRU[pack].first = TRU[0];
				VRU[pack].second = TRU[1];
			}
		}

		runsum += 1;

		if ( runsum % (64*1024*1024) == 0 )
			std::cerr << runsum/(1024*1024) << "/" << symsum << std::endl;
	}

	assert ( symsum == n );

	std::cerr << runsum << "/" << symsum << std::endl;

	uint64_t const size = RLgen->flush();
	RLgen.reset();

	std::cerr << "size=" << size << std::endl;

	optr.reset();

	RL.reset();


	libmaus2::rank::RunLengthBitVector::unique_ptr_type pvec(
		libmaus2::rank::RunLengthBitVector::load(rlfn)
	);
	std::atomic<int> parfailed(0);
	std::atomic<uint64_t> aprocessed(0);

	#if defined(_OPENMP)
	#pragma omp parallel for num_threads(numthreads) schedule(dynamic,1)
	#endif
	for ( uint64_t t = 0; t < numpacks; ++t )
	{
		try
		{
			uint64_t const low = t*packsize;
			uint64_t const high = std::min(low+packsize,n);
			uint64_t symsum = low;
			uint64_t runsum = 0;
			decoder_type::unique_ptr_type RL(new decoder_type(Vin,symsum,1));

			uint64_t RU[2] = { VRU[t].first,VRU[t].second };

			assert ( RU[0]+RU[1] == symsum );

			decoder_type::run_type R;
			while ( RL->decodeRun(R) )
			{
				int64_t const sym = R.sym;
				uint64_t const rlen = R.rlen;
				assert ( rlen );
				bool const b = (static_cast<uint64_t>(sym) >> 1) & 1;

				for ( uint64_t i = 0; i < rlen; ++i )
				{
					uint64_t const pos = symsum++;
					bool const decb = (*pvec)[pos];

					bool const bitok = (decb == b);
					if ( ! bitok )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "[E] expected " << b << " got " << decb << " at position " << pos << " inside [" << low << "," << high << ")" << std::endl;
						lme.finish();
						throw lme;
					}

					assert ( pvec->rankm0(pos) == RU[0] );
					assert ( pvec->rankm1(pos) == RU[1] );

					++RU[decb];

					assert ( pvec->rank0(pos) == RU[0] );
					assert ( pvec->rank1(pos) == RU[1] );

					#if 0
					if ( decb )
						assert ( pvec->select1(RU[decb]-1) == pos );
					else
						assert ( pvec->select0(RU[decb]-1) == pos );
					#endif

					uint64_t const lprocessed = ++aprocessed;
					if ( lprocessed % (1024*1024) == 0 )
					{
						libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
						std::cerr << lprocessed << "/" << n << std::endl;

					}

					if ( symsum == high )
						break;
				}

				if ( symsum == high )
					break;

				runsum += 1;
			}
		}
		catch(std::exception const & ex)
		{
			parfailed = 1;
			libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
			std::cerr << ex.what() << std::endl;
		}
	}

	if ( parfailed.load() )
	{
		throw libmaus2::exception::LibMausException("[E] parallel loop failed");
	}
	else
	{
		std::cerr << n << "/" << n << std::endl;
	}
}

int main(int argc, char * argv[])
{
	try
	{
		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> Vformat;
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("b","blocksize",true));
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("t","threads",true));
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("T","",true));
		libmaus2::util::ArgParser const arg(argc,argv,Vformat);

		if ( arg.size() < 1 )
		{
			std::cerr << "[E] usage: " << argv[0] << " <in.bwt>" << std::endl;
			return EXIT_FAILURE;
		}

		testRunLengthBitVector(arg);

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