/****
DIAMOND protein sequence aligner
Copyright (C) 2012-2026 Benjamin J. Buchfink

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/>.
****/
// SPDX-License-Identifier: GPL-3.0-or-later

#include <algorithm>
#include <utility>
#include <math.h>
#include <memory>
#include <numeric>
#include "extend.h"
#include "data/block/block.h"
#include "data/sequence_set.h"
#include "data/string_set.h"
#include "util/options/option.h"
#include "basic/config.h"
#include "target.h"
#include "util/log_stream.h"
#include "util/util.h"
#include "global_ranking/global_ranking.h"
#include "search/hit.h"
#include "load_hits.h"
#include "def.h"
#include "util/memory/mem_profile.h"

using std::accumulate;
using std::vector;
using std::list;
using std::array;
using std::pair;
using std::endl;
using std::make_move_iterator;
using std::any_of;
using std::numeric_limits;
using std::runtime_error;
using std::tie;
using std::min;

const SEMap<Extension::Mode> EnumTraits<Extension::Mode>::from_string = {
	{ "banded-fast", Extension::Mode::BANDED_FAST},
	{ "banded-slow", Extension::Mode::BANDED_SLOW},
	{ "full", Extension::Mode::FULL},
	{ "none", Extension::Mode::NONE},
	{ "global", Extension::Mode::GLOBAL}
};

