/*
    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_PROTEIN_PROTEINALIGNMENT_HPP)
#define LIBMAUS2_PROTEIN_PROTEINALIGNMENT_HPP

#include <memory>
#include <string>
#include <sstream>
#include <libmaus2/exception/LibMausException.hpp>
#include <libmaus2/autoarray/AutoArray.hpp>
#include <libmaus2/lcs/AlignmentTraceContainer.hpp>

#if defined(LIBMAUS2_HAVE_PARASAIL)
#include <parasail.h>
#include <parasail/matrices/blosum62.h>
#endif

namespace libmaus2
{
	namespace protein
	{
		struct ProteinAlignment : libmaus2::lcs::AlignmentTraceContainer
		{
			int64_t gapopen;
			int64_t gapextend;

			ProteinAlignment(int64_t const rgapopen = 11, int64_t const rgapextend = 1) : gapopen(rgapopen), gapextend(rgapextend)
			{
			}

			struct Result
			{
				std::size_t abpos;
				std::size_t aepos;
				std::size_t bbpos;
				std::size_t bepos;
				int64_t score;
				std::string cig;

				Result(
					std::size_t rabpos = 0,
					std::size_t raepos = 0,
					std::size_t rbbpos = 0,
					std::size_t rbepos = 0,
					int64_t rscore = 0,
					std::string const & rcig = std::string()
				) : abpos(rabpos), aepos(raepos), bbpos(rbbpos), bepos(rbepos), score(rscore), cig(rcig)
				{

				}

				std::size_t enumerateMatches(
					libmaus2::autoarray::AutoArray<
						std::pair<uint64_t,uint64_t>
					> & P
				) const
				{
					std::string::const_iterator it = cig.begin();

					std::size_t pos_a = abpos;
					std::size_t pos_b = bbpos;

					uint64_t o = 0;

					while ( it != cig.end() )
					{
						std::size_t l = 0;

						while ( it != cig.end() && ::isdigit(*it) )
						{

							l *= 10;
							l += (*(it++)) - '0';
						}

						if ( it == cig.end() )
						{
							libmaus2::exception::LibMausException lme;
							lme.getStream() << "[E] unable to parse " << cig << std::endl;
							lme.finish();
							throw lme;
						}

						char const op = (*(it++));

						for ( std::size_t i = 0; i < l; ++i )
						{
							switch ( op )
							{
								case '=':
									P.push(o,std::make_pair(pos_a++,pos_b++));
									break;
								case 'X':
									pos_a+=1;
									pos_b+=1;
									break;
								case 'I':
									pos_b+=1;
									break;
								case 'D':
									pos_a+=1;
									break;
								default:
								{
									libmaus2::exception::LibMausException lme;
									lme.getStream() << "[E] unhandled operation " << op << " in " << cig << std::endl;
									lme.finish();
									throw lme;
								}
							}
						}
					}

					return o;
				}


				std::vector < std::pair<std::size_t,std::size_t> >  enumerateMatches() const
				{
					libmaus2::autoarray::AutoArray<
						std::pair<uint64_t,uint64_t>
					> P;
					auto const o = enumerateMatches(P);
					return std::vector<std::pair<std::size_t,std::size_t> >(P.begin(),P.begin()+o);
				}

				std::string toString() const
				{
					std::ostringstream ostr;
					ostr << "([" << abpos << "," << aepos << "),[" << bbpos << "," << bepos << ")," << score << "," << cig << ")";
					return ostr.str();
				}

				bool operator<(Result const & R) const
				{
					     if ( abpos != R.abpos )
						return abpos < R.abpos;
					else if ( aepos != R.aepos )
						return aepos < R.aepos;
					else if ( bbpos != R.bbpos )
						return bbpos < R.bbpos;
					else if ( bepos != R.bepos )
						return bepos < R.bepos;
					else if ( score != R.score )
						return score < R.score;
					else if ( cig != R.cig )
						return cig < R.cig;
					else
						return false;
				}

				bool operator==(Result const & R) const
				{
					return
						abpos == R.abpos
						&&
						aepos == R.aepos
						&&
						bbpos == R.bbpos
						&&
						bepos == R.bepos
						&&
						score == R.score
						&&
						cig == R.cig;
				}
			};

			template<typename matrix_type>
			std::size_t computeMaxExtend(char const * ba, char const * be, matrix_type const * matrix) const
			{
				int64_t idscore = 0;

				for ( char const * b = ba; b != be; ++b )
				{
					int const csym = matrix->mapper[static_cast<int>(*b)];
					int64_t const symscore = matrix->matrix[csym*matrix->size+csym];
					idscore += symscore;
				}

				/*
				gapopen + i * gapextend <= idscore
				i * gapextend <= idscore-gapopen
				i             <= (idscore-gapopen)/gapextend
				*/
				std::size_t const maxextend =
					std::max
					(
						static_cast<int64_t>(
							(idscore - gapopen)/gapextend
						),
						static_cast<int64_t>(0)
					)
					;

				return maxextend;
			}

			Result align(
				#if defined(LIBMAUS2_HAVE_PARASAIL)
				char const * aa,
				char const * ae,
				std::size_t const seedposa,
				char const * ba,
				char const * be,
				std::size_t const seedposb,
				int64_t const rmaxextend = -1
				#else
				char const *,
				char const *,
				std::size_t const,
				char const *,
				char const *,
				std::size_t const,
				int64_t const = -1
				#endif
			)
			{
				#if defined(LIBMAUS2_HAVE_PARASAIL)
				auto const matrix = &parasail_blosum62;
				std::size_t const maxextend = (rmaxextend > 0) ? rmaxextend : computeMaxExtend(ba,be,matrix);

				std::ptrdiff_t const na = ae - aa;
				std::ptrdiff_t const nb = be - ba;
				std::size_t const aforw =
					std::min(
						static_cast<std::size_t>(na - seedposa),
						nb + maxextend
					)
					;
				std::size_t const bforw = nb - seedposb;

				assert ( aforw );
				assert ( bforw );

				std::shared_ptr<parasail_result_t> sresfor(
					parasail_sg_qe_trace(
						aa+seedposa,
						aforw,
						ba+seedposb,
						bforw,
						gapopen,
						gapextend,
						matrix
					),
					[](auto p){parasail_result_free(p);}
				)
				;
				std::shared_ptr<parasail_cigar_t> scigfor(
					parasail_result_get_cigar(
						sresfor.get(),
						aa+seedposa,
						aforw,
						ba+seedposb,
						bforw,
						matrix),
					[](auto p){parasail_cigar_free(p);}
				)
				;

				// compute number of AS clipped of at the end
				std::size_t aclipend = 0;
				for ( int i = scigfor->len-1; i >= 0 && (parasail_cigar_decode_op(scigfor->seq[i]) == 'I'); --i )
					aclipend += parasail_cigar_decode_len(scigfor->seq[i]);

				// compute end position on A and B
				std::size_t const alenforw = aforw - aclipend;
				std::size_t const aepos = seedposa + alenforw;
				std::size_t const bepos = nb;

				std::size_t const aback =
					std::min(
						static_cast<std::size_t>(aepos),
						nb + maxextend
					)
					;
				std::size_t const astart = aepos - aback;

				assert ( aepos > astart );
				assert ( bepos > 0 );

				std::shared_ptr<parasail_result_t> sresrev(
					parasail_sg_qb_trace(
						aa + astart,
						aepos - astart,
						ba,
						bepos,
						gapopen,
						gapextend,
						matrix
					),
					[](auto p){parasail_result_free(p);}
				)
				;
				std::shared_ptr<parasail_cigar_t> scigrev(
					parasail_result_get_cigar(
						sresrev.get(),
						aa + astart,
						aepos - astart,
						ba,
						bepos,
						matrix),
					[](auto p){parasail_cigar_free(p);}
				)
				;

				std::size_t aclipfront = 0;

				for ( int i = 0; i < scigrev->len && (parasail_cigar_decode_op(scigrev->seq[i]) == 'I'); ++i )
					aclipfront += parasail_cigar_decode_len(scigrev->seq[i]);

				std::size_t const abpos = astart + aclipfront;
				std::size_t const bbpos = 0;

				std::shared_ptr<parasail_result_t> sresfull(
					parasail_nw_trace(
						ba + bbpos,
						bepos - bbpos,
						aa + abpos,
						aepos - abpos,
						gapopen,
						gapextend,
						matrix
					),
					[](auto p){parasail_result_free(p);}
				)
				;

				std::shared_ptr<parasail_cigar_t> scig(
					parasail_result_get_cigar(
						sresfull.get(),
						ba + bbpos,
						bepos - bbpos,
						aa + abpos,
						aepos - abpos,
						matrix),
					[](auto p){parasail_cigar_free(p);}
				)
				;

				std::size_t oplen = 0;
				for ( int i = 0; i < scig->len; ++i )
				{
					uint32_t const c = scig->seq[i];
					char const op = parasail_cigar_decode_op(c);
					uint32_t const len = parasail_cigar_decode_len(c);

					switch ( op )
					{
						case '=':
						case 'X':
						case 'D':
						case 'I':
							break;
						default:
						{
							libmaus2::exception::LibMausException lme;
							lme.getStream() << "[E] invalid cigar op code " << static_cast<char>(op) << std::endl;
							lme.finish();
							throw lme;
						}
					}

					oplen += len;
				}

				libmaus2::lcs::AlignmentTraceContainer & ATC = *this;

				if ( ATC.capacity() < oplen )
					ATC.resize(oplen);

				ATC.reset();

				for ( int i = 0; i < scig->len; ++i )
				{
					uint32_t const c = scig->seq[scig->len-i-1];
					char const op = parasail_cigar_decode_op(c);
					uint32_t len = parasail_cigar_decode_len(c);

					while ( len-- )
					{
						switch ( op )
						{
							case '=':
								*(--ATC.ta) = ATC.STEP_MATCH;
								break;
							case 'X':
								*(--ATC.ta) = ATC.STEP_MISMATCH;
								break;
							case 'I':
								*(--ATC.ta) = ATC.STEP_INS;
								break;
							case 'D':
								*(--ATC.ta) = ATC.STEP_DEL;
								break;
						}
					}
				}

				std::ostringstream cigstream;
				for ( int i = 0; i < scig->len; ++i )
				{
					uint32_t const c = scig->seq[i];
					char const op = parasail_cigar_decode_op(c);
					uint32_t const len = parasail_cigar_decode_len(c);
					cigstream << len << op;
				}
				std::string const cigstr = cigstream.str();

				return Result(abpos,aepos,bbpos,bepos,sresfull->score,cigstr);
				#else
				libmaus2::exception::LibMausException lme;
				lme.getStream() << "[E] ProteinAlignment::align: libmaus2 is compiled without support for parasail" << std::endl;
				lme.finish();
				throw lme;
				#endif
			}
		};
	}
}
#endif
