#pragma once

#include "pipeline/stage.hpp"
#include "assembly_graph/graph_support/graph_processing_algorithm.hpp"
#include "assembly_graph/graph_support/basic_edge_conditions.hpp"
#include "modules/simplification/tip_clipper.hpp"
#include "projects/mts/contig_abundance.hpp"
#include "io/reads/osequencestream.hpp"

#include "llvm/Support/YAMLParser.h"
#include "llvm/Support/YAMLTraits.h"

namespace debruijn_graph {

struct SeriesAnalysisConfig {
    uint k;
    uint sample_cnt;
    uint frag_size;
    uint min_len;

    std::string kmer_mult, bin, bin_prof, edges_sqn, edges_mpl, edge_fragments_mpl;
};

}

namespace llvm { namespace yaml {

template<> struct MappingTraits<debruijn_graph::SeriesAnalysisConfig> {
    static void mapping(IO& io, debruijn_graph::SeriesAnalysisConfig& cfg) {
        io.mapRequired("k", cfg.k);
        io.mapRequired("sample_cnt", cfg.sample_cnt);
        io.mapRequired("kmer_mult", cfg.kmer_mult);
        io.mapRequired("bin", cfg.bin);
        io.mapRequired("bin_prof", cfg.bin_prof);
        io.mapRequired("min_len", cfg.min_len);
        io.mapRequired("edges_sqn", cfg.edges_sqn);
        io.mapRequired("edges_mpl", cfg.edges_mpl);
        io.mapRequired("edge_fragments_mpl", cfg.edge_fragments_mpl);
        io.mapRequired("frag_size", cfg.frag_size);
    }
};

} }

namespace debruijn_graph {

template<class graph_pack>
shared_ptr<visualization::graph_colorer::GraphColorer<typename graph_pack::graph_t>> DefaultGPColorer(
    const graph_pack& gp) {
    io::SingleRead genome("ref", gp.genome.str());
    auto mapper = MapperInstance(gp);
    auto path1 = mapper->MapRead(genome).path();
    auto path2 = mapper->MapRead(!genome).path();
    return visualization::graph_colorer::DefaultColorer(gp.g, path1, path2);
}

inline double l2_norm(const AbundanceVector& v) {
    double s = 0.;
    for (auto val : v) {
        s += val * val;
    }
    return std::sqrt(s);
}

inline double cosine_sim(const AbundanceVector& v1, const AbundanceVector& v2) {
    double s = 0.;
    for (size_t i = 0; i < v1.size(); ++i) {
        s += v1[i] * v2[i];
    }
    return s / (l2_norm(v1) * l2_norm(v2));
}

template<class Graph>
class EdgeAbundance: public omnigraph::GraphActionHandler<Graph> {
    typedef map<EdgeId, AbundanceVector> Storage;
    typedef Storage::const_iterator const_iterator;
    Storage edge_abundance_;
    const ContigAbundanceCounter& abundance_counter_;

public:
    EdgeAbundance(const Graph& g, const ContigAbundanceCounter& abundance_counter) :
        omnigraph::GraphActionHandler<Graph>(g, "EdgeAbundance"),
        abundance_counter_(abundance_counter){}

    void Fill() {
        for (auto it = this->g().ConstEdgeBegin(true); !it.IsEnd(); ++it) {
            HandleAdd(*it);
        }
    }

    virtual void HandleAdd(EdgeId e) override {
        auto ab = abundance_counter_(this->g().EdgeNucls(e).str());
        if (!ab) {
            INFO("Couldn't estimate abundance of edge " << this->g().str(e));
        } else {
            edge_abundance_[e] = *ab;
        }
    }

    const_iterator begin() const {
        return edge_abundance_.begin();
    }

    const_iterator end() const {
        return edge_abundance_.end();
    }

    const_iterator find(EdgeId e) const {
        return edge_abundance_.find(e);
    }

