#define USE_prefix_out
#define USE_encode
#define USE_mask
#define USE_header_format
#define USE_se_pe
#include "htqc/Options.h"
#include "htqc/App.h"
#include "htqc/MultiSeqFile.h"

#include "htio2/SeqIO.h"
#include "htio2/PairedEndIO.h"
#include "htio2/QualityUtil.h"

// We don't care strand in this program, so simply let them identical.
#define PE_STRAND_FILE htio2::STRAND_FWD, htio2::STRAND_FWD
#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 rejected sequences.", "FILE_PREFIX");

size_t OPT_cut_len(80);
htio2::Option OPT_cut_len_ENTRY("length", 'L', "Filter Options",
                                &OPT_cut_len, htio2::Option::FLAG_NONE,
                                "Cutoff for the length of high-quality base pairs.", "INT");

int16_t OPT_cut_qual(20);
htio2::Option OPT_cut_qual_ENTRY("quality", 'Q', "Filter Options",
                                 &OPT_cut_qual, htio2::Option::FLAG_NONE,
                                 "Cutoff for base pair quality.", "INT");

void parse_options(int argc, char** argv)
{
    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_se_ENTRY);
    opt_psr.add_option(OPT_pe_ENTRY);
    opt_psr.add_option(OPT_pe_itlv_ENTRY);

    opt_psr.add_option(OPT_cut_len_ENTRY);
    opt_psr.add_option(OPT_cut_qual_ENTRY);

    opt_psr.add_option(OPT_encode_ENTRY);
    opt_psr.add_option(OPT_mask_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] << " - filter sequences by length and quality"
             << endl << endl
             << opt_psr.format_document();
        exit(EXIT_SUCCESS);
    }

    if (OPT_version) show_version_and_exit();
}

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

    if (OPT_prefix_reject.length())
        tidy_prefix_out(OPT_prefix_reject);

    if (OPT_prefix_out == OPT_prefix_reject)
    {
        cerr << "conflict output prefix for accepted and rejected reads: \"" << OPT_prefix_reject << "\"" << endl;
        exit(EXIT_FAILURE);
    }
}

void show_options()
{
    show_se_pe();
    show_files_in();
    show_prefix_out();

    if (OPT_prefix_reject.length())
        cout << "# file prefix for rejected reads: " << OPT_prefix_reject << endl;

    cout << "# length cutoff: " << OPT_cut_len << endl;
    cout << "# quality cutoff: " << OPT_cut_qual << endl;
    show_encode();
}

bool validate_seq(const htio2::FastqSeq& seq)
{
    vector<int16_t> qual;
    htio2::decode_quality(seq.quality, qual, OPT_encode);

    size_t high_qual_len = 0;
    for (size_t i = 0; i < seq.length(); i++)
    {
        if (qual[i] >= OPT_cut_qual)
            high_qual_len++;
    }

    return (high_qual_len >= OPT_cut_len);
}

void filter_se()
{
    // input
    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);

    // output
    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 fh_out(htio2::SeqIO::New(file_out, htio2::WRITE, OPT_compress));

    // rejected
    htio2::SeqIO::Ptr fh_fail;
    if (OPT_prefix_reject.length())
    {
        string file_fail;
        add_fastq_suffix_se(OPT_prefix_reject, OPT_compress, file_fail);
        check_output_overwrite(OPT_files_in, file_fail);
        fh_fail = htio2::SeqIO::New(file_fail, htio2::WRITE, OPT_compress);
    }

    // traverse input
    htio2::FastqSeq curr_seq;
    int64_t n_accept = 0;
    int64_t n_seq = 0;

    while (fh_in.next_seq(curr_seq))
    {
        n_seq++;

        if (validate_seq(curr_seq))
        {
            fh_out->write_seq(curr_seq);
            n_accept++;
        }
        else
        {
            if (OPT_prefix_reject.length())
                fh_fail->write_seq(curr_seq);
        }

        if (!OPT_quiet && (n_seq % LOG_BLOCK == 0))
            cout << "\t" << n_seq << " reads, " << n_accept << " accept" << endl;
    }

    if (!OPT_quiet)
        cout << "\t" << n_seq << " reads, " << n_accept << " accept" << endl;
}

int64_t n_accept_pair = 0;
int64_t n_accept_single = 0;
int64_t n_pair = 0;

