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

#include <libmaus2/protein/Blosum62.hpp>
#include <libmaus2/autoarray/AutoArray.hpp>
#include <libmaus2/util/NumberSerialisation.hpp>

namespace libmaus2
{
	namespace protein
	{
		template<typename _score_type>
		struct Align
		{
			typedef _score_type score_type;

			int64_t const deletionscore;
			int64_t const insertionscore;
			libmaus2::autoarray::AutoArray<int64_t> AS;
			libmaus2::autoarray::AutoArray<unsigned char> AT;
			libmaus2::autoarray::AutoArray<unsigned char> ATO;

			enum trace_step { trace_m = 0, trace_d = 1, trace_i = 2, trace_x = 3 };

			Align(
				int64_t const rdeletionscore = -8,
				int64_t const rinsertionscore = -8
			) : deletionscore(rdeletionscore), insertionscore(rinsertionscore)
			{

			}

			int64_t getDeletionScore() const
			{
				return deletionscore;
			}

			int64_t getInsertionScore() const
			{
				return insertionscore;
			}

			static std::pair<int64_t,int64_t> antiDiagToBlock(std::pair<int64_t,int64_t> const P)
			{
				return std::pair<int64_t,int64_t>(P.first-P.second,P.second);
			}

			static std::pair<int64_t,int64_t> blockToAntiDiag(std::pair<int64_t,int64_t> const P)
			{
				return std::pair<int64_t,int64_t>(P.first+P.second,P.second);
			}

			static uint64_t antiDiagOffset(uint64_t const anti)
			{
				return ((anti+1)*anti)/2;
			}

			static uint64_t antiDiagOffset(uint64_t const anti, uint64_t const diag)
			{
				return antiDiagOffset(anti) + diag;
			}

			static char traceToChar(trace_step const t)
			{
				switch ( t )
				{
					case trace_m: return 'M';
					case trace_i: return 'I';
					case trace_d: return 'D';
					case trace_x: default: return 'X';
				}
			}

			std::ostream & printTrace(std::ostream & out, uint64_t const ATOo) const
			{
				for ( uint64_t i = 0; i < ATOo; ++i )
					out << traceToChar(static_cast<trace_step>(ATO[i]));
				return out;
			}

			std::string getTrace(uint64_t const ATOo) const
			{
				std::ostringstream ostr;
				printTrace(ostr,ATOo);
				return ostr.str();
			}

			std::pair<uint64_t,uint64_t> getStringLengthUsed(uint64_t const ATOo) const
			{
				uint64_t la = 0;
				uint64_t lb = 0;

				for ( uint64_t i = 0; i < ATOo; ++i )
					switch ( static_cast<trace_step>(ATO[i]) )
					{
						case trace_m: ++la; ++lb; break;
						case trace_i:       ++lb; break;
						case trace_d: ++la;       break;
						case trace_x: default:  break;
					}

				return std::pair<uint64_t,uint64_t>(la,lb);
			}

			uint64_t enumerateMatches(
				uint64_t const ATOo,
				uint64_t const pos_a,
				uint64_t const pos_b,
				libmaus2::autoarray::AutoArray< std::pair<uint64_t,uint64_t> > & O
			) const
			{
				uint64_t p_a = pos_a;
				uint64_t p_b = pos_b;
				uint64_t o = 0;

				for ( uint64_t i = 0; i < ATOo; ++i )
					switch ( static_cast<trace_step>(ATO[i]) )
					{
						case trace_m: O.push(o,std::make_pair(p_a,p_b)); ++p_a; ++p_b; break;
						case trace_i:        ++p_b; break;
						case trace_d: ++p_a;        break;
						case trace_x: default:      break;
					}

				return o;
			}

			template<typename iterator_a, typename iterator_b>
			int64_t getScore(uint64_t const ATOo, iterator_a a, iterator_b b) const
			{
				int64_t score = 0;

				for ( uint64_t i = 0; i < ATOo; ++i )
					switch ( static_cast<trace_step>(ATO[i]) )
					{
						case trace_m: score += score_type::getScore(*a,*b); ++a; ++b; break;
						case trace_i: score += getInsertionScore();              ++b; break;
						case trace_d: score += getDeletionScore();          ++a;      break;
						case trace_x: default:  break;
					}

				return score;
			}

