/*
    libmaus2
    Copyright (C) 2021 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/>.
*/
#if ! defined(LIBMAUS2_VCF_TABIXINDEX_HPP)
#define LIBMAUS2_VCF_TABIXINDEX_HPP

#include <libmaus2/bambam/BamAlignmentReg2Bin.hpp>
#include <libmaus2/lz/BgzfChunkInputStream.hpp>
#include <libmaus2/aio/InputStreamInstance.hpp>
#include <libmaus2/lz/PlainOrGzipStream.hpp>
#include <cassert>

namespace libmaus2
{
	namespace vcf
	{
		struct TabixIndex
		{
			int32_t n_ref;
			int32_t format;
			int32_t col_seq;
			int32_t col_beg;
			int32_t col_end;
			int32_t meta;
			int32_t skip;
			int32_t l_nm;
			std::vector<std::string> Vnames;
			std::map<std::string,std::size_t> Mnames;
			// positions are zero based if bedCoord is true
			bool bedCoord;
			uint16_t shortformat;

			enum tabix_format {
				tabix_format_generic = 0,
				tabix_format_SAM = 1,
				tabix_format_VCF = 2
			};

			std::map < int32_t, std::map< int32_t, std::vector<libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk> > > Mbin;

			/**
			 * get VCF stream for all chunks involved in a region. This does not imply the stream will not contain any records
			 * outside this region. The parameter pos is expected to be 1 based (as per VCF convention)
			 **/
			std::shared_ptr<std::istream> getVCFRegionStream(std::string const & fn, int32_t const ref, int64_t const pos, int64_t const len) const
			{
				// check prerequisites
				if ( shortformat != tabix_format_VCF || bedCoord || col_seq != 1 || col_beg != 2 || col_end != 0 )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] TabixIndex::getVCFRegionStream: index is unsuitable for accessing VCF" << std::endl;
					lme.getStream() << "[E] shortformat=" << shortformat << " tabix_format_VCF=" << tabix_format_VCF << std::endl;
					lme.getStream() << "[E] bedCoord=" << bedCoord << std::endl;
					lme.getStream() << "[E] col_seq=" << col_seq << std::endl;
					lme.getStream() << "[E] col_beg=" << col_beg << std::endl;
					lme.getStream() << "[E] col_end=" << col_end << std::endl;
					lme.finish();
					throw lme;
				}

				std::vector<libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk> Vchunk;

				// any bins for ref?
				auto const it_ref = Mbin.find(ref);
				if ( it_ref != Mbin.end() )
				{
					// convert 1 based to 0 based position
					int64_t const beg = pos - 1;
					// compute end of zero based interval (excluded)
					int64_t const end = beg + len;

					// get all bins involved
					libmaus2::autoarray::AutoArray<int> AB;
					uint64_t const AB_o = libmaus2::bambam::BamAlignmentReg2Bin::reg2bins(beg,end,AB);

					// extract all relevant chunks
					for ( uint64_t ib = 0; ib < AB_o; ++ib )
					{
						int const bin = AB[ib];

						auto const it_bin = it_ref->second.find(bin);

						if ( it_bin != it_ref->second.end() )
						{
							std::copy(
								it_bin->second.begin(),
								it_bin->second.end(),
								std::back_inserter(Vchunk)
							);
						}
					}
				}

				// sort chunks
				std::sort(Vchunk.begin(),Vchunk.end());

				// merge chunks
				std::size_t low = 0;
				std::size_t o = 0;
				while ( low < Vchunk.size() )
				{
					std::size_t high = low+1;

					while ( high < Vchunk.size() && Vchunk[high-1].cnk_end >= Vchunk[high].cnk_beg )
						++high;

					Vchunk[o++] = libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk(Vchunk[low].cnk_beg,Vchunk[high-1].cnk_end);

					low = high;
				}

				Vchunk.resize(o);

				std::shared_ptr<std::istream> tinptr(new libmaus2::aio::InputStreamInstance(fn));
				std::shared_ptr<std::istream> inptr(new libmaus2::lz::BgzfChunkInputStream(tinptr,Vchunk));

				return inptr;
			}

			bool haveRef(std::string const & ref) const
			{
				return Mnames.find(ref) != Mnames.end();
			}

