#include "assembly_graph/core/graph.hpp"
#include "assembly_graph/core/coverage.hpp"

#include "utils/stl_utils.hpp"
#include "utils/logger/log_writers.hpp"

#include "pipeline/graph_pack.hpp"
#include "pipeline/config_struct.hpp"

#include "modules/alignment/sequence_mapper.hpp"
#include "modules/alignment/pacbio/g_aligner.hpp"

#include "io/reads/io_helper.hpp"
#include "edlib/edlib.h"

#include "graphio.hpp"

#include <gtest/gtest.h>


using namespace debruijn_graph;

TEST(GraphAligner, EdlibSHWFULLTest) {

    std::string target = "abcd";
    std::string query = "ac";
    edlib::EdlibAlignResult result = edlib::edlibAlign(query.c_str(), (int) query.size(), target.c_str(), (int) target.size(),
                                                       edlib::edlibNewAlignConfig(10, edlib::EDLIB_MODE_SHW_EXTENDED, edlib::EDLIB_TASK_DISTANCE,
                                                                                  NULL, 0));
    EXPECT_EQ(result.status, edlib::EDLIB_STATUS_OK);
    EXPECT_NE(result.status, -1);
    INFO("Edlib: Result status = " << result.status)
    if (result.status == edlib::EDLIB_STATUS_OK) {
        for (int i = 0; i < result.numLocations; ++ i) {
            if (result.endLocations[i] == 0) {
                EXPECT_EQ(result.endScores[i], 1);           
            }
            if (result.endLocations[i] == 1) {
                EXPECT_EQ(result.endScores[i], 1);           
            }
            if (result.endLocations[i] == 2) {
                EXPECT_EQ(result.endScores[i], 1);           
            }
            if (result.endLocations[i] == 3) {
                EXPECT_EQ(result.endScores[i], 2);           
            }
        }
    }
    edlib::edlibFreeAlignResult(result);
}


TEST(GraphAligner, EdlibSHWFULLTestZero  ) {

    std::string target = "abcd";
    std::string query = "";
    edlib::EdlibAlignResult result = edlib::edlibAlign(query.c_str(), (int) query.size(), target.c_str(), (int) target.size()
                                                   , edlib::edlibNewAlignConfig(10, edlib::EDLIB_MODE_SHW_EXTENDED, edlib::EDLIB_TASK_DISTANCE,
                                                                         NULL, 0));
    EXPECT_EQ(result.status, edlib::EDLIB_STATUS_OK);
    EXPECT_NE(result.status, -1);
    INFO("Edlib: Result status = " << result.status)
    if (result.status == edlib::EDLIB_STATUS_OK) {
        INFO("Edlib: Result ed = " << result.editDistance)
        for (int i = 0; i < result.numLocations; ++ i) {
            INFO("Edlib: Result loc = " << result.endScores[i])
            if (result.endLocations[i] == 0) {
                EXPECT_EQ(result.endScores[i], 1);           
            }
        }
    }
    edlib::edlibFreeAlignResult(result);
}

debruijn_graph::config::pacbio_processor InitializePacBioProcessor() {
    debruijn_graph::config::pacbio_processor pb;  
    pb.internal_length_cutoff = 200; //500
    pb.compression_cutoff = 0.6; // 0.6
    pb.path_limit_stretching = 1.3; //1.3
    pb.path_limit_pressing = 0.7;//0.7
    pb.max_path_in_dijkstra = 15000; //15000
    pb.max_vertex_in_dijkstra = 15000; //2000
//gap_closer
    pb.long_seq_limit = 400; //400
    pb.pacbio_min_gap_quantity = 2; //2
    pb.contigs_min_gap_quantity = 1; //1
    return pb;
}