			static std::string lformatNumber(uint64_t const n, uint64_t const len)
			{
				std::ostringstream str;
				str << std::setw(len) << std::setfill(' ') << n;
				return str.str();
			}

			template<typename iterator_a, typename iterator_b>
			std::ostream & printTrace(
				std::ostream & out,
				uint64_t const ATOo,
				iterator_a a,
				iterator_b b,
				uint64_t const cols = std::numeric_limits<uint64_t>::max(),
				std::string const & label_a = std::string(),
				std::string const & label_b = std::string(),
				uint64_t const rpos_a = 0,
				uint64_t const rpos_b = 0
			)
			{
				std::ostringstream astr;
				std::ostringstream bstr;
				std::ostringstream cstr;

				uint64_t const label_a_size = label_a.size();
				uint64_t const label_b_size = label_b.size();
				uint64_t const max_label_size = std::max(label_a_size,label_b_size);
				uint64_t const max_label_padded = max_label_size ? (max_label_size+1) : 0;
				uint64_t const pad_a_size = max_label_padded - label_a_size;
				uint64_t const pad_b_size = max_label_padded - label_b_size;

				std::string const print_label_a = label_a + std::string(pad_a_size,' ');
				std::string const print_label_b = label_b + std::string(pad_b_size,' ');
				std::string const indent = std::string(max_label_padded,' ');

				int64_t score = 0;

				for ( uint64_t i = 0; i < ATOo; ++i )
					switch ( static_cast<trace_step>(ATO[i]) )
					{
						case trace_m: cstr << ((*a==*b) ? '=' : 'X'); score += score_type::getScore(*a,*b); astr << *(a++); bstr << *(b++); break;
						case trace_i: cstr << 'I'                   ; score += getInsertionScore();         astr << ' '   ; bstr << *(b++); break;
						case trace_d: cstr << 'D'                   ; score += getDeletionScore();          astr << *(a++); bstr << ' '   ; break;
						case trace_x: default:  break;
					}

				std::string const sa = astr.str();
				std::string const sb = bstr.str();
				std::string const sc = cstr.str();

				assert ( sa.size() == sb.size() );
				assert ( sb.size() == sc.size() );

				uint64_t r = sa.size();
				uint64_t pos_a = rpos_a;
				uint64_t pos_b = rpos_b;
				char const * ca = sa.c_str();
				char const * cb = sb.c_str();
				char const * cc = sc.c_str();

				uint64_t maxpos_a = rpos_a;
				uint64_t maxpos_b = rpos_b;

				while ( r )
				{
					uint64_t const s = std::min(cols,r);

					maxpos_a = pos_a;
					maxpos_b = pos_b;

					for ( uint64_t i = 0; i < s; ++i )
					{
						char const c = cc[i];

						switch ( c )
						{
							case '=':
							case 'X':
								pos_a++;
								pos_b++;
								break;
							case 'I':
								pos_b++;
								break;
							case 'D':
								pos_a++;
								break;
							default:
								break;
						}
					}

					cc += s;
					r -= s;
				}

				uint64_t const maxpos = std::max(maxpos_a,maxpos_b);
				uint64_t const numlen = libmaus2::util::NumberSerialisation::formatNumber(maxpos,0).size();

				ca = sa.c_str();
				cb = sb.c_str();
				cc = sc.c_str();
				r = sa.size();
				pos_a = rpos_a;
				pos_b = rpos_b;

				while ( r )
				{
					uint64_t const s = std::min(cols,r);

					std::string const spos_a = lformatNumber(pos_a,numlen);
					std::string const spos_b = lformatNumber(pos_b,numlen);
					std::string const spos_c(numlen,' ');

					for ( uint64_t i = 0; i < s; ++i )
					{
						char const c = cc[i];

						switch ( c )
						{
							case '=':
							case 'X':
								pos_a++;
								pos_b++;
								break;
							case 'I':
								pos_b++;
								break;
							case 'D':
								pos_a++;
								break;
							default:
								break;
						}
					}

					out << print_label_a << " " << spos_a << " "; out.write(ca,s); out.put('\n'); ca += s;
					out << print_label_b << " " << spos_b << " ";; out.write(cb,s); out.put('\n'); cb += s;
					out << indent        << " " << spos_c << " "; out.write(cc,s); out.put('\n'); cc += s;

					r -= s;
				}

				return out;
			}

