//     HTQC - a high-throughput sequencing quality control toolkit
//
//     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/>.

#include "htqc/MicroAssembler.h"

#include "htio2/QualityUtil.h"
#include "htio2/Kmer.h"

#include <sstream>

#include "stdlib.h"

using namespace std;

namespace htqc
{

MicroAssembler::MicroAssembler(int kmer_size, int confident_quality, int aln_identity, int conflict_quality_diff)
: km_sz(kmer_size)
, conf_qual(confident_quality)
, aln_iden(aln_identity)
, qual_diff(conflict_quality_diff)
{
}

MicroAssembler::~MicroAssembler()
{
}

bool MicroAssembler::assemble(const std::string& seq1,
                              const std::vector<int16_t>& qual_list_1,
                              const std::string& seq2,
                              const std::vector<int16_t>& qual_list_2)
{
    result_seq.clear();
    result_qual.clear();
    result_desc.clear();

    size_t len_a = seq1.length();
    size_t len_b = seq2.length();

    // generate quality mask
    vector<bool> qual_mask_a;
    vector<bool> qual_mask_b;
    gen_qual_mask(qual_list_1, qual_mask_a);
    gen_qual_mask(qual_list_2, qual_mask_b);

    // encode sequence
    // generate kmer profile
    htio2::EncodedSeq enc_a;
    htio2::EncodedSeq enc_b;
    vector<KmerT> kmers_a;
    vector<KmerT> kmers_b;
    multimap<KmerT, PosT> kmers_pos_b;

    htio2::encode_nt5(seq1, enc_a);
    htio2::encode_nt5(seq2, enc_b);
    htio2::gen_kmer_nt5(enc_a, kmers_a, km_sz);
    htio2::gen_kmer_nt5(enc_b, kmers_b, km_sz);
    htio2::summ_kmer_pos(kmers_b, kmers_pos_b);

    //
    // perform alignment
    //
    const ssize_t nk_a = kmers_a.size();

    bool success = false;

    // traverse A position
    for (PosT i_kmer_a = 0; i_kmer_a < nk_a; i_kmer_a++)
    {
        if (!qual_mask_a[i_kmer_a])
        {
            continue;
        }

        const KmerT kmer_a = kmers_a[i_kmer_a];
        pair<multimap<KmerT,PosT>::const_iterator, multimap<KmerT,PosT>::const_iterator> kb_pos_bound = kmers_pos_b.equal_range(kmer_a);

        // traverse B position
        for (multimap<KmerT,PosT>::const_iterator it = kb_pos_bound.first; it != kb_pos_bound.second; it++)
        {
            const PosT i_kmer_b = it->second;
            if (!qual_mask_b[i_kmer_b]) continue;

            const size_t len_a_before = i_kmer_a;
            const size_t len_b_before = i_kmer_b;
            const size_t len_a_after = len_a - len_a_before;
            const size_t len_b_after = len_b - len_b_before;


            // skip if B lays on left of A
            if (len_a_before < len_b_before && len_b_after < len_a_after)
            {
                continue;
            }

            // ib = ia + offset
            const PosT offset = i_kmer_b - i_kmer_a;

            //
            // check sequence
            //
            PosT ia_from = 0;
            PosT ib_from = 0;
            if (len_a_before < len_b_before)
                ib_from = ia_from + offset;

            else
                ia_from = ib_from - offset;

            PosT ia_to = ia_from;
            PosT ib_to = ib_from;

            uint16_t iden = 0;
            bool conflict = false;

            string aligned_part;
            vector<int16_t> aligned_part_qual;

            while (1)
            {
                int16_t qual_a = qual_list_1[ia_to];
                int16_t qual_b = qual_list_2[ib_to];

                // check base
                if (enc_a[ia_to] == enc_b[ib_to])
                {
                    if (qual_a >= conf_qual || qual_b >= conf_qual) iden++;
                }
                else
                {
                    // check quality
                    int curr_qual_diff = abs(qual_a - qual_b);
                    if (curr_qual_diff < qual_diff)
                    {
                        conflict = true;
                        break;
                    }
                }

                // store base and quality
                if (qual_a >= qual_b)
                {
                    aligned_part.push_back(seq1[ia_to]);
                    aligned_part_qual.push_back(qual_list_1[ia_to]);
                }
                else
                {
                    aligned_part.push_back(seq2[ib_to]);
                    aligned_part_qual.push_back(qual_list_2[ib_to]);
                }

                ia_to++;
                ib_to++;
                if (ia_to == len_a || ib_to == len_b) break;
            }

            // record the assembled sequence
            if (!conflict && iden >= aln_iden)
            {
                success = true;

                ostringstream handle;
                handle << "kmer_size=" << km_sz
                    << " kmer_pos_a=" << i_kmer_a + 1
                    << " kmer_pos_b=" << i_kmer_b + 1
                    << " num_identical=" << iden;
                result_desc = handle.str();

                // content before alignment
                if (ia_from > 0)
                {
                    result_seq.assign(seq1, 0, ia_from);
                    result_qual.assign(qual_list_1.begin(), qual_list_1.begin() + ia_from);
                }
                else if (ib_from > 0)
                {
                    result_seq.assign(seq2, 0, ib_from);
                    result_qual.assign(qual_list_2.begin(), qual_list_2.begin() + ib_from);
                }

                // contents that aligned
                result_seq += aligned_part;
                result_qual.insert(result_qual.end(), aligned_part_qual.begin(), aligned_part_qual.end());

                // contents after alignment
                if (ia_to != len_a)
                {
                    result_seq.append(seq1, ia_to, string::npos);
                    result_qual.insert(result_qual.end(), qual_list_1.begin() + ia_to, qual_list_1.end());
                }
                else if (ib_to != len_b)
                {
                    result_seq.append(seq2, ib_to, string::npos);
                    result_qual.insert(result_qual.end(), qual_list_2.begin() + ib_to, qual_list_2.end());
                }

                break;
            }
        }

        if (success) break;
    }

    return success;
}

void MicroAssembler::gen_qual_mask(const vector<int16_t>& quality, std::vector<bool>& mask)
{
    // init mask
    const size_t length = quality.size();

    if (length < km_sz)
    {
        mask.resize(0);
        return;
    }

    const size_t num_kmer = length - km_sz + 1;
    mask.resize(num_kmer);

    // check contents
    size_t i = 0;

    while (i < num_kmer)
    {
        // check the first kmer
        bool ok = true;
        size_t bad_i = 0;

        for (size_t j = 0; j < km_sz; j++)
        {
            size_t i_check = i + j;
            if (quality[i_check] < conf_qual)
            {
                ok = false;
                bad_i = i_check;
                break;
            }
        }

        if (ok)
        {
            mask[i] = true;
            // check next kmer
            // actually only check one base on the tail
            i++;
            for (; i < num_kmer; i++)
            {
                size_t i_tail = i + km_sz - 1;
                if (conf_qual <= quality[i_tail])
                    mask[i] = true;
                else
                {
                    // mark false within current kmer range
                    for (; i <= i_tail && i < num_kmer; i++)
                        mask[i] = false;
                    // now i should point to the one after i_tail
                    break;
                }
            }
        }
        else
        {
            // mark bad from i to bad_i
            for (; i <= bad_i && i < num_kmer; i++)
                mask[i] = false;
            // now i should point to the one after bad_i
        }

    }
}


} // namespace htqc

