//     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 "htio2/FastqIO.h"
#include "htio2/QualityUtil.h"

#define USE_prefix_out
#define USE_encode
#define USE_se_pe
#define USE_interleave
#include "htqc/Options.h"
#include "htqc/MicroAssembler.h"
#include "htqc/App.h"
#include "htqc/MultiSeqFile.h"

// we need to process reads in fwd strand
#define PE_STRAND_FILE OPT_pe_strand1(), OPT_pe_strand2()
#define PE_STRAND_OBJ htio2::STRAND_FWD, htio2::STRAND_FWD

using namespace std;
using namespace htqc;

string OPT_prefix_reject;
htio2::Option OPT_prefix_reject_ENTRY("reject", 'u', OPTION_GROUP_IO,
                                      &OPT_prefix_reject, htio2::Option::FLAG_NONE,
                                      "Output file prefix for non-overlapping reads.",
                                      "FILE_PREFIX");

#define OPTION_GROUP_PARAM "Overlapping Parameters"
size_t OPT_word_size = 7;
htio2::Option OPT_word_size_ENTRY("word-size", 'W', OPTION_GROUP_PARAM,
                                  &OPT_word_size, htio2::Option::FLAG_NONE,
                                  "Alignment is initiated from identical bases of this length.",
                                  "INT");

int16_t OPT_conf_qual = 20;
htio2::Option OPT_conf_qual_ENTRY("quality", 'Q', OPTION_GROUP_PARAM,
                                  &OPT_conf_qual, htio2::Option::FLAG_NONE,
                                  "Only the bases whose quality is higher than this are used for kmer generation and identity counting.",
                                  "INT");

size_t OPT_iden = 10;
htio2::Option OPT_iden_ENTRY("match", 'M', OPTION_GROUP_PARAM,
                             &OPT_iden, htio2::Option::FLAG_NONE,
                             "Number of matched bases required for assembly.",
                             "INT");

int16_t OPT_qual_diff = 15;
htio2::Option OPT_qual_diff_ENTRY("quality-diff", 'D', OPTION_GROUP_PARAM,
                                  &OPT_qual_diff, htio2::Option::FLAG_NONE,
                                  "Quality difference to accept mismatches.",
                                  "INT");

void parse_options(int argc, char** argv)
{
    // we are always processing PE
    OPT_pe = true;

    htio2::OptionParser opt_psr;

    opt_psr.add_option(OPT_files_in_ENTRY);
    opt_psr.add_option(OPT_prefix_out_ENTRY);
    opt_psr.add_option(OPT_prefix_reject_ENTRY);

    opt_psr.add_option(OPT_pe_itlv_ENTRY);
    opt_psr.add_option(OPT_pe_strand_ENTRY);
    opt_psr.add_option(OPT_encode_ENTRY);
    opt_psr.add_option(OPT_mask_ENTRY);

    opt_psr.add_option(OPT_word_size_ENTRY);
    opt_psr.add_option(OPT_conf_qual_ENTRY);
    opt_psr.add_option(OPT_iden_ENTRY);
    opt_psr.add_option(OPT_qual_diff_ENTRY);

    opt_psr.add_option(OPT_quiet_ENTRY);
    opt_psr.add_option(OPT_help_ENTRY);
    opt_psr.add_option(OPT_version_ENTRY);

    opt_psr.parse_options(argc, argv);

    //
    // show help
    //
    if (OPT_help)
    {
        cout << endl
             << argv[0] << " - assemble paired-ends into single sequences"
             << endl << endl
             << opt_psr.format_document();
        exit(EXIT_SUCCESS);
    }
    if (OPT_version) show_version_and_exit();
}

void validate_options()
{
    ensure_files_in();
    ensure_se_pe();
    ensure_encode();
    ensure_prefix_out();

    tidy_prefix_out(OPT_prefix_reject);
}

void show_options()
{
    show_files_in();
    show_prefix_out();
    if (OPT_prefix_reject.length())
        cout << "# file prefix for non-assembled reads: " << OPT_prefix_reject << endl;
    show_encode();
    cout << "# assembly parameters:" << endl;
    cout << "#   kmer size: " << OPT_word_size << endl;
    cout << "#   quality for confidental bases: " << OPT_conf_qual << endl;
    cout << "#   quality difference for acceptable conflict bases: " << OPT_qual_diff << endl;
    cout << "#   required number of identical bases: " << OPT_iden << endl;
}

