//***************************************************************************
//* Copyright (c) 2015 Saint Petersburg State University
//* Copyright (c) 2011-2014 Saint Petersburg Academic University
//* All Rights Reserved
//* See file LICENSE for details.
//***************************************************************************

#pragma once

#include "abstract_contig_corrector.hpp"

using namespace debruijn_graph;

namespace dipspades {

class CloseGapsCorrector : public AbstractContigCorrector{

	set<size_t> incorr_contigs;
	size_t num_corr;

	size_t connected_length_;
	size_t disconnected_length_;

	vector<EdgeId> ClosePathGap(vector<EdgeId> path, vector<size_t> gap_index){
		vector<EdgeId> new_path;
		size_t current_gap = 0;
		for(size_t i = 0; i < path.size() - 1; i++){
			EdgeId cur_edge = path[i];
			new_path.push_back(cur_edge);
			if(i == gap_index[current_gap]){
				VertexId start = g_.EdgeEnd(cur_edge);
				VertexId end = g_.EdgeStart(path[i + 1]);
				auto dijkstra = DijkstraHelper<Graph>::CreateTargeredBoundedDijkstra(g_,
						end, dsp_cfg::get().pbr.max_bulge_nucls_len); //DijkstraHelper<Graph>::CreateBoundedDijkstra(g_, dsp_cfg::get().pbr.max_bulge_nucls_len);
				dijkstra.Run(start);
				if(dijkstra.DistanceCounted(end)){
					vector<EdgeId> add_path = dijkstra.GetShortestPathTo(end);
					for(auto e = add_path.begin(); e != add_path.end(); e++)
						if(g_.EdgeStart(*e) != g_.EdgeEnd(*e))
							new_path.push_back(*e);
				}
				else{
					// second attempt
					VertexId prev_start = g_.EdgeStart(cur_edge);
					dijkstra.Run(prev_start);
	               if(dijkstra.DistanceCounted(end)){
	                    vector<EdgeId> add_path = dijkstra.GetShortestPathTo(end);
						new_path.erase(new_path.begin() + new_path.size() - 1);
                	    for(auto e = add_path.begin(); e != add_path.end(); e++)
                	    	if(g_.EdgeStart(*e) != g_.EdgeEnd(*e))
                	    		new_path.push_back(*e);
					}
				}
				current_gap++;
			}
		}
		new_path.push_back(path[path.size() - 1]);
		return new_path;
	}

	size_t CountContigsWithGaps(ContigStoragePtr storage) {
		size_t contigs_with_gaps = 0;
		for(size_t i = 0; i < storage->Size(); i++)
			if(!IsPathConnected(g_, (*storage)[i]->path_seq()))
				contigs_with_gaps++;
		return contigs_with_gaps;
	}

	void ProcessContigs(ContigStoragePtr storage) {
		double processed_perc = 0.1;
		double step = 0.1;
		for(size_t i = 0; i < storage->Size(); i++) {
			storage->ReplaceContig(Correct((*storage)[i]), i);
			double cur_process_perc = static_cast<double>(i) / static_cast<double>(storage->Size());
			if(cur_process_perc > processed_perc) {
				while(processed_perc + step <= cur_process_perc)
					processed_perc += step;
				INFO(ToString(processed_perc * 100.0) << "% contigs were processed");
				processed_perc += step;
			}
		}
		INFO("100% contigs were processed");
	}

public:
	CloseGapsCorrector(Graph &g) :
		AbstractContigCorrector(g),
		num_corr(0),
		connected_length_(0),
		disconnected_length_(0)	{ }

	virtual ContigStoragePtr Correct(ContigStoragePtr storage){

		INFO(ToString(CountContigsWithGaps(storage)) << " contigs from " <<
				ToString(storage->Size()) << " have gaps before correction");

		ProcessContigs(storage);

		INFO(ToString(num_corr) << " contigs from " <<
				ToString(storage->Size()) << " with total length " << ToString(connected_length_) + " are correct");
		INFO(ToString(storage->Size() - num_corr) << " contigs from "
				<< ToString(storage->Size()) << " with total length " <<
				ToString(disconnected_length_) + " have gaps after correction");

		storage->DeleteByIDs(incorr_contigs);
		return storage;
	}

	virtual MappingContigPtr Correct(MappingContigPtr contig){
		vector<EdgeId> path = contig->path_seq();
		if(path.size() <= 1){
			num_corr++;
			return contig;
		}
		vector<size_t> gap_indexes;
		for(size_t i = 0; i < path.size() - 1; i++){
			EdgeId e1 = path[i];
			EdgeId e2 = path[i + 1];
			if(!AreEdgesConnected(g_, e1, e2)){
				gap_indexes.push_back(i);
			}
		}

		TRACE("Contig " << contig->id() << " has " << gap_indexes.size() << " gaps");

		// contig is connected
		if(gap_indexes.size() == 0) {
			num_corr++;
			connected_length_ += GetPathLength(g_, contig->path_seq());
			return contig;
		}

		TRACE("Contig path before correction: " << SimplePathWithVerticesToString(g_, contig->path_seq()));

		vector<EdgeId> new_path = ClosePathGap(path, gap_indexes);
		if(IsPathConnected(g_, new_path)) {
			TRACE("Gaps were closed");
			TRACE("Contig path after correction: " << SimplePathWithVerticesToString(g_, new_path));
			num_corr++;
			connected_length_ += GetPathLength(g_, new_path);
			return MappingContigPtr(new ReplacedPathMappingContig(contig, new_path));
		}

		TRACE("Contig " << contig->id() << " remains incorrected!");
		incorr_contigs.insert(contig->id());
		disconnected_length_ += GetPathLength(g_, contig->path_seq());
		return contig;
	}

private:
	DECL_LOGGER("CloseGapsCorrector")
};

}