TEST(GraphAligner, TrivialTest ) {
    size_t K = 55;
    Graph g(K);
    graphio::ScanBasicGraph("./src/test/debruijn/graph_fragments/ecoli_400k/distance_estimation", g);

    std::string s = "[encoded data, 141663 characters]"; 
    std::string idealMapping = "17572 (7854,15914) 1565 (15914,8812) 20042 (8812,14180) 20044 (14180,10154) 20385 (10154,7187) 20449 (7187,9851) 2142 (9851,14410) 4373 (14410,14180) 20044 (14180,10154) 18816 (10154,7187) 20449 (7187,9851) 2142 (9851,14410) 18066 (14410,15914) 1565 (15914,8812) 19391 (8812,8959) 19476 (8959,8766) 20469 (8766,15855) 20420 (15855,12379) 20100 (12379,15855) 20420 (15855,12379) 20465 (12379,15430) 19870 (15430,9757) 19876 (9757,15430) 19870 (15430,9757) 19876 (9757,15430) 19870 (15430,9757) 19876 (9757,15430) 19870 (15430,9757) 20126 (9757,8767) 19477 (8767,8958) 18366 (8958,16718) ";
    // Sequence readSeq(s);
    io::SingleRead r("read", s);
    debruijn_graph::config::pacbio_processor pb = InitializePacBioProcessor();
    alignment::BWAIndex::AlignmentMode mode = alignment::BWAIndex::AlignmentMode::PacBio;
    sensitive_aligner::GapClosingConfig gap_cfg;
    gap_cfg.find_shortest_path = true;
    gap_cfg.updates_limit = 10000000;
    gap_cfg.penalty_ratio = 0.1;
    gap_cfg.queue_limit = 1000000;
    gap_cfg.iteration_limit = 1000000;
    gap_cfg.max_ed_proportion = 3;
    gap_cfg.ed_lower_bound = 500;
    gap_cfg.ed_upper_bound = 2000;
    gap_cfg.max_gs_states = 12000000;

    sensitive_aligner::EndsClosingConfig ends_cfg;
    ends_cfg.find_shortest_path = true;
    ends_cfg.updates_limit = 10000000;
    ends_cfg.penalty_ratio = 0.1;
    ends_cfg.queue_limit = 1000000;
    ends_cfg.iteration_limit = 1000000;
    ends_cfg.max_ed_proportion = 3;
    ends_cfg.ed_lower_bound = 500;
    ends_cfg.ed_upper_bound = 2000;
    ends_cfg.max_restorable_length = 5000;

    sensitive_aligner::GAligner galigner(g, sensitive_aligner::GAlignerConfig(pb, mode, gap_cfg, ends_cfg));
    auto current_read_mapping = galigner.GetReadAlignment(r);
    const auto& aligned_edges = current_read_mapping.edge_paths;
    std::string pathStr = "";
    if (aligned_edges.size() > 0){
        for (const auto &path : aligned_edges){
            for (const auto &edgeid: path) {
                VertexId v1 = g.EdgeStart(edgeid);
                VertexId v2 = g.EdgeEnd(edgeid);
                pathStr += std::to_string(edgeid.int_id()) + " (" + std::to_string(v1.int_id()) + "," + std::to_string(v2.int_id()) + ") ";
            }
            EXPECT_EQ(idealMapping, pathStr);
        }
    }
    
}


TEST(GraphAligner, DijkstraOneEdgeTest ) {
    size_t K = 55;
    Graph g(K);
    graphio::ScanBasicGraph("./src/test/debruijn/graph_fragments/ecoli_400k/distance_estimation", g);

    std::string s = "CCCTGAAGGTACGTGCAAAGATAATCGATTCTTTTTCGATTGTCTGGCTGTATGCGTCAACGTGAAACCGGCACCGGAAGTGCGTGAATTCTGGGGCTGGTGGATGGAGCTGGAAGCGCAGGAATCCCGTTTTACCTACAGTTACCAGTTTGGTCTGTTCGATAAAGCAGGCGACTGGAAGAGTGTTCCGGTAAAAGACA";
    int ideal_score = 2;
    std::unordered_map<VertexId, size_t> vertex_pathlen;
    int path_maxlen = 100500;
    int eid_int = 17572;
    EdgeId eid;
    for (auto it = g.ConstEdgeBegin(); !it.IsEnd(); ++it) {
        if (g.int_id(*it) == eid_int){
            eid = *it;
            break;
        }
    }
    sensitive_aligner::GapClosingConfig gap_cfg;
    gap_cfg.find_shortest_path = true;
    gap_cfg.updates_limit = 10000000;
    gap_cfg.penalty_ratio = 200;
    gap_cfg.queue_limit = 1000000;
    gap_cfg.iteration_limit = 1000000;

    sensitive_aligner::DijkstraGapFiller gap_filler = sensitive_aligner::DijkstraGapFiller(g, gap_cfg, s, eid, eid, 0, (int) s.size(), path_maxlen, vertex_pathlen);
    gap_filler.CloseGap();
    int score = gap_filler.edit_distance();
    EXPECT_EQ(ideal_score, score);
}


TEST(GraphAligner, DijkstraGrowEndsTest ) {
    size_t K = 55;
    Graph g(K);
    graphio::ScanBasicGraph("./src/test/debruijn/graph_fragments/ecoli_400k/distance_estimation", g);

    std::string s = "CCCTGAAGGTACGTGCAAAGATAATCGATTCTTTTTCGATTGTCTGGCTGTATGCGTCAACGTGAAACCGGCACCGGAAGTGCGTGAATTCTGGGGCTGGTGGATGGAGCTGGAAGCGCAGGAATCCCGTTTTACCTACAGTTACCAGTTTGGTCTGTTCGATAAAGCAGGCGACTGGAAGAGTGTTCCGGTAAAAGACA";
    int ideal_score = 2;
    int path_maxlen = 100500;
    int eid_int = 17572;
    EdgeId eid;
    for (auto it = g.ConstEdgeBegin(); !it.IsEnd(); ++it) {
        if (g.int_id(*it) == eid_int){
            eid = *it;
            break;
        }
    }
    sensitive_aligner::EndsClosingConfig gap_cfg;
    gap_cfg.find_shortest_path = true;
    gap_cfg.updates_limit = 10000000;
    gap_cfg.penalty_ratio = 0.1;
    gap_cfg.queue_limit = 1000000;
    gap_cfg.iteration_limit = 1000000;
    sensitive_aligner::DijkstraEndsReconstructor ends_filler = sensitive_aligner::DijkstraEndsReconstructor(g, gap_cfg, s, eid, 0, path_maxlen);
    ends_filler.CloseGap();
    int score = ends_filler.edit_distance();
    EXPECT_EQ(ideal_score, score);
}