size_t n_seq = 0;
size_t n_asm = 0;

void process_one(const htio2::FastqSeq& seq1,
                 const htio2::FastqSeq& seq2,
                 htio2::SeqIO& fh_asm,
                 htio2::PairedEndIO* fh_nasm,
                 htqc::MicroAssembler& assembler)
{
    n_seq++;

    vector<int16_t> qual_a;
    vector<int16_t> qual_b;
    htio2::decode_quality(seq1.quality, qual_a, OPT_encode);
    htio2::decode_quality(seq2.quality, qual_b, OPT_encode);

    htio2::FastqSeq result;

    bool re = assembler.assemble(seq1.seq, qual_a, seq2.seq, qual_b);

    if (re)
    {
        result.id = seq1.id;
        result.desc = assembler.result_desc;
        result.seq = assembler.result_seq;
        htio2::encode_quality(assembler.result_qual, result.quality, OPT_encode);
        fh_asm.write_seq(result);
        n_asm++;
    }
    else if (fh_nasm)
    {
        fh_nasm->write_pair(seq1, seq2, PE_STRAND_OBJ);
    }

    if (!OPT_quiet && n_seq % LOG_BLOCK == 0)
        cout << "  " << n_seq << " pairs, " << n_asm << " assembled" << endl;
}

void execute()
{
    // output handle
    string file_out;
    add_fastq_suffix_se(OPT_prefix_out, OPT_compress, file_out);
    check_output_overwrite(OPT_files_in, file_out);
    htio2::SeqIO::Ptr o_merge(htio2::SeqIO::New(file_out, htio2::WRITE, OPT_compress));

    // handle for non-overlap reads
    htio2::PairedEndIO* o_fail = NULL;

    if (OPT_prefix_reject.length())
    {
        if (OPT_pe_itlv)
        {
            // generate file name for non-overlap reads
            string file_u;
            add_fastq_suffix_se(OPT_prefix_reject, OPT_compress, file_u);
            check_output_overwrite(OPT_files_in, file_u);

            // create handle
            o_fail = new htio2::PairedEndIO(file_u,
                                            htio2::WRITE,
                                            PE_STRAND_FILE,
                                            OPT_compress);
        }
        else
        {
            // generate file name for non-overlap reads
            string file_u1;
            string file_u2;
            add_fastq_suffix_pe(OPT_prefix_reject, OPT_compress, file_u1, file_u2);
            check_output_overwrite(OPT_files_in, file_u1);
            check_output_overwrite(OPT_files_in, file_u2);

            // create handle
            o_fail = new htio2::PairedEndIO(file_u1,
                                            file_u2,
                                            htio2::WRITE,
                                            PE_STRAND_FILE,
                                            OPT_compress);
        }
    }

    // traverse input
    MicroAssembler assembler(OPT_word_size,
                             OPT_conf_qual,
                             OPT_iden,
                             OPT_qual_diff);

    if (!OPT_quiet)
        cout << "beginning assembly" << endl;

    htio2::FastqSeq seq1;
    htio2::FastqSeq seq2;

    if (OPT_pe_itlv)
    {
        htqc::MultiSeqFilePE fh_in(OPT_files_in,
                                    PE_STRAND_FILE,
                                    OPT_compress);

        while (fh_in.next_pair(seq1, seq2, PE_STRAND_OBJ))
            process_one(seq1, seq2, *o_merge, o_fail, assembler);
    }
    else
    {
        vector<string> files_in_a;
        vector<string> files_in_b;
        separate_paired_files(OPT_files_in, files_in_a, files_in_b);

        htqc::MultiSeqFilePE fh_in(files_in_a,
                                    files_in_b,
                                    PE_STRAND_FILE,
                                    OPT_compress);

        while (fh_in.next_pair(seq1, seq2, PE_STRAND_OBJ))
            process_one(seq1, seq2, *o_merge, o_fail, assembler);
    }

    if (!OPT_quiet)
        cout << "  " << n_seq << " pairs, " << n_asm << " assembled" << endl;

    if (o_fail)
        delete o_fail;
}