    size_t count(EdgeId e) const {
        return edge_abundance_.count(e);
    }

private:
    DECL_LOGGER("EdgeAbundance");
};

template<class Graph>
class AggressiveClearing: public omnigraph::EdgeProcessingAlgorithm<Graph> {
    typedef typename Graph::EdgeId EdgeId;
    const EdgeAbundance<Graph>& edge_abundance_;
    const AbundanceVector base_profile_;
    const double similarity_threshold_;
    const double norm_ratio_threshold_;
    EdgeRemover<Graph> edge_remover_;
    func::TypedPredicate<EdgeId> topological_condition_;

protected:
    virtual bool ProcessEdge(EdgeId e) override {
        DEBUG("Processing edge " << this->g().str(e));
        if (!topological_condition_(e)) {
            DEBUG("Topological condition failed");
            return false;
        }
        auto it = edge_abundance_.find(e);
        if (it == edge_abundance_.end()) {
            DEBUG("Edge " << this->g().str(e) << " did not have valid abundance profile");
            return false;
        }
        const auto& profile = it->second;
        DEBUG("Edge profile " << PrintVector(profile));
        double sim = cosine_sim(profile, base_profile_);
        double norm_ratio = l2_norm(profile) / l2_norm(base_profile_);

        DEBUG("Similarity between edge and base profiles " << sim);
        DEBUG("Norm ratio " << norm_ratio);
        if (math::ls(norm_ratio, norm_ratio_threshold_)
                || math::ls(sim, similarity_threshold_)) {
            DEBUG("Removing edge " << this->g().str(e));

            edge_remover_.DeleteEdge(e);
            return true;
        }
        return false;
    }

public:
    AggressiveClearing(Graph &g,
                       const EdgeAbundance<Graph>& edge_abundance,
                       const AbundanceVector& base_profile,
                       double similarity_threshold,
                       double norm_ratio_threshold,
                       const std::function<void(EdgeId)> &removal_handler = 0) :
        EdgeProcessingAlgorithm<Graph>(g, true),
        edge_abundance_(edge_abundance),
        base_profile_(base_profile),
        similarity_threshold_(similarity_threshold),
        norm_ratio_threshold_(norm_ratio_threshold),
        edge_remover_(g, removal_handler),
        topological_condition_(func::Or(AlternativesPresenceCondition<Graph>(g), TipCondition<Graph>(g))) {
            DEBUG("Base profile " << PrintVector(base_profile_));
        }
private:
    DECL_LOGGER("AggressiveClearing");
};

class SeriesAnalysis : public spades::AssemblyStage {

    boost::optional<AbundanceVector> InferAbundance(const std::string& bin_mult_fn,
                                                    const std::string& b_id) const {
        path::CheckFileExistenceFATAL(bin_mult_fn);

        ifstream is(bin_mult_fn);
        vector<AbundanceVector> abundances;
        while (true) {
            string name;
            is >> name;
            if (!is.fail()) {
                AbundanceVector vec(SampleCount(), 0.0);
                for (size_t i = 0; i < SampleCount(); ++i) {
                    is >> vec[i];
                    VERIFY(!is.fail());
                }
                if (name == b_id) {
                    abundances.push_back(vec);
                }
            } else {
                INFO("Read " << abundances.size() << " profiles for bin " << b_id);
                break;
            }
        }
        return boost::optional<AbundanceVector>(MeanVector(abundances));
    }

    void PrintEdgeFragmentProfiles(const conj_graph_pack &gp, const ContigAbundanceCounter &abundance_counter, 
                                   size_t split_length, size_t min_len, std::ostream &os) const {
        for (auto it = gp.g.ConstEdgeBegin(true); !it.IsEnd(); ++it) {
            EdgeId e = *it;
            io::SingleRead full_contig(ToString(gp.g.int_id(e)), gp.g.EdgeNucls(e).str());
            for (size_t i = 0; i < full_contig.size(); i += split_length) {
                if (full_contig.size() - i < min_len) {
                    DEBUG("Fragment shorter than min_length_bound " << min_len);
                    break;
                }

                io::SingleRead contig = full_contig.Substr(i, std::min(i + split_length, full_contig.size()));

                DEBUG("Processing fragment # " << (i / split_length) << " with id " << contig.name());

                auto abundance_vec = abundance_counter(contig.GetSequenceString(), contig.name());

                if (abundance_vec) {
                    size_t len = contig.GetSequenceString().size();
                    os << contig.name() << " " << len << " " << PrintVector(*abundance_vec) << std::endl;
                    //copy(abundance_vec->begin(), abundance_vec->begin() + config.sample_cnt,
                    //     ostream_iterator<Mpl>(ss, " "));
                    DEBUG("Successfully estimated abundance of " << contig.name());
                } else {
                    DEBUG("Failed to estimate abundance of " << contig.name());
                }
            }
        }
    }

public:
    SeriesAnalysis() : AssemblyStage("Series Analysis", "series_analysis") { }