			template<
				typename iterator_a,
				typename iterator_b
			>
			uint64_t operator()(iterator_a aa, iterator_a ae, iterator_b ba, iterator_b be)
			{
				int64_t const la = ae-aa;
				int64_t const lb = be-ba;

				uint64_t o = 0;
				uint64_t ot = 0;
				AS.push(o,0);
				AT.push(ot,trace_x);

				for ( int64_t antidiag = 1; antidiag <= la+lb; ++antidiag )
					for ( int64_t diag = 0; diag <= antidiag; ++ diag )
					{
						std::pair<int64_t,int64_t> const antiDiagPair(antidiag,diag);
						std::pair<int64_t,int64_t> const blockPair = antiDiagToBlock(antiDiagPair);

						assert ( blockPair.first >= 0 );
						assert ( blockPair.second >= 0 );

						int64_t score = 0;
						unsigned char trace = trace_x;

						if ( blockPair.first <= la && blockPair.second <= lb )
						{
							bool valid = false;

							if ( blockPair.first )
							{
								std::pair<int64_t,int64_t> const topBlockPair(blockPair.first-1,blockPair.second);
								std::pair<int64_t,int64_t> const topAntiDiagPair = blockToAntiDiag(topBlockPair);
								assert ( topAntiDiagPair.first < antiDiagPair.first );
								assert ( topAntiDiagPair.first >= 0 );
								assert ( topAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(topAntiDiagPair.first,topAntiDiagPair.second);
								assert ( offset < o );
								int64_t const topscore = AS[offset] + getDeletionScore();

								if ( valid )
								{
									if ( topscore > score )
									{
										score = topscore;
										trace = trace_d;
									}
								}
								else
								{
									score = topscore;
									trace = trace_d;
									valid = true;
								}
							}
							if ( blockPair.second )
							{
								std::pair<int64_t,int64_t> const leftBlockPair(blockPair.first,blockPair.second-1);
								std::pair<int64_t,int64_t> const leftAntiDiagPair = blockToAntiDiag(leftBlockPair);
								assert ( leftAntiDiagPair.first < antiDiagPair.first );
								assert ( leftAntiDiagPair.first >= 0 );
								assert ( leftAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(leftAntiDiagPair.first,leftAntiDiagPair.second);
								assert ( offset < o );
								int64_t const leftscore = AS[offset] + getInsertionScore();

								if ( valid )
								{
									if ( leftscore > score )
									{
										score = leftscore;
										trace = trace_i;
									}
								}
								else
								{
									score = leftscore;
									trace = trace_i;
									valid = true;
								}
							}
							if ( blockPair.first && blockPair.second )
							{
								std::pair<int64_t,int64_t> const diagBlockPair(blockPair.first-1,blockPair.second-1);
								std::pair<int64_t,int64_t> const diagAntiDiagPair = blockToAntiDiag(diagBlockPair);
								assert ( diagAntiDiagPair.first < antiDiagPair.first );
								assert ( diagAntiDiagPair.first >= 0 );
								assert ( diagAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(diagAntiDiagPair.first,diagAntiDiagPair.second);
								assert ( offset < o );
								int64_t const asym = aa[diagBlockPair.first];
								int64_t const bsym = ba[diagBlockPair.second];

								int64_t const diagscore = AS[offset] + score_type::getScore(asym,bsym);

								if ( valid )
								{
									if ( diagscore > score )
									{
										score = diagscore;
										trace = trace_m;
									}
								}
								else
								{
									score = diagscore;
									trace = trace_m;
									valid = true;
								}
							}

							assert ( valid );
						}

						AS.push(o,score);
						AT.push(ot,trace);
					}

				std::pair<int64_t,int64_t> P = blockToAntiDiag(std::pair<int64_t,int64_t>(la,lb));

				uint64_t ATOo = 0;
				while ( P != std::pair<int64_t,int64_t>(0,0) )
				{
					uint64_t const offset = antiDiagOffset(P.first,P.second);
					std::pair<int64_t,int64_t> B(antiDiagToBlock(P));
					trace_step const trace = static_cast<trace_step>(AT[offset]);

					ATO.push(ATOo,trace);

					switch ( trace )
					{
						case trace_m:
						{
							assert ( B.first );
							assert ( B.second );

							B.first -= 1;
							B.second -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						case trace_d:
						{
							assert ( B.first );

							B.first -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						case trace_i:
						{
							assert ( B.second );

							B.second -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						default:
						{
							assert ( false );
						}
					}
				}

				std::reverse(ATO.begin(),ATO.begin()+ATOo);

				return ATOo;
			}

			template<
				typename iterator_a,
				typename iterator_b
			>
			uint64_t alignPrefix(iterator_a aa, iterator_a ae, iterator_b ba, iterator_b be)
			{
				int64_t const za = ae-aa;
				int64_t const lb = be-ba;

				// lb*max_score + i*deletionscore < 0
				// lb*max_score < i * -deletionscore
				// i > lb*max_score/-deletionscore
				int64_t const maxscore = score_type::getMaxScore();
				int64_t const delpen = -getDeletionScore();
				int64_t const maxdel = (lb * maxscore + delpen - 1) / delpen + 1;
				int64_t const maxusea = maxdel + lb;
				int64_t const la = std::min(za,maxusea);

				uint64_t o = 0;
				uint64_t ot = 0;
				AS.push(o,0);
				AT.push(ot,trace_x);

				bool avalid = false;
				int64_t amax = 0;
				int64_t amaxscore = 0;

				for ( int64_t antidiag = 1; antidiag <= la+lb; ++antidiag )
					for ( int64_t diag = 0; diag <= antidiag; ++ diag )
					{
						std::pair<int64_t,int64_t> const antiDiagPair(antidiag,diag);
						std::pair<int64_t,int64_t> const blockPair = antiDiagToBlock(antiDiagPair);

						assert ( blockPair.first >= 0 );
						assert ( blockPair.second >= 0 );

						int64_t score = 0;
						unsigned char trace = trace_x;

						if ( blockPair.first <= la && blockPair.second <= lb )
						{
							bool valid = false;

							if ( blockPair.first )
							{
								std::pair<int64_t,int64_t> const topBlockPair(blockPair.first-1,blockPair.second);
								std::pair<int64_t,int64_t> const topAntiDiagPair = blockToAntiDiag(topBlockPair);
								assert ( topAntiDiagPair.first < antiDiagPair.first );
								assert ( topAntiDiagPair.first >= 0 );
								assert ( topAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(topAntiDiagPair.first,topAntiDiagPair.second);
								assert ( offset < o );
								int64_t const topscore = AS[offset] + getDeletionScore();

								if ( valid )
								{
									if ( topscore > score )
									{
										score = topscore;
										trace = trace_d;
									}
								}
								else
								{
									score = topscore;
									trace = trace_d;
									valid = true;
								}
							}
							if ( blockPair.second )
							{
								std::pair<int64_t,int64_t> const leftBlockPair(blockPair.first,blockPair.second-1);
								std::pair<int64_t,int64_t> const leftAntiDiagPair = blockToAntiDiag(leftBlockPair);
								assert ( leftAntiDiagPair.first < antiDiagPair.first );
								assert ( leftAntiDiagPair.first >= 0 );
								assert ( leftAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(leftAntiDiagPair.first,leftAntiDiagPair.second);
								assert ( offset < o );
								int64_t const leftscore = AS[offset] + getInsertionScore();

								if ( valid )
								{
									if ( leftscore > score )
									{
										score = leftscore;
										trace = trace_i;
									}
								}
								else
								{
									score = leftscore;
									trace = trace_i;
									valid = true;
								}
							}
							if ( blockPair.first && blockPair.second )
							{
								std::pair<int64_t,int64_t> const diagBlockPair(blockPair.first-1,blockPair.second-1);
								std::pair<int64_t,int64_t> const diagAntiDiagPair = blockToAntiDiag(diagBlockPair);
								assert ( diagAntiDiagPair.first < antiDiagPair.first );
								assert ( diagAntiDiagPair.first >= 0 );
								assert ( diagAntiDiagPair.second >= 0 );
								uint64_t const offset = antiDiagOffset(diagAntiDiagPair.first,diagAntiDiagPair.second);
								assert ( offset < o );
								int64_t const asym = aa[diagBlockPair.first];
								int64_t const bsym = ba[diagBlockPair.second];

								int64_t const diagscore = AS[offset] + score_type::getScore(asym,bsym);

								if ( valid )
								{
									if ( diagscore > score )
									{
										score = diagscore;
										trace = trace_m;
									}
								}
								else
								{
									score = diagscore;
									trace = trace_m;
									valid = true;
								}
							}

							assert ( valid );

							if ( blockPair.second == lb )
							{
								if ( ! avalid )
								{
									avalid = true;
									amax = blockPair.first;
									amaxscore = score;
								}
								else if ( score > amaxscore )
								{
									amax = blockPair.first;
									amaxscore = score;
								}
							}
						}


						AS.push(o,score);
						AT.push(ot,trace);
					}

				assert ( avalid );

				std::pair<int64_t,int64_t> P = blockToAntiDiag(std::pair<int64_t,int64_t>(amax,lb));

				uint64_t ATOo = 0;
				while ( P != std::pair<int64_t,int64_t>(0,0) )
				{
					uint64_t const offset = antiDiagOffset(P.first,P.second);

					assert ( offset < o );
					assert ( offset < ot );

					std::pair<int64_t,int64_t> B(antiDiagToBlock(P));
					trace_step const trace = static_cast<trace_step>(AT[offset]);

					assert ( trace != trace_x );

					ATO.push(ATOo,trace);

					switch ( trace )
					{
						case trace_m:
						{
							assert ( B.first );
							assert ( B.second );

							B.first -= 1;
							B.second -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						case trace_d:
						{
							assert ( B.first );

							B.first -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						case trace_i:
						{
							assert ( B.second );

							B.second -= 1;

							P = blockToAntiDiag(B);

							break;
						}
						default:
						{
							assert ( false );
						}
					}
				}

				std::reverse(ATO.begin(),ATO.begin()+ATOo);

				return ATOo;
			}

			struct AlignResult
			{
				uint64_t startpos_a;
				uint64_t endpos_a;
				uint64_t startpos_b;
				uint64_t endpos_b;
				uint64_t al_O;
				int64_t score;

				AlignResult() : startpos_a(0), endpos_a(0), startpos_b(0), endpos_b(0), al_O(0), score(std::numeric_limits<int64_t>::min()) {}
				AlignResult(
					uint64_t rstartpos_a,
					uint64_t rendpos_a,
					uint64_t rstartpos_b,
					uint64_t rendpos_b,
					uint64_t ral_O,
					int64_t rscore
				) : startpos_a(rstartpos_a), endpos_a(rendpos_a), startpos_b(rstartpos_b), endpos_b(rendpos_b), al_O(ral_O), score(rscore)
				{}
			};

			template<
				typename it_a,
				typename it_b
			>
			AlignResult align(
				it_a const aa,
				it_a const ae,
				uint64_t const seedposa,
				it_b const bb,
				it_b const be,
				uint64_t const seedposb
			)
			{
				uint64_t const al_O_0 = alignPrefix(aa+seedposa,ae,bb+seedposb,be);

				std::pair<uint64_t,uint64_t> const SL_0 = getStringLengthUsed(al_O_0);

				assert ( be - (bb+seedposb) == static_cast<int64_t>(SL_0.second) );

				uint64_t const endposa = seedposa + SL_0.first;
				uint64_t const endposb = seedposb + SL_0.second;

				std::reverse_iterator<it_a> raa(aa + endposa);
				std::reverse_iterator<it_a> rae(aa);
				std::reverse_iterator<it_b> rba(bb + endposb);
				std::reverse_iterator<it_b> rbe(bb);

				uint64_t const al_O_1 = alignPrefix(raa,rae,rba,rbe);

				std::pair<uint64_t,uint64_t> const SL_1 = getStringLengthUsed(al_O_1);

				assert ( static_cast<int64_t>(SL_1.second) == be-bb );

				uint64_t const startposa = endposa - SL_1.first;
				uint64_t const startposb = endposb - SL_1.second;

				uint64_t const al_O_2 = (*this)(aa + startposa,aa + endposa,bb + startposb,bb + endposb);

				int64_t const score = getScore(al_O_2,aa+startposa,bb+startposb);

				AlignResult const R(startposa,endposa,startposb,endposb,al_O_2,score);

				return R;
			}

		};

		typedef Align< ::libmaus2::protein::Blosum62> AlignBlosum62;
	}
}
#endif