			std::shared_ptr<std::istream> getVCFRegionStream(std::string const & fn, std::string const & ref, int64_t const pos, int64_t const len) const
			{
				auto const ref_it = Mnames.find(ref);

				if ( ref_it == Mnames.end() )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] TabixIndex::getVCFRegionStream: reference name " << ref << " is not contained in index" << std::endl;
					lme.finish();
					throw lme;
				}

				std::shared_ptr<std::istream> ptr(getVCFRegionStream(fn,ref_it->second,pos,len));

				return ptr;
			}

			// get stream for decoding a single bin
			std::shared_ptr<std::istream> getBinStream(std::string const & fn, int32_t const ref, int32_t const bin) const
			{
				std::vector<libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk> Vchunk;

				auto const it_ref = Mbin.find(ref);
				if ( it_ref != Mbin.end() )
				{
					auto const it_bin = it_ref->second.find(bin);

					if ( it_bin != it_ref->second.end() )
					{
						Vchunk = it_bin->second;
					}
				}

				std::shared_ptr<std::istream> tinptr(new libmaus2::aio::InputStreamInstance(fn));
				std::shared_ptr<std::istream> inptr(new libmaus2::lz::BgzfChunkInputStream(tinptr,Vchunk));

				return inptr;
			}

			// get vector of bins for a reference sequence
			std::vector<int32_t> getBinVector(int32_t const ref) const
			{
				std::vector<int32_t> V;

				auto const it_ref = Mbin.find(ref);

				if ( it_ref != Mbin.end() )
					for ( auto P : it_ref->second )
						V.emplace_back(P.first);
				return V;
			}

			template<typename stream_type, typename integer_type>
			integer_type getLittleEndian(stream_type & in)
			{
				integer_type v = integer_type(0);

				for ( std::size_t i = 0; i < sizeof(integer_type); ++i )
				{
					int const lv = in.get();

					if ( (! in) || (lv == std::istream::traits_type::eof()) )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "[E] TabixIndex::getLittleEndian: found unexpected EOF" << std::endl;
						lme.finish();
						throw lme;
					}

					v |= static_cast<integer_type>(lv) << (8*i);
				}

				return v;
			}

			// load character data
			std::string loadData(std::istream & in, std::size_t const n)
			{
				std::vector<char> Vdata(n);

				for ( std::size_t i = 0; i < n; ++i )
				{
					int const c = in.get();

					if ( (! in) || (c == std::istream::traits_type::eof()) )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "[E] TabixIndex::init: found unexpected EOF" << std::endl;
						lme.finish();
						throw lme;
					}

					Vdata[i] = static_cast<char>(c);
				}

				return std::string(Vdata.begin(),Vdata.end());
			}

			void init(std::istream & rin)
			{
				libmaus2::lz::PlainOrGzipStream in(rin);

				std::string const fmagic(loadData(in,4));

				if ( fmagic != std::string("TBI\1") )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] TabixIndex::init: wrong file magic, found " << fmagic << std::endl;
					lme.finish();
					throw lme;
				}

				n_ref = getLittleEndian<std::istream,int32_t>(in);
				format = getLittleEndian<std::istream,int32_t>(in);
				col_seq = getLittleEndian<std::istream,int32_t>(in);
				col_beg = getLittleEndian<std::istream,int32_t>(in);
				col_end = getLittleEndian<std::istream,int32_t>(in);
				meta = getLittleEndian<std::istream,int32_t>(in);
				skip = getLittleEndian<std::istream,int32_t>(in);
				l_nm = getLittleEndian<std::istream,int32_t>(in);
				bedCoord = (format & 0x10000) != 0;
				shortformat = (format & 0xffff);

				// std::cerr << "col_seq=" << col_seq << " col_beg=" << col_beg << " col_end=" << col_end << " bedCoord=" << bedCoord << std::endl;

				std::string const Vdata(loadData(in,l_nm));

				if ( Vdata.size() )
				{
					if ( Vdata.back() )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "[E] TabixIndex::init: ref seq name array is not null terminated" << std::endl;
						lme.finish();
						throw lme;
					}

					std::size_t low = 0;

					while ( low < Vdata.size() )
					{
						std::size_t high = low;
						while ( high < Vdata.size() && Vdata[high] )
							++high;

						assert ( high < Vdata.size() && Vdata[high] == 0 );

						std::string const refname =
							std::string(
								Vdata.begin()+low,
								Vdata.begin()+high
							);

						if ( Mnames.find(refname) == Mnames.end() )
							Mnames[refname] = Vnames.size();
						else
						{
							libmaus2::exception::LibMausException lme;
							lme.getStream() << "[E] TabixIndex::init: reference sequence name " << refname << " is not unique" << std::endl;
							lme.finish();
							throw lme;
						}

						Vnames.push_back(refname);


						low = high+1;
					}
				}

				if ( static_cast<int32_t>(Vnames.size()) != n_ref )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] TabixIndex::init: ref seq name array size " << Vnames.size() << " is not consistent with n_ref=" << n_ref << std::endl;
					lme.finish();
					throw lme;
				}

				#if 0
				for ( std::size_t i = 0; i < Vnames.size(); ++i )
					std::cerr << "Vnames[i]=" << Vnames[i] << std::endl;

				// std::cerr << "bedCoord=" << bedCoord << std::endl;
				#endif

				for ( int32_t r = 0; r < n_ref; ++r )
				{
					int32_t const n_bin = getLittleEndian<std::istream,int32_t>(in);

					std::map < int32_t, std::vector<libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk> > & Mref = Mbin[r];

					for ( int32_t i = 0; i < n_bin; ++i )
					{
						int32_t const bin = getLittleEndian<std::istream,int32_t>(in);
						uint32_t const n_chunk = getLittleEndian<std::istream,uint32_t>(in);

						std::vector<libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk> & V = Mref[bin];

						for ( uint32_t c = 0; c < n_chunk; ++c )
						{
							uint64_t const cnk_beg = getLittleEndian<std::istream,uint64_t>(in);
							uint64_t const cnk_end = getLittleEndian<std::istream,uint64_t>(in);

							V.push_back(libmaus2::lz::BgzfChunkInputStreamBuffer::Chunk(cnk_beg,cnk_end));
						}

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

					int32_t const n_intv = getLittleEndian<std::istream,int32_t>(in);

					for ( int32_t i = 0; i < n_intv; ++i )
					{
						/* uint64_t const ioff = */ getLittleEndian<std::istream,uint64_t>(in);
					}
				}
			}

			/**
			 * constructor by name of tabix file (not the name of the indexed file)
			 **/
			TabixIndex(std::string const & fn)
			{
				libmaus2::aio::InputStreamInstance ISI(fn);
				init(ISI);
			}
		};
	}
}
#endif