namespace Extension {

Query::Query(BlockId block_id, Statistics& stats, const Search::Config& cfg, std::pmr::memory_resource& pool) :
	block_id(block_id),
	title(cfg.query->ids()[block_id]),
	source_length(align_mode.query_translated ? (int)cfg.query->source_seqs()[block_id].length() : (int)cfg.query->seqs()[block_id].length()),
	true_aa_length(0),
	self_alignment_score(cfg.query->has_self_aln() ? cfg.query->self_aln_score(block_id) : 0.0)
{
	MEM_SCOPE("extend/query");
	const unsigned contexts = align_mode.query_contexts;
	sequence.reserve(contexts);
	for (unsigned i = 0; i < contexts; ++i)
		sequence.push_back(cfg.query->seqs()[block_id * contexts + i]);
	true_aa_length = ::Stats::count_true_aa(sequence.front());

	const Loc qlen = sequence.front().length();
	if (qlen <= config.short_query_max_len)
		ungapped_cutoff = score_matrix.rawscore(config.short_query_ungapped_bitscore);
	else if (cfg.ungapped_evalue > 0.0)
		ungapped_cutoff = score_matrix.rawscore(score_matrix.bitscore_norm(cfg.ungapped_evalue, qlen));
	else
		ungapped_cutoff = 1;
	ungapped_cutoff = std::max(ungapped_cutoff, 1);

	const unsigned cbs = static_cast<unsigned>(config.comp_based_stats_.get(Stats::DEFAULT_CBS));
	if (::Stats::CBS::hauser(cbs)) {
		hauser_correction.reserve(contexts);
		for (const Sequence& seq : sequence)
			hauser_correction.emplace_back(seq);
	}
	if (::Stats::CBS::matrix_adjust(cbs))
		composition = ::Stats::composition(sequence.front());
	if ((config.reseek_diags || config.lin_stage1_query || config.lin_stage1_target) && (cfg.extension_mode == Extension::Mode::BANDED_FAST || cfg.extension_mode == Extension::Mode::BANDED_SLOW)) {
		if (sequence.size() > 1)
			throw runtime_error("Reseeking diagonals is not supported for translated queries.");
		MEM_SCOPE("extend/query-seqindex");
		seqindex.reset(new Seqindex(sequence.front(), cfg.keyword_length, cfg.keyword_threshold, pool, stats));
	}
}

const std::map<Sensitivity, Mode> default_ext_mode = {
	{ Sensitivity::FASTER, Mode::BANDED_FAST},
	{ Sensitivity::FAST, Mode::BANDED_FAST},
	{ Sensitivity::SHAPES6x10, Mode::BANDED_FAST},
	{ Sensitivity::SHAPES30x10, Mode::BANDED_FAST},
	{ Sensitivity::LINCLUST_20, Mode::BANDED_FAST},
	{ Sensitivity::LINCLUST_40, Mode::BANDED_FAST},
	{ Sensitivity::DEFAULT, Mode::BANDED_FAST},
	{ Sensitivity::MID_SENSITIVE, Mode::BANDED_FAST},
	{ Sensitivity::SENSITIVE, Mode::BANDED_FAST},
	{ Sensitivity::MORE_SENSITIVE, Mode::BANDED_SLOW},
	{ Sensitivity::VERY_SENSITIVE, Mode::BANDED_SLOW},
	{ Sensitivity::ULTRA_SENSITIVE, Mode::BANDED_SLOW}
};

constexpr int64_t MAX_CHUNK_SIZE = 400, MIN_CHUNK_SIZE = 128, MAPANY_CHUNK_SIZE = 16;

static int64_t ranking_chunk_size(int64_t target_count, const int64_t ref_letters, const int64_t max_target_seqs) {
	if (config.no_ranking || config.global_ranking_targets > 0)
		return target_count;
	if (config.ext_chunk_size > 0)
		return config.ext_chunk_size;
	if (config.mapany)
		return MAPANY_CHUNK_SIZE;
	const double default_letters = config.sensitivity >= Sensitivity::VERY_SENSITIVE ? 800 * 1e6 : 2 * 1e9;
	const int64_t block_mult = std::max(int64_t(std::round((double)ref_letters / default_letters)), (int64_t)1);
	if (config.toppercent.present())
		return MIN_CHUNK_SIZE * block_mult;
	const int64_t size = std::max(MIN_CHUNK_SIZE, std::min(make_multiple(max_target_seqs, (int64_t)32), MAX_CHUNK_SIZE)) * block_mult;
	return config.target_hard_cap ? std::min(size, config.target_hard_cap) : size;
}

static bool have_filters(const Search::Config& cfg) {
	return config.min_id > 0 || config.approx_min_id.get(0.0) > 0 || config.query_cover > 0 || config.subject_cover > 0 || config.query_or_target_cover > 0;
}

static HspValues first_round_hspv(const Search::Config& cfg) {
	HspValues first_round = HspValues::NONE;
	if (config.min_id > 0)
		first_round |= HspValues::IDENT | HspValues::LENGTH;
	if (config.query_cover > 0)
		first_round |= HspValues::QUERY_COORDS;
	if (config.subject_cover > 0)
		first_round |= HspValues::TARGET_COORDS;
	if (config.cluster_threshold.present())
		first_round |= cfg.output_format->hsp_values;
	return first_round;
}

static bool ranking_terminate(bool new_hits, int last_tail_score, int tail_score, int64_t targets_processed, int64_t targets_aligned) {
	if (config.target_hard_cap && targets_processed >= config.target_hard_cap)
		return true;
	if (config.mapany && config.toppercent.blank() && targets_aligned > 0)
		return true;
	return !new_hits && (last_tail_score == 0
		|| double(tail_score) / (double)last_tail_score <= config.ranking_score_drop_factor
		|| score_matrix.bitscore(tail_score) < config.ranking_cutoff_bitscore);
}

Match Match::self_match(BlockId query_id, Sequence query_seq) {
	Match m(query_id, query_seq, nullptr, 0, numeric_limits<Score>::max(), 0.0);
	m.hsp.emplace_back();
	m.hsp.back().evalue = 0.0;
	m.hsp.back().score = numeric_limits<Score>::max();
	m.hsp.back().bit_score = DBL_MAX;
	m.hsp.back().query_range = { 0,query_seq.length() };
	m.hsp.back().query_source_range = { 0,query_seq.length() };
	m.hsp.back().subject_range = { 0,query_seq.length() };
	return m;
}

static bool add_self_aln(const Search::Config& cfg) {
	return config.add_self_aln && ((config.self && cfg.current_ref_block == 0) || (!config.self && cfg.current_query_block == cfg.current_ref_block));
}

static Hsp seed_only_hsp(const SeedHit& hit, unsigned query_source_len) {
	Hsp hsp;
	hsp.seed_only = true;
	hsp.score = hit.score;
	hsp.evalue = 0.0;
	hsp.d_begin = hsp.d_end = hit.diag();
	hsp.set_begin(hit.i, hit.j, Frame(hit.frame), query_source_len);
	hsp.set_end(hit.i + 1, hit.j + 1, Frame(hit.frame), query_source_len);
	hsp.subject_source_range = hsp.subject_range;
	return hsp;
}

static vector<Match> seed_only_matches(const Query& query, const Search::Config& cfg, SeedHitList& l) {
	vector<Match> matches;
	matches.reserve(l.target_scores.size());
	for (const TargetScore& target_score : l.target_scores) {
		const BlockId target_block_id = l.target_block_ids[target_score.target];
		Match match(target_block_id, cfg.target->seqs()[target_block_id], nullptr, target_score.score, target_score.score, 0.0);
		for (auto i = l.seed_hits.begin(target_score.target); i < l.seed_hits.end(target_score.target); ++i)
			match.hsp.push_back(seed_only_hsp(*i, query.source_length));
		match.hsp.sort();
		if (config.max_hsps > 0 && match.hsp.size() > config.max_hsps)
			match.hsp.resize(config.max_hsps);
		matches.push_back(std::move(match));
	}
	culling(matches, cfg);
	return matches;
}

static size_t lazy_masking(std::vector<uint32_t>::const_iterator target_block_ids, vector<uint32_t>::const_iterator target_block_ids_end, Block& targets, const MaskingAlgo algo) {
	MEM_SCOPE("extend/lazy-masking");
	if (algo == MaskingAlgo::NONE)
		return 0;
	vector<Letter> seq;
	const Masking& masking = Masking::get();
	size_t n = 0;
	for (auto i = target_block_ids; i != target_block_ids_end; ++i)
		if (targets.fetch_seq_if_unmasked(*i, seq)) {
			masking(seq.data(), seq.size(), algo, *i);
			targets.write_masked_seq(*i, seq);
			++n;
		}
	return n;
}

static TargetList extend_chunk(const Query& query,
	FlatArray<SeedHit>::Iterator seed_hits,
	FlatArray<SeedHit>::Iterator seed_hits_end,
	vector<uint32_t>::const_iterator target_block_ids,
	const Search::Config& cfg,
	Statistics& stat,
	DP::Flags flags,
	const HspValues hsp_values,
	std::pmr::memory_resource& pool)
{
	static const Loc GAPPED_FILTER_MIN_QLEN = 85;
	const int64_t n = seed_hits_end - seed_hits;
	stat.inc(Statistics::TARGET_HITS2, n);
	TaskTimer timer(flag_any(flags, DP::Flags::PARALLEL) ? config.target_parallel_verbosity : UINT_MAX);

	if (cfg.lazy_masking && !config.global_ranking_targets)
		stat.inc(Statistics::MASKED_LAZY, lazy_masking(target_block_ids, target_block_ids + n, *cfg.target, cfg.target_masking));

	pair<FlatArray<SeedHit>, vector<uint32_t>> gf;
	if (cfg.gapped_filter_evalue > 0.0 && config.global_ranking_targets == 0 && (!align_mode.query_translated || query.sequence[0].length() >= GAPPED_FILTER_MIN_QLEN)) {
		timer.go("Computing gapped filter");
		gf = gapped_filter(query, seed_hits, seed_hits_end, target_block_ids, stat, flags, cfg);
		if (!flag_any(flags, DP::Flags::PARALLEL))
			stat.inc(Statistics::TIME_GAPPED_FILTER, timer.microseconds());
		seed_hits = gf.first.begin();
		seed_hits_end = gf.first.end();
		target_block_ids = gf.second.cbegin();
	}
	stat.inc(Statistics::TARGET_HITS3, seed_hits_end - seed_hits);

	timer.go("Computing chaining");
	vector<WorkTarget> targets = ungapped_stage(query, seed_hits, seed_hits_end, target_block_ids, flags, stat, *cfg.target, cfg.extension_mode, pool, cfg);
	if (!flag_any(flags, DP::Flags::PARALLEL))
		stat.inc(Statistics::TIME_CHAINING, timer.microseconds());

	auto ret = align(targets, query, flags, hsp_values, cfg.extension_mode, *cfg.thread_pool, cfg, stat, pool);
	return ret;
}

vector<Match> extend(
	const Query& query,
	const Search::Config& cfg,
	Statistics& stat,
	DP::Flags flags,
	SeedHitList& l,
	std::pmr::memory_resource& pool)
{
	const unsigned UNIFIED_TARGET_LEN = 50;
	if (config.log_query || (flag_any(flags, DP::Flags::PARALLEL) && !config.swipe_all))
		*log_stream << "Query=" << query.title << " Hits=" << l.seed_hits.data_size() << endl;

	const unsigned query_len = (unsigned)query.sequence.front().length();
	const size_t target_count = l.target_block_ids.size();
	if (cfg.extension_mode == Mode::NONE) { std::sort(l.target_scores.begin(), l.target_scores.end()); return seed_only_matches(query, cfg, l); }
	const int64_t chunk_size = ranking_chunk_size(target_count, cfg.target->seqs().letters(), cfg.max_target_seqs);
	vector<TargetScore>::const_iterator i0 = l.target_scores.cbegin(), i1 = i0 + std::min((ptrdiff_t)chunk_size, l.target_scores.cend() - i0);

	if (config.toppercent.blank() && config.min_bit_score == 0.0 && (i1 - i0) < cfg.max_target_seqs && (config.ext_chunk_size == 0 || config.lin_stage1_query))
#ifdef EVAL_TARGET
		while (i1 < l.target_scores.cend() && i1->evalue <= config.max_evalue && size_t(i1 - i0) < config.max_alignments) ++i1;
#else
		while (i1 < l.target_scores.cend() && score_matrix.evalue(i1->score, query_len, UNIFIED_TARGET_LEN) <= config.max_evalue) i1 += min((ptrdiff_t)16, l.target_scores.cend() - i1);
#endif
	
	const HspValues first_round_hspv = config.anchored_swipe ||
		(config.sensitivity <= Sensitivity::SHAPES30x10 && flag_only(cfg.output_format->hsp_values, HspValues::COORDS) && config.toppercent.blank() && (size_t)cfg.max_target_seqs >= l.target_scores.size())
		? HspValues::COORDS : HspValues::NONE;
	const bool first_round_culling = !have_filters(cfg) || config.toppercent.present();
	bool new_hits_ev = false;
	int tail_score = 0, previous_tail_score = 0;
	FlatArray<SeedHit> seed_hits_chunk;
	vector<uint32_t> target_block_ids_chunk;
	vector<Match> matches;

	do {

		TargetList aligned_targets(&pool);
		bool new_hits;

		do {
			const int64_t current_chunk_size = i1 - i0;
			const bool multi_chunk = current_chunk_size < (int64_t)l.target_scores.size();

			if (multi_chunk) {
				MEM_SCOPE("extend/chunk-buffers");
				target_block_ids_chunk.clear();
				seed_hits_chunk.clear();
				target_block_ids_chunk.reserve(i1 - i0);
				seed_hits_chunk.reserve(i1 - i0, accumulate(i0, i1, (int64_t)0, [&l](int64_t n, const TargetScore& s) { return l.seed_hits.count(s.target) + n; }));
				for (vector<TargetScore>::const_iterator j = i0; j < i1; ++j) {
					target_block_ids_chunk.push_back(l.target_block_ids[j->target]);
					seed_hits_chunk.push_back(l.seed_hits.begin(j->target), l.seed_hits.end(j->target));
				}
			}

			TargetList v = extend_chunk(
				query,
				multi_chunk ? seed_hits_chunk.begin() : l.seed_hits.begin(),
				multi_chunk ? seed_hits_chunk.end() : l.seed_hits.end(),
				multi_chunk ? target_block_ids_chunk.cbegin() : l.target_block_ids.cbegin(),
				cfg,
				stat,
				flags,
				first_round_hspv,
				pool);

			stat.inc(Statistics::TARGET_HITS4, v.size());
			new_hits = new_hits_ev = v.size() > 0;
			if (multi_chunk)
				new_hits = append_hits(aligned_targets, v.begin(), v.end(), first_round_culling, cfg);
			else
				aligned_targets = std::move(v);

			i0 = i1;
			i1 += std::min((ptrdiff_t)chunk_size, l.target_scores.cend() - i1);
			previous_tail_score = tail_score;
			if (new_hits)
				tail_score = (i1 - 1)->score;
		} while (i0 < l.target_scores.cend() && !ranking_terminate(new_hits, previous_tail_score, (i1 - 1)->score, i1 - l.target_scores.cbegin(), aligned_targets.size()));

		if (config.swipe_all)
			aligned_targets = full_db_align(query, flags, HspValues::NONE, stat, *cfg.target, pool);

		culling(aligned_targets, !first_round_culling, cfg);
		stat.inc(Statistics::TARGET_HITS5, aligned_targets.size());
		
		{
			MEM_SCOPE("extend/matches");
			vector<Match> round_matches = align(aligned_targets, matches.size(), query, flags, first_round_hspv, first_round_culling, stat, cfg, pool);
			matches.insert(matches.end(), make_move_iterator(round_matches.begin()), make_move_iterator(round_matches.end()));
		}
	} while (config.toppercent.blank() && (int64_t)matches.size() < config.max_target_seqs_.get(DEFAULT_MAX_TARGET_SEQS)
		&& i0 < l.target_scores.cend() && new_hits_ev && (!config.mapany || (config.mapany && matches.empty())));

	if (add_self_aln(cfg) && !any_of(matches.cbegin(), matches.cend(), [&query](const Match& m) { return m.target_block_id == query.block_id; }))
		matches.push_back(Match::self_match(query.block_id, query.sequence[0]));

	culling(matches, cfg);

	return matches;
}

vector<Match> extend(BlockId query_id, Search::Hit* begin, Search::Hit* end, const Search::Config &cfg, Statistics &stat, DP::Flags flags, std::pmr::memory_resource& pool) {
	const Query query(query_id, stat, cfg, pool);
	TaskTimer timer(flag_any(flags, DP::Flags::PARALLEL) ? config.target_parallel_verbosity : UINT_MAX);
	timer.go("Loading seed hits");
	SeedHitList l = load_hits(begin, end, cfg.target->seqs());
	stat.inc(Statistics::TARGET_HITS0, l.target_block_ids.size());
	stat.inc(Statistics::TIME_LOAD_HIT_TARGETS, timer.microseconds());
	timer.finish();

	const int64_t target_count = (int64_t)l.target_block_ids.size();
	if (target_count == 0 && !config.swipe_all) {
		if (cfg.extension_mode == Mode::NONE)
			return vector<Match>();
		if (add_self_aln(cfg)) {
			vector<Match> r;
			Match match = Match::self_match(query.block_id, query.sequence[0]);
			r.push_back(std::move(match));
			return r;
		}
		return vector<Match>();
	}
	const int64_t chunk_size = ranking_chunk_size(target_count, cfg.target->seqs().letters(), cfg.max_target_seqs);

	if (chunk_size < target_count || config.global_ranking_targets > 0) {
		timer.go("Sorting targets by score");
		std::sort(l.target_scores.begin(), l.target_scores.end());
		stat.inc(Statistics::TIME_SORT_TARGETS_BY_SCORE, timer.microseconds());
		timer.finish();
		if (config.global_ranking_targets > 0)
			return GlobalRanking::ranking_list(query_id, l.target_scores.begin(), l.target_scores.end(), l.target_block_ids.begin(), l.seed_hits, cfg);
	}

	vector<Match> r = extend(query, cfg, stat, flags, l, pool);
	return r;
}

}