void filter_one_pe(const htio2::FastqSeq& seq1,
                   const htio2::FastqSeq& seq2,
                   htio2::PairedEndIO* fh_out,
                   htio2::SeqIO* fh_single,
                   htio2::PairedEndIO* fh_fail_pair,
                   htio2::SeqIO* fh_fail_single)
{
    n_pair++;

    bool accept_a = validate_seq(seq1);
    bool accept_b = validate_seq(seq2);

    if (accept_a)
    {
        if (accept_b)
        {
            fh_out->write_pair(seq1, seq2, PE_STRAND_OBJ);
            n_accept_pair++;
        }
        else
        {
            fh_single->write_seq(seq1);
            if (OPT_prefix_reject.length())
                fh_fail_single->write_seq(seq2);
            n_accept_single++;
        }
    }
    else
    {
        if (accept_b)
        {
            if(OPT_prefix_reject.length())
                fh_fail_single->write_seq(seq1);
            fh_single->write_seq(seq2);
            n_accept_single++;
        }
        else
        {
            if (OPT_prefix_reject.length())
                fh_fail_pair->write_pair(seq1, seq2, PE_STRAND_OBJ);
        }
    }

    if (!OPT_quiet && (n_pair % LOG_BLOCK == 0))
        cout << "\t" << n_pair << " pairs, " << n_accept_pair << " accept by pair, " << n_accept_single << " accept by single" << endl;
}

void filter_pe()
{
    // generate output file name for one-side-accepted reads
    string file_out_s;
    add_fastq_suffix_se(OPT_prefix_out + "_s", OPT_compress, file_out_s);
    check_output_overwrite(OPT_files_in, file_out_s);

    string file_fail_s;
    if (OPT_prefix_reject.length())
    {
        add_fastq_suffix_se(OPT_prefix_reject + "_s", OPT_compress, file_fail_s);
        check_output_overwrite(OPT_files_in, file_fail_s);
    }

    // tmp sequence
    htio2::FastqSeq seq1;
    htio2::FastqSeq seq2;

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

        // output
        string file_out_p;
        add_fastq_suffix_se(OPT_prefix_out + "_p", OPT_compress, file_out_p);
        check_output_overwrite(OPT_files_in, file_out_p);

        htio2::PairedEndIO::Ptr fh_out_p(new htio2::PairedEndIO(file_out_p,
                                                              htio2::WRITE,
                                                              PE_STRAND_FILE,
                                                              OPT_compress)
                                       );

        htio2::SeqIO::Ptr fh_out_s(htio2::SeqIO::New(file_out_s,
                                                   htio2::WRITE,
                                                   OPT_compress)
                                 );

        // fail
        htio2::PairedEndIO::Ptr fh_fail_p;
        htio2::SeqIO::Ptr fh_fail_s;

        if (OPT_prefix_reject.length())
        {
            string file_fail_p;
            add_fastq_suffix_se(OPT_prefix_reject + "_p", OPT_compress, file_fail_p);
            check_output_overwrite(OPT_files_in, file_fail_p);

            fh_fail_p = new htio2::PairedEndIO(file_fail_p, htio2::WRITE, PE_STRAND_FILE, OPT_compress);
            fh_fail_s = htio2::SeqIO::New(file_fail_s, htio2::WRITE, OPT_compress);
        }

        // traverse input files
        while (fh_in.next_pair(seq1, seq2, PE_STRAND_OBJ))
            filter_one_pe(seq1, seq2, fh_out_p.get(), fh_out_s.get(), fh_fail_p.get(), fh_fail_s.get());
    }
    else
    {
        // input
        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);


        // output
        string file_out_a;
        string file_out_b;
        add_fastq_suffix_pe(OPT_prefix_out, OPT_compress, file_out_a, file_out_b);
        check_output_overwrite(OPT_files_in, file_out_a);
        check_output_overwrite(OPT_files_in, file_out_b);

        htio2::PairedEndIO::Ptr fh_out_p(new htio2::PairedEndIO(file_out_a,
                                                              file_out_b,
                                                              htio2::WRITE,
                                                              PE_STRAND_FILE,
                                                              OPT_compress)
                                       );

        htio2::SeqIO::Ptr fh_out_s(htio2::SeqIO::New(file_out_s,
                                                   htio2::WRITE,
                                                   OPT_compress)
                                 );

        // fail
        htio2::PairedEndIO::Ptr fh_fail_p;
        htio2::SeqIO::Ptr fh_fail_s;
        if (OPT_prefix_reject.length())
        {
            string file_fail_a, file_fail_b;
            add_fastq_suffix_pe(OPT_prefix_reject, OPT_compress, file_fail_a, file_fail_b);
            check_output_overwrite(OPT_files_in, file_fail_a);
            check_output_overwrite(OPT_files_in, file_fail_b);

            fh_fail_p = new htio2::PairedEndIO(file_fail_a, file_fail_b, htio2::WRITE, PE_STRAND_FILE, OPT_compress);
            fh_fail_s = htio2::SeqIO::New(file_fail_s, htio2::WRITE, OPT_compress);
        }

        // traverse input files
        while (fh_in.next_pair(seq1, seq2, PE_STRAND_OBJ))
            filter_one_pe(seq1, seq2, fh_out_p.get(), fh_out_s.get(), fh_fail_p.get(), fh_fail_s.get());
    }

    if (!OPT_quiet)
        cout << "\t" << n_pair << " input, " << n_accept_pair << " accept by pair, " << n_accept_single << " accept by single" << endl;
}

void execute()
{
    if (OPT_se)
        filter_se();
    else if (OPT_pe)
        filter_pe();
    else
        throw runtime_error("neither single-end nor paired-end");
}