    void load(conj_graph_pack &, const std::string &, const char *) { }

    void save(const conj_graph_pack &, const std::string &, const char *) const { }

    void run(conj_graph_pack &gp, const char *) {
        std::string cfg = cfg::get().series_analysis;
        INFO("Series analysis enabled with config " << cfg);

        auto Buf = llvm::MemoryBuffer::getFile(cfg);
        VERIFY_MSG(Buf, "Failed to load config file " + cfg);

        llvm::yaml::Input yin(*Buf.get());
        SeriesAnalysisConfig config;
        yin >> config;

        SetSampleCount(config.sample_cnt);

        ContigAbundanceCounter abundance_counter(config.k, 
                                                 SingleClusterAnalyzer(2., 0.4),
                                                 cfg::get().tmp_dir);

        DEBUG("Initiating abundance counter");
        abundance_counter.Init(config.kmer_mult);
        DEBUG("Abundance counter ready");

        if (!config.edges_sqn.empty()) {
            io::osequencestream oss(config.edges_sqn);
            for (auto it = gp.g.ConstEdgeBegin(true); !it.IsEnd(); ++it) {
                EdgeId e = *it;
                string s = gp.g.EdgeNucls(e).str();
                oss << io::SingleRead(io::MakeContigId(gp.g.int_id(e), s.size()), s);
            }
        }

        if (!config.edges_mpl.empty()) {
            ofstream os(config.edges_mpl);
            PrintEdgeFragmentProfiles(gp, abundance_counter, -1ul, config.min_len, os);
        }

        if (!config.edge_fragments_mpl.empty()) {
            ofstream os(config.edge_fragments_mpl);
            PrintEdgeFragmentProfiles(gp, abundance_counter, config.frag_size, config.min_len, os);
        }

        boost::optional<AbundanceVector> bin_profile = InferAbundance(config.bin_prof, config.bin);
        if (!bin_profile) {
            ERROR("Couldn't estimate profile of bin");
            return;
        }

        EdgeAbundance<Graph> edge_abundance(gp.g, abundance_counter);
        edge_abundance.Fill();

        gp.EnsureBasicMapping();
        gp.FillQuality();
        visualization::graph_labeler::DefaultLabeler<Graph> labeler(gp.g, gp.edge_pos);
        auto colorer = DefaultGPColorer(gp);
        path::make_dir(cfg::get().output_dir + "pictures/");
        QualityEdgeLocalityPrintingRH<Graph> qual_removal_handler(gp.g, gp.edge_qual, labeler, colorer,
                                       cfg::get().output_dir + "pictures/");

        INFO("Launching aggressive graph clearing");
        //positive quality edges removed (folder colored_edges_deleted)
        AggressiveClearing<Graph> clearing(gp.g, edge_abundance,
                                            *bin_profile, 0.8, 0.3, [&](EdgeId e) {
                        qual_removal_handler.HandleDelete(e);});
        clearing.Run();
        INFO("Graph clearing finished");

        INFO("Drawing edges with failed abundance estimate")
        path::make_dir(cfg::get().output_dir + "pictures_no_ab/");
        QualityEdgeLocalityPrintingRH<Graph> qual_removal_handler2(gp.g, gp.edge_qual, labeler, colorer,
                                       cfg::get().output_dir + "pictures_no_ab/");

        for (auto it = gp.g.ConstEdgeBegin(true); !it.IsEnd(); ++it) {
            EdgeId e = *it;
            if (edge_abundance.count(e) == 0) {
                qual_removal_handler2.HandleDelete(e);
            }
        }
    }
};

}
